"""Per-prompt gate: probe on the k=0 workspace state predicts 'will looping help?'

1. Features: L30 residual at the last prompt position (plain forward), MBPP
   train items; labels easy(0)/hard(1) from the STaR pass (free supervision).
2. Logistic probe (d=1536 -> 1), class-balanced.
3. Gated eval on MBPP test: predicted-easy -> k=0, predicted-hard -> k=4 with
   the dedicated loop adapter. Reports gated pass@1 vs uniform-k references.
"""

import json
import sys
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import torch

from loop_common import BandLooper, MergeAdapter
from prep_mbpp import DIRECT_SUFFIX, extract_code, mbpp_prompt, run_tests

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

OUT = Path(__file__).resolve().parent.parent / "results-loop"
LAYER = 30
K_HARD = 4


@torch.no_grad()
def collect_states(model, tok, items, batch=16):
    feats = []
    for i in range(0, len(items), batch):
        chunk = items[i : i + batch]
        enc = tok([mbpp_prompt(tok, it, DIRECT_SUFFIX) for it in chunk],
                  return_tensors="pt", padding=True,
                  add_special_tokens=False).to("cuda")
        with ResidualCapture(model, layers=[LAYER]) as cap:
            model(**enc, use_cache=False, logits_to_keep=1)
        h = cap.acts[LAYER]  # (B, T, d); left padding -> last pos is prompt end
        feats.append(h[:, -1].float().cpu())
    return torch.cat(feats)


def fit_probe(X, y, epochs=300, lr=0.05):
    mu, sd = X.mean(0), X.std(0) + 1e-6
    Xn = (X - mu) / sd
    w = torch.zeros(X.shape[1], requires_grad=True)
    b = torch.zeros(1, requires_grad=True)
    opt = torch.optim.Adam([w, b], lr=lr)
    pos_w = (y == 0).sum() / max(1, (y == 1).sum())
    for _ in range(epochs):
        z = Xn @ w + b
        loss = torch.nn.functional.binary_cross_entropy_with_logits(
            z, y.float(), pos_weight=pos_w)
        opt.zero_grad()
        loss.backward()
        opt.step()
    return w.detach(), b.detach(), mu, sd


def main():
    model, tok = load_model(dtype=torch.bfloat16)
    tok.padding_side = "left"
    looper = BandLooper(model)
    adapter = MergeAdapter().cuda()
    adapter.load_state_dict(torch.load(OUT / "adapter_code_e399.pt",
                                       map_location="cuda"))

    data = json.load(open(OUT / "mbpp_data.json"))
    train = [it for it in data if it["split"] == "train"
             and it["label"] in ("easy", "hard")]
    test = [it for it in data if it["split"] == "test"][:250]

    print(f"collecting states: {len(train)} train, {len(test)} test", flush=True)
    Xtr = collect_states(model, tok, train)
    ytr = torch.tensor([it["label"] == "hard" for it in train]).long()
    Xte = collect_states(model, tok, test)

    w, b, mu, sd = fit_probe(Xtr, ytr)
    ptr = torch.sigmoid(((Xtr - mu) / sd) @ w + b)
    acc_tr = ((ptr > 0.5).long() == ytr).float().mean()
    pte = torch.sigmoid(((Xte - mu) / sd) @ w + b)
    pred_hard = (pte > 0.5).tolist()
    yte = [it["label"] == "hard" for it in test]
    tp = sum(p and t for p, t in zip(pred_hard, yte))
    print(f"probe: train_acc={acc_tr:.3f}  test: pred_hard="
          f"{sum(pred_hard)} (true hard={sum(yte)}, tp={tp})", flush=True)

    # gated generation: k=0 for predicted-easy, K_HARD for predicted-hard
    codes = [None] * len(test)
    for k, sel in ((0, [i for i, ph in enumerate(pred_hard) if not ph]),
                   (K_HARD, [i for i, ph in enumerate(pred_hard) if ph])):
        for i in range(0, len(sel), 8):
            idxs = sel[i : i + 8]
            chunk = [test[j] for j in idxs]
            enc = tok([mbpp_prompt(tok, it, DIRECT_SUFFIX) for it in chunk],
                      return_tensors="pt", padding=True,
                      add_special_tokens=False).to("cuda")
            gen = looper.generate_frozen_prompt(
                adapter, tok, enc["input_ids"], k, max_new_tokens=220,
                attention_mask=enc["attention_mask"])
            for jj, j in enumerate(idxs):
                codes[j] = extract_code(
                    tok.decode(gen[jj, enc["input_ids"].shape[1]:],
                               skip_special_tokens=True))
    with ThreadPoolExecutor(8) as ex:
        oks = list(ex.map(lambda ci: run_tests(ci[0], ci[1]),
                          zip(codes, test)))
    acc = sum(oks) / len(test)
    by = {}
    for it, ok in zip(test, oks):
        d = by.setdefault(it["label"], [0, 0])
        d[0] += ok
        d[1] += 1
    print(f"GATED pass@1={acc:.3f}  "
          f"by_label={ {l: round(c/n,3) for l,(c,n) in by.items()} }", flush=True)
    json.dump({"acc": acc,
               "by_label": {l: c / n for l, (c, n) in by.items()},
               "probe_train_acc": float(acc_tr),
               "pred_hard_test": int(sum(pred_hard))},
              open(OUT / "eval_gated.json", "w"), indent=1)
    print("wrote", OUT / "eval_gated.json")


if __name__ == "__main__":
    main()