"""Steering direction 'hesitation' from the model's own reasoning traces (no labels, no external model).

At every paragraph break ("\n\n") inside <think>, the residual stream of each layer is recorded at the break token,
i.e. at the moment the model decides how to continue. Paragraphs that open with a hesitation/re-check marker
("Wait", "Alternatively", "Hmm", "Let me double-check", ...) form class R, all others class E.
direction[layer] = mean_R - mean_E. Saved as .pt (all layers) for building llama.cpp control vectors.

usage: extract_cvec.py --gen results/dev_orig/dev.jsonl --items data/train/dev.jsonl --out results/cvec/hes.pt
"""
import argparse, json, re, 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()
ap.add_argument("--gen", required=True)
ap.add_argument("--items", required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--model", default="models/LFM3-1.7B")
ap.add_argument("--max-resp", type=int, default=200)
ap.add_argument("--max-len", type=int, default=8192)
args = ap.parse_args()

HES = re.compile(r"^(\*\*)?(wait|but wait|hmm|alternatively|actually|hold on|oh wait|however, wait|let me (double[- ]check|"
                 r"check (again|this|that)|verify|re-?examine|reconsider|think again|re-?check|re-?read|make sure)|"
                 r"but let me|let's (double[- ]check|verify|check again))", re.I)
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 = model.config.num_hidden_layers
items = {json.loads(l)["id"]: json.loads(l) for l in open(args.items, encoding="utf-8")}
recs = [json.loads(l) for l in open(args.gen, encoding="utf-8")]
recs = [r for r in recs if items[r["id"]]["task"] == "math" and "tokens" in r][: args.max_resp]

cap = {}
hooks = [layer.register_forward_hook(lambda m, i, o, k=k: cap.__setitem__(k, (o[0] if isinstance(o, tuple) else o)[0]))
         for k, layer in enumerate(model.model.layers)]
sums = {c: torch.zeros(L, model.config.hidden_size, dtype=torch.float64) for c in "RE"}
cnt = {"R": 0, "E": 0}
norm = torch.zeros(L, dtype=torch.float64)
for n_r, r in enumerate(recs):
    p = tok.apply_chat_template([{"role": "user", "content": items[r["id"]]["prompt"]}], tokenize=True,
                                add_generation_prompt=True, return_dict=False)
    resp = r["tokens"]
    pieces = tok.convert_ids_to_tokens(resp)
    pieces = [tok.convert_tokens_to_string([x]) for x in pieces]
    text, pos, labels = "", [], []
    for i, s in enumerate(pieces):
        text += s
        if "</think>" in text:
            break
        if s.endswith("\n\n") and i + 1 < len(resp):
            nxt = tok.decode(resp[i + 1: i + 12]).lstrip()
            pos.append(len(p) + i)
            labels.append("R" if HES.match(nxt) else "E")
    ids = (p + resp)[: args.max_len]
    keep = [(q, c) for q, c in zip(pos, labels) if q < len(ids)]
    if not keep:
        continue
    with torch.no_grad():
        model(input_ids=torch.tensor([ids], device="cuda"), use_cache=False)
    qi = torch.tensor([q for q, _ in keep], device="cuda")
    isr = torch.tensor([c == "R" for _, c in keep])
    for k in range(L):
        h = cap[k][qi].double().cpu()
        sums["R"][k] += h[isr].sum(0)
        sums["E"][k] += h[~isr].sum(0)
        norm[k] += h.norm(dim=-1).sum()
    for _, c in keep:
        cnt[c] += 1
    if n_r % 20 == 0:
        print(n_r, cnt, flush=True)
mR, mE = sums["R"] / cnt["R"], sums["E"] / cnt["E"]
d = (mR - mE).float()
avg_norm = (norm / (cnt["R"] + cnt["E"])).float()
for k in range(L):
    print(f"layer {k:2d}  |dir| {d[k].norm():8.2f}  |h| {avg_norm[k]:8.2f}  ratio {d[k].norm() / avg_norm[k]:.3f}  "
          f"cos(mR,mE) {torch.nn.functional.cosine_similarity(mR[k], mE[k], 0):.3f}")
torch.save({"dir": d, "avg_norm": avg_norm, "count": cnt}, args.out)
print("saved", args.out, cnt)