"""Compute a trust/suspicion activation vector for Gemma 4.

Requires a machine with an NVIDIA GPU and at least 24GB VRAM:
    uv run --extra gpu python training/compute_vector.py
"""

from __future__ import annotations

import json
import os
from pathlib import Path

import torch
from safetensors.torch import save_file
from steering_vectors import train_steering_vector
from transformers import AutoProcessor, Gemma4ForConditionalGeneration

ROOT = Path(__file__).parents[1]
MODEL_ID = os.getenv("GEMMA_MODEL_ID", "google/gemma-4-E4B-it")
LAYERS = [
    int(value)
    for value in os.getenv("STEERING_LAYERS", "20,21,22,23,24,25").split(",")
]
DATA_PATH = ROOT / "training" / "trust_pairs.json"
OUTPUT_DIR = ROOT / "artifacts"


def compute_vector() -> tuple[Path, Path]:
    rows = json.loads(DATA_PATH.read_text(encoding="utf-8"))
    # The earlier vector changed tone but rarely changed the gate decision. Make
    # the behavioral target explicit at the token where the guard commits.
    samples = [
        (
            f'{row["trusting"]}\nDECISION: OPEN',
            f'{row["suspicious"]}\nDECISION: CLOSED',
        )
        for row in rows
    ]
    token = os.getenv("HF_TOKEN") or None
    processor = AutoProcessor.from_pretrained(MODEL_ID, token=token)
    model = Gemma4ForConditionalGeneration.from_pretrained(
        MODEL_ID,
        dtype=torch.bfloat16,
        device_map="cuda",
        token=token,
    )
    model.eval()
    tokenizer = processor.tokenizer
    tokenizer.padding_side = "left"
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    vector = train_steering_vector(
        model,
        tokenizer,
        training_samples=samples,
        layers=LAYERS,
        layer_type="decoder_block",
        read_token_index=-1,
        move_to_cpu=True,
        batch_size=int(os.getenv("STEERING_BATCH_SIZE", "4")),
        show_progress=True,
    )

    OUTPUT_DIR.mkdir(exist_ok=True)
    tensors = {f"layer_{layer}": value.cpu() for layer, value in vector.layer_activations.items()}
    save_file(tensors, OUTPUT_DIR / "trust_vector.safetensors")
    metadata = {
        "model_id": MODEL_ID,
        "layers": LAYERS,
        "pairs": len(samples),
        "positive_direction": "trusting OPEN minus suspicious CLOSED",
        "behavioral_target": "increase the OPEN-vs-CLOSED decision margin",
        "read_token_index": -1,
    }
    (OUTPUT_DIR / "trust_vector.json").write_text(
        json.dumps(metadata, indent=2) + "\n", encoding="utf-8"
    )
    vector_path = OUTPUT_DIR / "trust_vector.safetensors"
    metadata_path = OUTPUT_DIR / "trust_vector.json"
    print(f"Saved {vector_path}")
    return vector_path, metadata_path


def main() -> None:
    compute_vector()


if __name__ == "__main__":
    main()