"""Score control vectors against a fitted fine-tune: how much of the fine-tune's shift away from stock a
vector-steered stock reproduces.

For each arm, stock runs with the arm's vector added at the end of every layer (as llama.cpp adds it:
nothing at layer 0, every position including BOS), and its next-token distributions are compared with
the fine-tune's at the reference positions of both splits:

  recovery       1 - KL(tune || stock + vector) / KL(tune || stock): 1 is a perfect copy, 0 is no better
                 than stock, below 0 is further from the fine-tune than stock is
  top1_recovery  the same for top-1 agreement: the share of the fine-tune's changed top tokens restored

Choose the strength on the fit split and read the held-out split; `best` in the output does that for
each vector (a --vector at several strengths is one vector; each --gguf and --blend is its own).

    python -m ecce_vector.score --base gemma-4-31B-it/ --ref fits/some-finetune.npz \\
        --vector fits/some-finetune.npz::vectors_track --alphas 0.5 0.75 1.0 1.25 1.5 --out score.json
    python -m ecce_vector.score --base gemma-4-31B-it/ --ref fits/merge.npz \\
        --gguf vectors/merge.gguf --blend vectors/a.gguf:0.5,vectors/b.gguf:0.5

An arm is one stream through each stock layer, computed as a difference from the recorded stock stream
(fp16 matmul on the difference, which reproduces full-precision recoveries to about 1e-4), so each extra
arm costs a fraction of a full pass.
"""
from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

import mlx.core as mx
import numpy as np
from transformers import AutoTokenizer

from .corpus import Corpus
from .io import N_LAYERS, Checkpoint, read_gguf
from .model import compare, embed, gather_rows, log_probs, make_layer, model_args

SPLITS = ("fit", "held_out")
REF_KEY = {"fit": "ref_fit", "held_out": "ref_ho"}
GROUP = 2           # arms per GPU call


def load_vector(spec: str) -> np.ndarray:
    """A (60, d) vector from FILE.npz[::KEY] (default key vectors_track) or FILE.gguf; layer 0 zeroed."""
    if spec.endswith(".gguf"):
        V = read_gguf(spec)
    else:
        path, key = spec.split("::") if "::" in spec else (spec, "vectors_track")
        V = np.load(path)[key].astype(np.float32).copy()
    V[0] = 0.0          # llama.cpp has no layer-0 vector
    return V


def reference(ref_path: Path):
    """The fine-tune's side of the comparison: log-probs if it has its own head, else hidden states + norm."""
    z = np.load(ref_path)
    sibling = ref_path.with_name(ref_path.stem + ".logp.npz")
    if sibling.exists():
        return np.load(sibling), None
    if "logp_file" in z.files:          # fits written by the original workspace tools
        return np.load(str(z["logp_file"])), None
    return z, (z["norm"] if "norm" in z.files else None)


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--base", required=True, help="stock google/gemma-4-31B-it checkpoint directory")
    ap.add_argument("--ref", required=True, type=Path, help="the fit's NAME.npz (its reference)")
    ap.add_argument("--vector", action="append", default=[], help="FILE.npz[::KEY] scored at each of --alphas")
    ap.add_argument("--alphas", type=float, nargs="+", default=[0.25, 0.5, 0.75, 1.0, 1.25, 1.5])
    ap.add_argument("--gguf", action="append", default=[], help="a control-vector GGUF, scored at 1.0 as released")
    ap.add_argument("--blend", action="append", default=[], help="A.gguf:w1,B.gguf:w2,...  summed, as llama.cpp does")
    ap.add_argument("--out", type=Path, help="write the results here (JSON)")
    a = ap.parse_args(argv)

    arms, family = {}, {}           # family: the vector an arm belongs to; `best` is chosen within each
    for spec in a.vector:
        V = load_vector(spec)
        for alpha in a.alphas:
            arms[f"{spec}@{alpha:g}"] = alpha * V
            family[f"{spec}@{alpha:g}"] = spec
    for spec in a.gguf:
        arms[spec] = load_vector(spec)
        family[spec] = spec
    for spec in a.blend:
        parts = [p.rsplit(":", 1) for p in spec.split(",")]
        arms[spec] = sum(float(w) * load_vector(f) for f, w in parts)
        family[spec] = spec
    assert arms, "nothing to score: give --vector, --gguf or --blend"

    t0 = time.time()
    args = model_args(a.base)
    stock = Checkpoint(a.base)
    corpus = Corpus(AutoTokenizer.from_pretrained(a.base))
    names = list(arms)
    groups = [names[i:i + GROUP] for i in range(0, len(names), GROUP)]
    h = {}
    for s in SPLITS:
        e0 = embed(stock, corpus.ids[s], args.hidden_size)
        h[("stock", s)] = e0
        for k in names:
            h[(k, s)] = e0
    for L in range(N_LAYERS):
        tl = time.time()
        layer, sides = make_layer(args, L, stock.layer(L))
        new = {}
        for s in SPLITS:
            for w in sides.values():
                w.mode = "rec"
            y = layer(h[("stock", s)], mask="causal")[0]
            mx.eval(y, [(w.xs, w.ys) for w in sides.values()])
            new[("stock", s)] = y
            for w in sides.values():
                w.mode = "diff"
            n = y.shape[0]
            for grp in groups:
                ys = layer(mx.concatenate([h[(k, s)] for k in grp], 0), mask="causal")[0]
                for j, k in enumerate(grp):
                    new[(k, s)] = ys[j * n:(j + 1) * n] + mx.array(arms[k][L])
                mx.eval([new[(k, s)] for k in grp])
            for w in sides.values():
                w.xs = w.ys = None
        h = new
        del layer, sides
        mx.clear_cache()
        print(f"L{L:02d} {time.time() - tl:.1f}s (total {time.time() - t0:.0f}s)", flush=True)

    norm0 = mx.array(stock.get("norm.weight"))
    emb = mx.array(stock.get("embed_tokens.weight")).astype(mx.bfloat16)
    ref, tune_norm = reference(a.ref)
    res = {"ref": str(a.ref), "stock": {}, "arms": {k: {} for k in names}}
    for s in SPLITS:
        lq_stock = log_probs(args, gather_rows(h[("stock", s)], corpus.rows[s]), norm0, emb)
        lp_key = REF_KEY[s].replace("ref", "logp")
        if lp_key in ref.files:
            lt = mx.array(ref[lp_key]).astype(mx.float32)
        else:
            lt = log_probs(args, mx.array(ref[REF_KEY[s]]),
                           mx.array(tune_norm) if tune_norm is not None else norm0, emb)
        ks, ts, kps = compare(lt, lq_stock, corpus.spans[s])
        res["stock"][s] = {"kl": ks, "top1_agree": ts, "kl_per_sequence": kps}
        for k in names:
            kv, tv, kpv = compare(lt, log_probs(args, gather_rows(h[(k, s)], corpus.rows[s]), norm0, emb),
                                  corpus.spans[s])
            res["arms"][k][s] = {"kl": kv, "recovery": 1 - kv / ks, "top1_agree": tv,
                                 "top1_recovery": (tv - ts) / (1 - ts) if ts < 1 else 0.0,
                                 "recovery_per_sequence": [1 - x / y for x, y in zip(kpv, kps)]}
        del lt, lq_stock
    # for each vector, the strength with the lowest fit-split KL, reported on the held-out split
    res["best"] = {}
    for fam in dict.fromkeys(family.values()):
        best = min((k for k in names if family[k] == fam), key=lambda k: res["arms"][k]["fit"]["kl"])
        res["best"][fam] = {"arm": best, **res["arms"][best]["held_out"]}
    print(f"KL(tune || stock) held out {res['stock']['held_out']['kl']:.4f}")
    for k in names:
        r = res["arms"][k]
        print(f"{k:60s} fit: rec {r['fit']['recovery']:+.4f}  held out: rec {r['held_out']['recovery']:+.4f} "
              f"top-1 {r['held_out']['top1_recovery']:+.3f}")
    for fam, b in res["best"].items():
        print(f"{fam}: chosen on the fit split {b['arm']} -> held-out recovery {b['recovery']:+.4f}, "
              f"top-1 {b['top1_recovery']:+.3f}")
    print(f"({time.time() - t0:.0f}s)")
    if a.out:
        a.out.write_text(json.dumps(res, indent=1))


if __name__ == "__main__":
    main()