"""Calibrate four residual directions using only the grouped training split."""

import argparse
import json
from pathlib import Path
from codelock_affect.channels import CHANNELS
from codelock_affect.config import ProjectConfig
from codelock_affect.extraction import extract_vector
from codelock_affect.layers import decoder_layers, parse_layers
from codelock_affect.model import load
from codelock_affect.vectors import VectorBank
from codelock_affect.evaluation import environment, sha256_file, write_json


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--config", default="configs/default.json")
    p.add_argument("--model")
    p.add_argument("--revision")
    p.add_argument("--data", default="data/affect_pairs.jsonl")
    p.add_argument("--output", default="vectors/multichannel")
    p.add_argument("--four-bit", action="store_true")
    args = p.parse_args()
    cfg = ProjectConfig.load(args.config)
    pairs = [
        json.loads(line)
        for line in Path(args.data).read_text(encoding="utf-8").splitlines()
        if line.strip()
    ]
    model, tok = load(args.model or cfg.model, args.four_bit, revision=args.revision)
    layers = parse_layers(cfg.layers, len(decoder_layers(model)))
    vectors = {}
    for channel in CHANNELS:
        examples = [
            row
            for row in pairs
            if row["channel"] == channel and row["split"] == "train"
        ]
        vector = extract_vector(
            model, tok, examples, layers, channel, lambda text: print(text, flush=True)
        )
        vector.meta.update(
            {
                "dataset_file_sha256": sha256_file(args.data),
                "quantized_4bit": args.four_bit,
                "environment": environment(),
                "requested_revision": args.revision,
            }
        )
        vectors[channel] = vector
        # Save each completed channel immediately; bank.json is the commit point.
        vector.save(Path(args.output) / f"{channel}.safetensors")
    bank = VectorBank(vectors)
    bank.save(args.output)
    write_json(
        Path(args.output) / "extraction.json",
        {
            "environment": environment(),
            "config": cfg.to_dict(),
            "dataset_sha256": sha256_file(args.data),
            "channels": {c: v.meta for c, v in vectors.items()},
        },
    )
    print(f"Saved {len(vectors)} channels to {args.output}", flush=True)


if __name__ == "__main__":
    main()