"""deviant_bed.py — trains the DEVIANT ROSTER candidates on the certified byte
bed. #TAG:deviant_bed #TAG:loss_campaign
One arm per gate-cleared candidate (inventory/DEVIANT_ROSTER.md) + the
gate-refused confidence penalty as the DESIGNATED CONTROL (tests whether gate
refusals predict training reality) + the sparsemax coupling-axis probe on the
addr_head collapse configuration. Certified operating point throughout:
wikitext-2 bytes, block 256, batch 32, 2000 steps, pure Adam 3e-4 wd=0,
fp32/TF32-off, crc32 seeds. Ledgers -> tools/deviant_runs/*.jsonl.

Baselines on this exact protocol (this session): ce 2.4769 mean bpb (3 seeds),
fac_lsh 4.1285, ce_fixedcode 3.8060, fac_none 3.9547.

Run:  python tools/deviant_bed.py --arm <name> --seed N [--steps 2000]
      python tools/deviant_bed.py --list
"""
import json
import math
import os
import sys
import time
import zlib

import torch
import torch.nn as nn
import torch.nn.functional as F


def _root():
    d = os.path.abspath(os.getcwd())
    while True:
        if os.path.exists(os.path.join(d, "MANIFEST.md")):
            return d
        p = os.path.dirname(d)
        if p == d:
            return os.getcwd()
        d = p

ROOT = _root()
if os.path.join(ROOT, "tools") not in sys.path:
    sys.path.insert(0, os.path.join(ROOT, "tools"))

import ar_differentiation_bed as bed                       # noqa: E402
from loss_forms import (dev_geomean_accum, dev_softmax_accum, fac_loss_link,
                        prim_ce, pwa_weights, sparsemax_loss)  # noqa: E402

torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
DEV = "cuda" if torch.cuda.is_available() else "cpu"
if DEV == "cuda":
    torch.cuda.set_per_process_memory_fraction(0.73)

DATA_ROOT = os.environ.get("GEOLIP_DATA", "./data")
RUNS = os.path.join(ROOT, "tools", "deviant_runs")
REF_CKPT = os.path.join(DATA_ROOT, "fac_ckpts", "ce_s0_t2000.pt")

def seed_for(name):
    return zlib.crc32(name.encode()) & 0x7FFFFFFF

ARMS = ("softmax_accum", "geomean_accum", "label_smooth", "focal",
        "sparsemax", "anti_curr", "fac_tanh", "fac_cauchy",
        "conf_penalty_CONTROL", "p4_sparsemax")
FAC_ARMS = ("fac_tanh", "fac_cauchy")


def build(arm, seed):
    torch.manual_seed(seed_for(f"deviant:{arm}:{seed}"))
    model_arm = "addr_head" if arm == "p4_sparsemax" else "addr_msl64"
    lm = bed.ByteLM(model_arm).to(DEV)
    ctx = {}
    if arm in FAC_ARMS:
        lm.head = nn.Identity()                       # forward returns feats
        gp = torch.Generator().manual_seed(seed_for("deviant-fac-frame"))
        ctx["R"] = torch.linalg.qr(
            torch.randn(256, 256, generator=gp))[0][:64].to(DEV)
        gsm = torch.randn(258, 64, generator=gp)      # lsh: box-3 smoothed
        sm = (gsm[:-2] + gsm[1:-1] + gsm[2:]) / 3.0
        ctx["code"] = ((sm > 0).float() * 2 - 1).to(DEV)
    if arm == "anti_curr":
        ref = bed.ByteLM("addr_msl64").to(DEV)
        sd = torch.load(REF_CKPT, map_location="cpu", weights_only=True)
        ref.load_state_dict(sd["state_dict"], strict=True)
        ref.eval()
        for p_ in ref.parameters():
            p_.requires_grad_(False)
        ctx["ref"] = ref
    return lm, ctx


def loss_of(arm, lm, ctx, x, y):
    out = lm(x)
    if arm in FAC_ARMS:
        link = "tanh" if arm == "fac_tanh" else "cauchy"
        return fac_loss_link(out, ctx["R"], ctx["code"][y], link=link).mean()
    if arm == "sparsemax" or arm == "p4_sparsemax":
        return sparsemax_loss(out.reshape(-1, 256), y.reshape(-1)).mean()
    ce_tok = prim_ce(out, y)
    if arm == "softmax_accum":
        return dev_softmax_accum(ce_tok, T=0.5)
    if arm == "geomean_accum":
        return dev_geomean_accum(ce_tok)
    if arm == "label_smooth":
        return F.cross_entropy(out.reshape(-1, 256), y.reshape(-1),
                               label_smoothing=0.1)
    if arm == "focal":
        w = (1 - F.softmax(out, -1).gather(-1, y.unsqueeze(-1))
             .squeeze(-1).detach()) ** 2
        return (w * ce_tok).sum() / w.sum().clamp_min(1e-9)
    if arm == "anti_curr":
        with torch.no_grad():
            pi = F.softmax(ctx["ref"](x), -1).gather(
                -1, y.unsqueeze(-1)).squeeze(-1)
        m = (pi > 0.6).float()
        return (ce_tok * m).sum() / m.sum().clamp_min(1.0)
    if arm == "conf_penalty_CONTROL":
        p = F.softmax(out, -1)
        ent = -(p * p.clamp_min(1e-12).log()).sum(-1).mean()
        return ce_tok.mean() - 0.1 * ent
    raise ValueError(arm)


@torch.no_grad()
def evaluate(arm, lm, ctx, va, g):
    """bpb-of-record (softmax over available scores), decoded acc, vitals."""
    lm.eval()
    tot_ce, tot_ok, n = 0.0, 0, 0
    for _ in range(8):
        x, y = bed._batch(va, 32, 256, DEV, g)
        out = lm(x)
        if arm in FAC_ARMS:
            s = F.normalize(out, dim=-1) @ ctx["R"].t()
            scores = s @ ctx["code"].t()
        else:
            scores = out
        lp = F.log_softmax(scores, -1)
        tot_ce += float(-lp.gather(-1, y.unsqueeze(-1)).sum())
        tot_ok += int((scores.argmax(-1) == y).sum())
        n += y.numel()
    x, _ = bed._batch(va, 8, 256, DEV, g)
    _ = lm(x)
    v = lm.head_addr.vitals(lm.head_proj(lm._last_h).view(
        *lm._last_h.shape[:-1], lm.n_slots, 4)) \
        if getattr(lm, "n_slots", 0) else lm.head_addr.vitals(lm._last_h)
    lm.train()
    return (tot_ce / n) / math.log(2), tot_ok / n, v


def train_arm(arm, seed, steps=2000):
    assert arm in ARMS, f"unknown arm {arm} (see --list)"
    os.makedirs(RUNS, exist_ok=True)
    tr, va = bed._wikitext_bytes(DATA_ROOT)
    lm, ctx = build(arm, seed)
    g = torch.Generator().manual_seed(seed_for(f"deviant-data:{arm}:{seed}"))
    ge = torch.Generator().manual_seed(seed_for("deviant-eval"))
    opt = torch.optim.Adam((p_ for p_ in lm.parameters() if p_.requires_grad),
                           lr=3e-4, weight_decay=0.0)
    t0 = time.time()
    for step in range(steps):
        x, y = bed._batch(tr, 32, 256, DEV, g)
        opt.zero_grad(set_to_none=True)
        L = loss_of(arm, lm, ctx, x, y)
        L.backward()
        opt.step()
        if step == 10 and DEV == "cuda":
            print(f"[{arm} s{seed}] step10 loss {float(L.detach()):.4f} "
                  f"peak {torch.cuda.max_memory_allocated()/2**30:.2f}GB",
                  flush=True)
    bpb, acc, vit = evaluate(arm, lm, ctx, va, ge)
    rec = {"arm": arm, "seed": seed, "steps": steps, "bpb": round(bpb, 4),
           "decoded_acc": round(acc, 4), "vitals": vit,
           "wall_s": round(time.time() - t0, 1),
           "peak_gb": round(torch.cuda.max_memory_allocated() / 2**30, 2)
           if DEV == "cuda" else 0.0}
    out = os.path.join(RUNS, f"{arm}_s{seed}_t{steps}.jsonl")
    with open(out, "a", encoding="utf-8") as f:
        f.write(json.dumps(rec) + "\n")
    print(f"[DONE {arm} s{seed}] bpb {bpb:.4f} acc {acc:.4f} "
          f"usage_ppl {vit['aliveness']['usage_ppl']:.1f} "
          f"win|cos| {vit['win_cos_mean']:.3f} ({rec['wall_s']}s)", flush=True)
    return rec


if __name__ == "__main__":
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("--arm", default=None)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--steps", type=int, default=2000)
    ap.add_argument("--list", action="store_true")
    a, _ = ap.parse_known_args()
    if a.list or not a.arm:
        print("arms:", " ".join(ARMS))
        sys.exit(0)
    train_arm(a.arm, a.seed, a.steps)