import base64
import os
from compression import zstd

import numpy as np
from openai import OpenAI

LAYER = 15


def decode(tensor):
    raw = zstd.decompress(base64.b64decode(tensor["data"]))
    array = np.frombuffer(raw, dtype=tensor["dtype"]).reshape(tensor["shape"])
    if tensor["original_dtype"] == "torch.bfloat16":
        array = (array.view(np.uint16).astype(np.uint32) << 16).view(np.float32)
    return array


def mean_over_tokens(prompt):
    completion = client.completions.create(
        model="llama-8b",
        prompt=prompt,
        max_tokens=1,
        extra_body={"output_residual_stream": [LAYER]},
    )
    residual_stream = decode(completion.activations["residual_stream"])
    return residual_stream[0, 1:].mean(axis=0)


client = OpenAI(base_url=os.environ["ACS_API_BASE"], api_key=os.environ["ACS_API_KEY"])
pirate = [
    "Arr, the wind be fair and the sea be calm.",
    "Ahoy, me hearties, hoist the sails!",
    "Shiver me timbers, that be a fine ship.",
    "Avast, ye scallywags, the treasure be ours!",
]
plain = [
    "The wind is fair and the sea is calm.",
    "Hello, friends, raise the sails!",
    "Goodness, that is a fine ship.",
    "Stop, you rascals, the treasure is ours!",
]
contrast = np.mean([mean_over_tokens(p) for p in pirate], axis=0) - np.mean(
    [mean_over_tokens(p) for p in plain], axis=0
)
print(contrast.shape, round(float(np.linalg.norm(contrast)), 2))

for scale in [0, 4, 8]:
    completion = client.completions.create(
        model="llama-8b",
        prompt="My advice for your first day at a new job is",
        max_tokens=16,
        temperature=0,
        extra_body={
            "apply_steering_vectors": [{"activations": [contrast.tolist()], "layer_indices": [LAYER], "scale": scale}]
        },
    )
    print(scale, repr(completion.choices[0].text))