"""Metacognitive halting module for LFM3 (an internal component, the rest of the model is untouched).

The output row of </think> is untied from the input embedding and extended by a learned vector D:
    logit_</think>(h) = (e_</think> + D) . h          h = final normalised hidden state
so the model's own probability of closing its reasoning becomes a learned function of its internal state ("is the
answer I would give now already right?"). Nothing else changes: all other logits, all weights, the input embedding of
</think>. Integration is native: a GGUF with a separate output.weight runs in stock llama.cpp, zero extra compute.

Training signal (no labels from outside the model except the reference answer of each problem):
  * sentence ends where the model, forced to answer, is right                  -> close (y=1)
  * sentence ends where it is still wrong                                      -> keep thinking (y=0)
  * the position where the model itself closed thinking                        -> close (y=1), keep it intact
  * random mid-sentence positions                                              -> keep thinking (y=0), don't disturb
Only D (2048 numbers) is optimised, on the exact sampling probability (full-vocabulary softmax at temperature T).

  extract: halt_module.py extract --rec results/rec/raw.jsonl --items "data/coc/*.jsonl" --exclude data/coc/mini_math.jsonl
           --out results/halt/feats.pt
  train:   halt_module.py train --feats results/halt/feats.pt --out results/halt/D.pt [--pos-weight 1 --l2 1e-3]
  export:  halt_module.py export --D results/halt/D.pt --scale 1.0 --gguf-in gguf/lfm3-orig-q5km.gguf --gguf-out gguf/halt.gguf
"""
import argparse
import glob
import json
import os
import random
import sys

import torch

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
THINK_END = 124902
ap = argparse.ArgumentParser()
ap.add_argument("cmd", choices=["extract", "train", "export"])
ap.add_argument("--rec", nargs="+", default=[])
ap.add_argument("--items", nargs="+", default=[])
ap.add_argument("--exclude", nargs="*", default=[])
ap.add_argument("--feats")
ap.add_argument("--out")
ap.add_argument("--model", default="models/LFM3-1.7B")
ap.add_argument("--temp", type=float, default=0.6)
ap.add_argument("--n-random", type=int, default=12)
ap.add_argument("--pos-weight", type=float, default=1.0)
ap.add_argument("--l2", type=float, default=1e-3)
ap.add_argument("--epochs", type=int, default=400)
ap.add_argument("--D")
ap.add_argument("--scale", type=float, default=1.0)
ap.add_argument("--gguf-in")
ap.add_argument("--gguf-out")
args = ap.parse_args()


def extract():
    from transformers import AutoModelForCausalLM, AutoTokenizer
    import transformers.integrations.sdpa_attention as _sdpa
    _sdpa.use_gqa_in_sdpa = lambda *a, **k: False  # memory-efficient kernel on Windows
    from grade import grade_one
    tok = AutoTokenizer.from_pretrained(args.model)
    model = AutoModelForCausalLM.from_pretrained(args.model, dtype=torch.bfloat16, attn_implementation="sdpa",
                                                 trust_remote_code=True).cuda().eval()
    W = model.get_output_embeddings().weight  # [V, H]
    ends = set()
    for i in range(len(tok)):
        t = tok.convert_ids_to_tokens(i)
        if t is not None and (t.endswith(".") or t.endswith("Ċ")):
            ends.add(i)
    items = {}
    for pat in args.items:
        for f in glob.glob(pat):
            for l in open(f, encoding="utf-8"):
                r = json.loads(l)
                items[r["id"]] = r
    excl = set()
    for f in args.exclude:
        excl |= {json.loads(l)["id"] for l in open(f, encoding="utf-8")}
    rng = random.Random(0)
    H, LSE, BASE, Y, KIND, TRAJ = [], [], [], [], [], []
    n_traj = 0
    for path in args.rec:
        for l in open(path, encoding="utf-8"):
            r = json.loads(l)
            if r["id"] in excl or r["id"] not in items or not r.get("probe_log"):
                continue
            it = items[r["id"]]
            p_ids = tok.apply_chat_template([{"role": "user", "content": it["prompt"]}], tokenize=True,
                                            add_generation_prompt=True, return_dict=False)
            text = r["completion"][len("<think>"):]
            resp = tok(text, add_special_tokens=False)["input_ids"][:8192]
            pos, ys, kinds = [], [], []
            # probe cuts (indices into the original generation; re-tokenised text may shift a few tokens:
            # snap to the nearest sentence end at or before the cut)
            for p in r["probe_log"]:
                c = min(p["cut"], len(resp))
                while c > 1 and resp[c - 1] not in ends:
                    c -= 1
                if c <= 1:
                    continue
                ok = grade_one({**r, "completion": "<think>x</think>" + p["text"], "n_tokens": 0}, it)["correct"]
                pos.append(c - 1)  # hidden state at the sentence-end token predicts the next token
                ys.append(1.0 if ok else 0.0)
                kinds.append("cut")
            if THINK_END in resp:
                j = resp.index(THINK_END)
                pos.append(j - 1)
                ys.append(1.0)
                kinds.append("natural")
            mid = [j for j in range(8, min(len(resp), 8192) - 1) if resp[j] not in ends and resp[j] != THINK_END]
            for j in rng.sample(mid, min(args.n_random, len(mid))):
                pos.append(j)
                ys.append(0.0)
                kinds.append("random")
            if not pos:
                continue
            ids = p_ids + resp
            with torch.no_grad():
                h = model.model(input_ids=torch.tensor([ids], device="cuda"), use_cache=False).last_hidden_state[0]
                hp = h[[len(p_ids) + q for q in pos]].float()
                logits = hp @ W.float().t() / args.temp
                base = logits[:, THINK_END].clone()
                logits[:, THINK_END] = -1e30
                lse = torch.logsumexp(logits, -1)
            H.append(hp.cpu()); LSE.append(lse.cpu()); BASE.append(base.cpu())
            Y += ys; KIND += kinds; TRAJ += [n_traj] * len(pos)
            n_traj += 1
            if n_traj % 25 == 0:
                print(n_traj, "trajectories", len(Y), "points", flush=True)
    os.makedirs(os.path.dirname(args.out), exist_ok=True)
    torch.save({"h": torch.cat(H), "lse": torch.cat(LSE), "base": torch.cat(BASE), "y": torch.tensor(Y),
                "kind": KIND, "traj": torch.tensor(TRAJ), "temp": args.temp}, args.out)
    print("saved", args.out, n_traj, "trajectories", len(Y), "points",
          {k: (sum(1 for x in KIND if x == k), sum(y for x, y in zip(KIND, Y) if x == k)) for k in set(KIND)})


def auc(s, y):
    order = torch.argsort(s)
    r = torch.empty_like(order, dtype=torch.float)
    r[order] = torch.arange(1, len(s) + 1, dtype=torch.float)
    n1 = y.sum()
    n0 = len(y) - n1
    return ((r[y == 1].sum() - n1 * (n1 + 1) / 2) / (n1 * n0)).item() if n1 > 0 and n0 > 0 else float("nan")


def train():
    d = torch.load(args.feats)
    h, lse, base, y, traj, T = d["h"], d["lse"], d["base"], d["y"], d["traj"], d["temp"]
    kind = d["kind"]
    is_cut = torch.tensor([k == "cut" for k in kind])
    n_tr = int(traj.max().item()) + 1
    g = torch.Generator().manual_seed(0)
    val_traj = set(torch.randperm(n_tr, generator=g)[: n_tr // 5].tolist())
    val = torch.tensor([t.item() in val_traj for t in traj])
    D = torch.zeros(h.shape[1], requires_grad=True)
    opt = torch.optim.Adam([D], lr=1e-2)
    w = torch.where(y > 0.5, torch.full_like(y, args.pos_weight), torch.ones_like(y))
    tr = ~val
    for ep in range(args.epochs):
        z = base + (h @ D) / T - lse  # log-odds of sampling </think>
        loss = torch.nn.functional.binary_cross_entropy_with_logits(z[tr], y[tr], weight=w[tr]) + args.l2 * D.pow(2).sum()
        opt.zero_grad()
        loss.backward()
        opt.step()
    with torch.no_grad():
        z = base + (h @ D) / T - lse
        p = torch.sigmoid(z)
        p0 = torch.sigmoid(base - lse)
    for name, m in [("train", tr), ("val", val)]:
        c = m & is_cut
        rnd = m & torch.tensor([k == "random" for k in kind])
        nat = m & torch.tensor([k == "natural" for k in kind])
        print(f"{name}: cut AUC {auc(z[c], y[c]):.3f} (base {auc((base - lse)[c], y[c]):.3f})  "
              f"p_close at cuts: right {p[c & (y > .5)].mean():.3f} wrong {p[c & (y < .5)].mean():.3f} "
              f"(base {p0[c].mean():.4f}) | random pos {p[rnd].mean():.5f} (base {p0[rnd].mean():.5f}) | "
              f"natural end {p[nat].mean():.3f} (base {p0[nat].mean():.3f})")
    torch.save({"D": D.detach(), "args": vars(args)}, args.out)
    print("saved", args.out, "|D|", D.norm().item())


def export():
    import numpy as np
    import gguf
    from safetensors.torch import load_file
    D = torch.load(args.D)["D"] * args.scale
    sd = {}
    for f in glob.glob(os.path.join(args.model, "*.safetensors")):
        sd.update(load_file(f))
    E = sd["model.embed_tokens.weight"].float()  # tied embedding = original output matrix
    E[THINK_END] += D
    r = gguf.GGUFReader(args.gguf_in)
    arch = r.fields["general.architecture"].contents()
    w = gguf.GGUFWriter(args.gguf_out, arch)
    for f in r.fields.values():
        if f.name == "general.architecture" or f.name.startswith("GGUF."):
            continue
        vt = f.types[0]
        st = f.types[-1] if vt == gguf.GGUFValueType.ARRAY else None
        w.add_key_value(f.name, f.contents(), vt, sub_type=st)
    w.add_string("general.halt_module", f"untied </think> output row + learned D, scale {args.scale}")
    q = gguf.quants.quantize(E.numpy().astype(np.float32), gguf.GGMLQuantizationType.Q8_0)
    tensors = [(t.name, t.data, t.tensor_type) for t in r.tensors] + [("output.weight", q, gguf.GGMLQuantizationType.Q8_0)]
    for name, data, tt in tensors:
        w.add_tensor_info(name, data.shape, data.dtype, data.nbytes, tt)
    w.write_header_to_file(); w.write_kv_data_to_file(); w.write_ti_data_to_file()
    for name, data, tt in tensors:
        w.write_tensor_data(data)
    w.close()
    print("wrote", args.gguf_out)


{"extract": extract, "train": train, "export": export}[args.cmd]()