"""Extract the pain direction from contrastive pairs and save it as safetensors.

    python scripts/extract_pain_vector.py \
        --model Qwen/Qwen3-8B --pairs data/contrastive_pairs.jsonl \
        --layers mid --out vectors/pain_qwen3-8b.safetensors
"""

import argparse
import json
import os
import sys
import time

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from codelock_affect import model as cm  # noqa: E402
from codelock_affect.steering import (
    decoder_layers,
    extract_vector,
    parse_layers,
)  # noqa: E402


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default=cm.DEFAULT_MODEL)
    ap.add_argument("--pairs", default="data/contrastive_pairs.jsonl")
    ap.add_argument("--layers", default="mid", help="'mid', '10-26' or '8,12,16'")
    ap.add_argument("--out", default="vectors/pain.safetensors")
    ap.add_argument("--four-bit", action="store_true")
    ap.add_argument(
        "--model-id",
        default=None,
        help="hub id of the base model recorded in the vector metadata (defaults to --model)",
    )
    args = ap.parse_args()

    with open(args.pairs, encoding="utf-8") as f:
        pairs = [json.loads(line) for line in f if line.strip()]
    print(f"loading {args.model} ...")
    model, tok = cm.load(args.model, four_bit=args.four_bit)
    layers = parse_layers(args.layers, len(decoder_layers(model)))
    print(f"{len(pairs)} pairs, layers {layers[0]}-{layers[-1]}")

    t0 = time.time()
    vec = extract_vector(model, tok, pairs, layers, name="pain", log=print)
    vec.meta["base_model"] = args.model_id or args.model
    os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
    vec.save(args.out)
    print(f"saved {args.out} in {time.time() - t0:.1f}s")
    for i in layers:
        print(f"  layer {i:2d}  |h| = {vec.norms[i]:8.2f}")


if __name__ == "__main__":
    main()