"""Item 30: the carried state as a metacognitive signal (Nils's design).

The program's evidence: the carry's inference-time value is deferral /
completion-state (430 divergence; drop-bucket p=0.0094), and gate quality
is worth ~9.6 points with classifier quality binding (item 20). This
exploits both: a small answer-readiness head g(s) on the carried L30
state at scratchpad line boundaries. The frozen model reads nothing new
— g's decisions reach it as tokens.

Stages:
  harvest --split train|test [--feedforward]
      Roll out rung-A generation (k=2, p by label), teacher-forced replay
      to capture carried states at line boundaries, then fork a forced
      "Answer: " completion per boundary and label it by correctness.
      Saves metacog_<split>[_ff].pt: states (N,1536), labels, item idx,
      boundary index, natural-end correctness.
  head
      Train class-balanced logistic + 2-layer MLP on train states; report
      test AUC for carry and (if present) FF states.
  sweep
      Offline theta sweep on the test LUT: answer at first boundary with
      g >= theta, else natural end. Accuracy / mean boundary / oracle.
"""
import argparse
import json
import os
import sys
import time
from pathlib import Path

import torch
import torch.nn.functional as F

from carry_common import PAUSE_ID, _pos_ids, carry_logits, generate_carry_c
from loop_common import (BandLooper, MergeAdapter, chat_prompt,
                         DIRECT_SUFFIX, last_number, num_eq)

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
from jlens.core import load_model  # noqa: E402

OUT = Path(os.environ.get("LOOP_OUT",
    Path(__file__).resolve().parent.parent / "results-loop"))
K = 2
P_BY_LABEL = {"easy": 2, "hard": 6, "drop": 6}
NL_MAX_BOUND = 6


def newline_positions(tok, gen_ids, start):
    """Positions (absolute) of '\n' tokens within the generated span."""
    out = []
    for j, t in enumerate(gen_ids):
        if "\n" in tok.decode([t]):
            out.append(start + j)
    return out


@torch.no_grad()
def harvest(args, model, tok, looper, adapter):
    items = [it for it in json.load(open(OUT / "star_data.json"))
             if it["split"] == args.split]
    if args.split == "test":
        items = items[:256]
    if args.limit:
        items = items[:args.limit]
    rows = []
    t0 = time.time()
    for bi in range(0, len(items), args.batch):
        chunk = items[bi:bi + args.batch]
        p = max(P_BY_LABEL[it["label"]] for it in chunk)
        enc = tok([chat_prompt(tok, it["question"], DIRECT_SUFFIX)
                   for it in chunk], return_tensors="pt", padding=True,
                  add_special_tokens=False).to("cuda")
        gen = generate_carry_c(looper, adapter, tok, enc["input_ids"],
                               enc["attention_mask"], K, p,
                               max_new_tokens=96,
                               feedforward=args.feedforward)
        n_pre = enc["input_ids"].shape[1] + p
        # teacher-forced replay of each rollout to capture carried states
        for b, it in enumerate(chunk):
            seq = gen[b]
            # strip left padding
            keep = seq != (tok.pad_token_id or 0)
            keep[enc["attention_mask"].shape[1]:] = True
            seq = seq[keep]
            n_prompt = int(enc["attention_mask"][b].sum())
            # trim at end-of-turn
            gen_span = seq[n_prompt + p:]
            eot = tok.convert_tokens_to_ids("<end_of_turn>")
            ends = (gen_span == eot).nonzero()
            g_end = int(ends[0, 0]) if ends.numel() else len(gen_span)
            gen_span = gen_span[:g_end]
            full = seq[: n_prompt + p + g_end][None].cuda()
            msk = torch.ones_like(full)
            _, S = carry_logits(looper, adapter, full, msk,
                                torch.tensor([n_prompt], device="cuda"), K,
                                feedforward=args.feedforward,
                                return_states=True)
            text = tok.decode(gen_span, skip_special_tokens=True)
            nat_ok = bool(num_eq(last_number(text), it["gold"]))
            nls = newline_positions(tok, gen_span.tolist(), n_prompt + p)
            nls = [x for x in nls if x < full.shape[1] - 2][:NL_MAX_BOUND]
            for k_b, pos in enumerate(nls):
                # forced answer from this prefix
                pre = seq[: pos + 1].cuda()
                fa = tok("Answer: ", add_special_tokens=False,
                         return_tensors="pt")["input_ids"][0].cuda()
                inp = torch.cat([pre, fa])[None].cuda()
                am = torch.ones_like(inp)
                g2 = generate_carry_c(looper, adapter, tok, inp, am, 0, 0,
                                      max_new_tokens=8,
                                      feedforward=args.feedforward)
                ans = tok.decode(g2[0, inp.shape[1]:],
                                 skip_special_tokens=True)
                ok = bool(num_eq(last_number("Answer: " + ans), it["gold"]))
                rows.append({"idx": it["idx"], "label": it["label"],
                             "bound": k_b, "ok": ok, "nat_ok": nat_ok,
                             "state": S[0, pos].float().cpu()})
        if bi % (args.batch * 8) == 0:
            print(f"[{bi+len(chunk)}/{len(items)}] rows={len(rows)} "
                  f"({time.time()-t0:.0f}s)", flush=True)
    states = torch.stack([r.pop("state") for r in rows])
    tag = f"metacog_{args.split}" + ("_ff" if args.feedforward else "")
    torch.save({"states": states, "rows": rows}, OUT / f"{tag}.pt")
    pos_rate = sum(r["ok"] for r in rows) / max(1, len(rows))
    print(f"saved {len(rows)} boundary rows -> {tag}.pt "
          f"(answer-ready rate {pos_rate:.3f})", flush=True)


def auc(scores, labels):
    order = sorted(range(len(scores)), key=lambda i: scores[i])
    r = {}
    for rank, i in enumerate(order):
        r[i] = rank + 1
    pos = [i for i, l in enumerate(labels) if l]
    neg = [i for i, l in enumerate(labels) if not l]
    if not pos or not neg:
        return float("nan")
    s = sum(r[i] for i in pos)
    return (s - len(pos) * (len(pos) + 1) / 2) / (len(pos) * len(neg))


def train_head(tr_s, tr_y, dev="cuda", hidden=64, steps=400, seed=0):
    torch.manual_seed(seed)
    d = tr_s.shape[1]
    net = torch.nn.Sequential(
        torch.nn.LayerNorm(d), torch.nn.Linear(d, hidden),
        torch.nn.GELU(), torch.nn.Linear(hidden, 1)).to(dev)
    opt = torch.optim.AdamW(net.parameters(), lr=1e-3, weight_decay=0.01)
    y = tr_y.float().to(dev)
    w = torch.where(y > 0, (1 - y.mean()) / y.mean().clamp(min=1e-3),
                    torch.ones_like(y))
    x = tr_s.to(dev)
    for i in range(steps):
        logit = net(x).squeeze(-1)
        loss = F.binary_cross_entropy_with_logits(logit, y, weight=w)
        opt.zero_grad(); loss.backward(); opt.step()
    return net


def head(args):
    def load(tag):
        d = torch.load(OUT / f"{tag}.pt", map_location="cpu")
        y = torch.tensor([r["ok"] for r in d["rows"]])
        return d["states"], y, d["rows"]
    results = {}
    for name, tr_tag, te_tag in [
            ("carry", "metacog_train", "metacog_test"),
            ("ff", "metacog_train_ff", "metacog_test_ff")]:
        if not (OUT / f"{tr_tag}.pt").exists():
            continue
        tr_s, tr_y, _ = load(tr_tag)
        te_s, te_y, te_rows = load(te_tag)
        net = train_head(tr_s, tr_y)
        with torch.no_grad():
            sc = torch.sigmoid(net(te_s.cuda()).squeeze(-1)).cpu()
        a = auc(sc.tolist(), te_y.tolist())
        results[name] = a
        print(f"{name}: test AUC={a:.3f} (n={len(te_y)}, "
              f"pos-rate={te_y.float().mean():.3f})", flush=True)
        if name == "carry":
            torch.save({"net": net.state_dict(),
                        "scores": sc, }, OUT / "metacog_head.pt")
            for r, s_ in zip(te_rows, sc.tolist()):
                r["g"] = s_
            json.dump(te_rows, open(OUT / "metacog_lut.json", "w"))
    json.dump(results, open(OUT / "metacog_auc.json", "w"))


def sweep(args):
    rows = json.load(open(OUT / "metacog_lut.json"))
    by_item = {}
    for r in rows:
        by_item.setdefault(r["idx"], []).append(r)
    n = len(by_item)
    nat = sum(1 for rs in by_item.values() if rs[0]["nat_ok"]) / n
    print(f"items={n}  natural-end accuracy={nat:.3f}")
    print("theta  acc    mean-bound  answered-early")
    for th in (0.5, 0.6, 0.7, 0.8, 0.85, 0.9, 0.95, 0.99):
        hits, bsum, early = 0, 0, 0
        for rs in by_item.values():
            rs = sorted(rs, key=lambda r: r["bound"])
            fired = next((r for r in rs if r["g"] >= th), None)
            if fired is not None:
                hits += fired["ok"]; bsum += fired["bound"] + 1; early += 1
            else:
                hits += rs[0]["nat_ok"]; bsum += len(rs) + 1
        print(f"{th:.2f}   {hits/n:.3f}  {bsum/n:.2f}        {early}/{n}")
    oracle, obsum = 0, 0
    for rs in by_item.values():
        best = max([r["ok"] for r in rs] + [rs[0]["nat_ok"]])
        oracle += best
        obsum += min([r["bound"] + 1 for r in rs if r["ok"]],
                     default=len(rs) + 1)
    print(f"ORACLE stop: acc={oracle/n:.3f}  mean-bound={obsum/n:.2f}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("stage", choices=("harvest", "head", "sweep"))
    ap.add_argument("--split", default="train")
    ap.add_argument("--feedforward", action="store_true")
    ap.add_argument("--adapter", default=None)
    ap.add_argument("--batch", type=int, default=8)
    ap.add_argument("--limit", type=int, default=0)
    args = ap.parse_args()
    if args.stage == "harvest":
        model, tok = load_model(dtype=torch.bfloat16)
        tok.padding_side = "left"
        looper = BandLooper(model)
        adapter = MergeAdapter(
            d=model.config.get_text_config().hidden_size).cuda()
        default = OUT / ("adapter_carrycot_ff_e400.pt" if args.feedforward
                         else "adapter_carrycot_e400.pt")
        adapter.load_state_dict(torch.load(args.adapter or default,
                                           map_location="cuda"))
        adapter.eval()
        harvest(args, model, tok, looper, adapter)
    elif args.stage == "head":
        head(args)
    else:
        sweep(args)


if __name__ == "__main__":
    main()