"""'Correct-reasoning' direction from discordant pairs: problems where one of the model's two samples is right and
the other wrong. Difficulty is controlled (same problem), length is controlled (both responses are averaged over the
same number of leading tokens, at most --n-tok). direction[layer] = mean over pairs of (h_correct - h_wrong).
usage: extract_dir_correct.py --gen results/dev_orig/dev.jsonl --graded results/dev_orig/graded.jsonl
                              --items data/train/dev.jsonl --out results/cvec/cor.pt"""
import argparse, json, collections, torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import transformers.integrations.sdpa_attention as _sdpa

# Windows torch has no flash kernel: masked/GQA SDPA would fall back to the O(L^2)-memory math kernel.
# Expanding K/V instead keeps it on the memory-efficient kernel.
_sdpa.use_gqa_in_sdpa = lambda *a, **k: False
ap = argparse.ArgumentParser()
for a in ["--gen", "--graded", "--items", "--out"]:
    ap.add_argument(a, required=True)
ap.add_argument("--model", default="models/LFM3-1.7B")
ap.add_argument("--n-tok", type=int, default=2048)
args = ap.parse_args()
tok = AutoTokenizer.from_pretrained(args.model)
model = AutoModelForCausalLM.from_pretrained(args.model, dtype=torch.bfloat16, attn_implementation="sdpa",
                                             trust_remote_code=True).cuda().eval()
L, H = model.config.num_hidden_layers, model.config.hidden_size
items = {json.loads(l)["id"]: json.loads(l) for l in open(args.items, encoding="utf-8")}
ok = {(g["id"], g["sample"]): g["correct"] for g in map(json.loads, open(args.graded, encoding="utf-8"))}
recs = collections.defaultdict(dict)
for l in open(args.gen, encoding="utf-8"):
    r = json.loads(l)
    if "tokens" in r and (r["id"], r["sample"]) in ok:
        recs[r["id"]][r["sample"]] = r
pairs = []
for pid, s in recs.items():
    if len(s) == 2 and ok[(pid, 0)] != ok[(pid, 1)]:
        c, w = (s[0], s[1]) if ok[(pid, 0)] else (s[1], s[0])
        pairs.append((pid, c, w))
print(len(pairs), "discordant pairs", collections.Counter(items[p]["task"] for p, _, _ in pairs), flush=True)
cap = {}
for k, layer in enumerate(model.model.layers):
    layer.register_forward_hook(lambda m, i, o, k=k: cap.__setitem__(k, (o[0] if isinstance(o, tuple) else o)[0]))


def mean_states(prompt_ids, resp, n):
    ids = prompt_ids + resp[:n]
    with torch.no_grad():
        model(input_ids=torch.tensor([ids], device="cuda"), use_cache=False)
    return torch.stack([cap[k][len(prompt_ids):].float().mean(0) for k in range(L)]).double().cpu()


diff = torch.zeros(L, H, dtype=torch.float64)
norm = torch.zeros(L, dtype=torch.float64)
per_pair = []
for i, (pid, c, w) in enumerate(pairs):
    p = tok.apply_chat_template([{"role": "user", "content": items[pid]["prompt"]}], tokenize=True,
                                add_generation_prompt=True, return_dict=False)
    n = min(len(c["tokens"]), len(w["tokens"]), args.n_tok)
    hc, hw = mean_states(p, c["tokens"], n), mean_states(p, w["tokens"], n)
    diff += hc - hw
    norm += 0.5 * (hc.norm(dim=-1) + hw.norm(dim=-1))
    per_pair.append((hc - hw).float())
    if i % 10 == 0:
        print(i, flush=True)
d = (diff / len(pairs)).float()
avg_norm = (norm / len(pairs)).float()
P = torch.stack(per_pair)  # pairs x L x H
for k in range(L):
    loo = (P[:, k].sum(0, keepdim=True) - P[:, k]) / (len(P) - 1)  # mean of the OTHER pairs (leave-one-out)
    agree = ((P[:, k] * loo).sum(-1) > 0).float().mean()  # held-out agreement; 0.5 = no consistent direction
    print(f"layer {k:2d} |dir| {d[k].norm():6.3f} |h| {avg_norm[k]:6.2f} ratio {d[k].norm() / avg_norm[k]:.3f} pair-agreement {agree:.2f}")
torch.save({"dir": d, "avg_norm": avg_norm, "count": len(pairs)}, args.out)
print("saved", args.out)