"""The promotion gate, old readout vs corrected readout, on one checkpoint.

Three things were wrong with how CE(q||p) was measured, all of them inflating
it:

  prefix  data.py applied instance values only when train=True, so eval scored
          p on a sequence whose earlier empty cells held the unique solution.
          The model eliminates candidates from that -- correctly -- and was
          then charged for not matching a q that conditions on nothing.
  depth   evaluater.py sets k_budget = K, so one full-depth prediction was
          compared against every stage's q, including stages training far
          fewer slots.
  q       the stored pool stops at coverage (~5 grids, 1.4 by stage 11) and
          deliberately picks rare candidates, so its histogram is flattened
          toward uniform. The gate spends log|S| - H(q) as headroom, and that
          quantity is meaningless at n = 5: measured H rises from 0.70 at
          n = 3 to 0.99 at n = 7 against log|S| = 1.10, still climbing.

Reports the gate statistic under each combination so the contribution of each
correction is visible.

  python probe_gate_fixed.py CKPT STAGE [NPUZ] [BIGQ_NPZ]
"""

import functools
import os
import sys

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)

CKPT = sys.argv[1]
S = int(sys.argv[2])
NPUZ = int(sys.argv[3]) if len(sys.argv) > 3 else 300
BIGQ = (sys.argv[4] if len(sys.argv) > 4
        else "/egr/research-slim/ghoshavr/ckpt_probe/bigq/test_bigq300.npz")
POP = np.array([bin(i).count("1") for i in range(512)], np.int16)

import jax                                                    # noqa: E402
import jax.numpy as jnp                                       # noqa: E402
from jax import random                                        # noqa: E402
from flax.training import checkpoints                         # noqa: E402
from flax import jax_utils                                    # noqa: E402
from train import main as train_main                          # noqa: E402
from train import data, model, trainer, evaluater             # noqa: E402

config = train_main.get_config()
K = int(config.num_latent_slots)
mc = model.TransformerConfig(
    dtype=config.dtype, vocab_size=config.vocab_size, seq_len=config.seq_len,
    num_heads=config.num_heads, num_layers=config.num_layers,
    emb_dim=config.emb_dim, qkv_dim=config.qkv_dim, mlp_dim=config.mlp_dim,
    dropout_rate=config.dropout_rate,
    attention_dropout_rate=config.attention_dropout_rate,
    deterministic=False, num_latent_slots=K,
    inject_latents=bool(int(getattr(config, "recurrent_latent", 1))))
rng = jax.random.PRNGKey(config.seed)
rng, ir, _ = random.split(rng, 3)
rng, dr = jax.random.split(rng)
net = model.TransformerLMHeadModel(mc)
_, iv = jax.jit(net.init_with_output)(
    {"params": ir, "dropout": dr},
    jnp.ones((config.minibatch_size, config.seq_len), jnp.int32),
    jnp.zeros((config.minibatch_size, K, config.emb_dim), mc.dtype),
    jnp.zeros((config.minibatch_size, K), jnp.int32),
    jnp.zeros((config.minibatch_size, K), bool))
st, _ = trainer.get_state(config, net, iv)
rs = checkpoints.restore_checkpoint(os.path.abspath(CKPT), st)
print(f"restored step {int(rs.step)}  stage {S}", flush=True)
state = jax_utils.replicate(rs)
pes = jax.pmap(functools.partial(
    evaluater.eval_step, config=mc.replace(deterministic=True)),
    axis_name="batch")

cur = data.CurriculumState(stage=S, max_stage=K)
ds = data.SudokuDataset(config, train=False, curriculum=cur)
cm = np.load(config.test_cand_masks_path, mmap_mode="r")
z = np.load(BIGQ)
bigq, biggrid = z["counts"], z["grids"]
sidx = S - 1
BS = config.minibatch_size
nd = jax.local_device_count()


def fwd(sq, lv, act, sp):
    def sh(x):
        x = np.asarray(x)
        return x.reshape((nd, x.shape[0] // nd) + x.shape[1:])
    lg, hd, _ = pes(state, sh(sq), sh(lv), sh(sp), sh(act))
    lg, hd = np.array(lg), np.array(hd)
    return lg.reshape(-1, *lg.shape[2:]), hd.reshape(-1, *hd.shape[2:])


def probs(seqs, si3, depth):
    """Digit probabilities with `depth` latent slots active and built."""
    n = len(seqs)
    sq = np.stack(seqs)
    s3 = np.asarray(si3)
    pad = (-n) % BS
    if pad:
        sq = np.concatenate([sq, np.zeros((pad, config.seq_len), sq.dtype)])
        s3 = np.concatenate([s3, np.full(pad, s3[0])])
    out = []
    for b in range(0, len(sq), BS):
        x = sq[b:b + BS].astype(np.int32)
        t3 = s3[b:b + BS]
        sp = t3[:, None] + np.arange(K)[None, :]
        slots = np.arange(K)[None, :]
        af = np.ones((len(x), K), bool)
        bi = np.arange(len(x))
        lv = np.zeros((len(x), K, config.emb_dim), np.float32)
        mk = x.copy()
        for j in range(len(x)):
            mk[j, t3[j] + K:] = 0
        for j in range(depth):
            _, hd = fwd(mk, lv, af & (slots < j), sp)
            lv[:, j] = hd[bi, t3 - 1 + j]
        lg, _ = fwd(x, lv, af & (slots < depth), sp)
        zz = lg - lg.max(-1, keepdims=True)
        p = np.exp(zz)
        p /= p.sum(-1, keepdims=True)
        out.append(p[:, :, 1:10].astype(np.float64))
    return np.concatenate(out)[:n]


# ---- Build every sequence variant once -------------------------------------
R = int(biggrid.shape[2])
rows, meta, sel = [], [], []
for idx in range(min(NPUZ, bigq.shape[0])):
    n5 = int(ds.inst_counts[idx, sidx])
    base = np.asarray(ds.eval_inputs[idx])
    si = int(ds.eval_start_index[idx, 0])
    rows.append(ds.insert_latent_slots(base, si))
    meta.append((idx, "sol", 0))
    for r in range(R):
        g = biggrid[idx, sidx, r]
        rows.append(ds.insert_latent_slots(ds._apply_values(base, g), si))
        meta.append((idx, "big", r))
    lo = int(ds.inst_starts[idx, sidx])
    n5c = max(n5, 0)
    for r in range(min(n5c, 4)):
        rows.append(ds.insert_latent_slots(
            ds.apply_instance(base, idx, sidx, inst_row=lo + r), si))
        meta.append((idx, "pool", r))
    sel.append((idx, si, min(n5c, 4)))

si3 = np.array([3 * int(ds.eval_start_index[m[0], 0]) for m in meta])
store = {}
for depth, tag in ((S, "d"), (K, "K")):
    P = probs(rows, si3, depth)
    for i, (idx, kind, r) in enumerate(meta):
        store.setdefault((tag, idx, kind), {})[r] = P[i]

# ---- Score ------------------------------------------------------------------
# Each combination names (which q, which prefix, which depth).
COMBO = [
    ("OLD  gate as it runs today", "pool5", "sol", "K"),
    ("  + stage depth only", "pool5", "sol", "d"),
    ("  + instance prefix", "pool5", "big", "d"),
    ("NEW  + large-sample q", "big", "big", "d"),
    ("     (large q, solution prefix)", "big", "sol", "d"),
]
agg = {c[0]: dict(ce=0.0, lk=0.0, hq=0.0, n=0.0, pt=0.0, un=0.0, spr=0.0,
                  mass=0.0, qt=0.0, puz=[]) for c in COMBO}

for idx, si, n5c in sel:
    seq = np.asarray(ds.insert_latent_slots(ds.eval_inputs[idx], si))
    bits = np.asarray(cm[idx, sidx], np.int64)
    tb = np.asarray(cm[idx, -1], np.int64)
    q5 = np.zeros((81, 9), np.float64)
    lo = int(ds.inst_starts[idx, sidx])
    n5 = int(ds.inst_counts[idx, sidx])
    if n5 > 0:
        A = np.asarray(ds.instances[lo:lo + n5], np.int64)
        for c in range(81):
            col = A[:, c]
            v = (col >= 1) & (col <= 9)
            np.add.at(q5[c], col[v] - 1, 1)
    qb = bigq[idx, sidx].astype(np.float64)
    per = {c[0]: dict(ce=0.0, lk=0.0, n=0.0) for c in COMBO}
    base0 = 3 * si + K
    for t in range(81 - si):
        vp = base0 + 3 * t + 2
        if vp >= config.seq_len:
            break
        r_, c_ = int(seq[vp - 2]), int(seq[vp - 1])
        if not (0 <= r_ < 9 and 0 <= c_ < 9):
            continue
        cell = r_ * 9 + c_
        k = int(POP[int(bits[cell])])
        if k < 2:
            continue
        inS = np.array([(int(bits[cell]) >> d) & 1 for d in range(9)], bool)
        true = int(tb[cell]).bit_length()
        lk = np.log(k)
        for name, qsrc, psrc, tag in COMBO:
            qc = q5[cell] if qsrc == "pool5" else qb[cell]
            tot = qc.sum()
            if tot <= 0:
                continue
            q = qc / tot
            if psrc == "sol":
                dig = store[(tag, idx, "sol")][0][vp - 1]
            else:
                src = "big" if psrc == "big" else "pool"
                d = store[(tag, idx, src)]
                dig = np.mean([d[r] for r in sorted(d)], axis=0)[vp - 1]
            ce = -float((q * np.log(np.maximum(dig, 1e-12))).sum())
            qn = q[q > 0]
            T = agg[name]
            T["ce"] += ce
            T["lk"] += lk
            T["hq"] += -float((qn * np.log(qn)).sum())
            T["n"] += 1
            T["pt"] += float(dig[true - 1])
            T["qt"] += float(q[true - 1])
            T["un"] += 1.0 / k
            m = float(dig[inS].sum())
            pin = dig[inS] / max(m, 1e-30)
            T["mass"] += m
            T["spr"] += float(-(pin * np.log(np.maximum(pin, 1e-30))).sum()) / lk
            per[name]["ce"] += ce
            per[name]["lk"] += lk
            per[name]["n"] += 1
    for name in per:
        if per[name]["n"] > 0:
            agg[name]["puz"].append(max(
                0.0, per[name]["ce"] / per[name]["n"]
                - per[name]["lk"] / per[name]["n"]))

BAR = 0.05
print(f"\n{len(sel)} held-out puzzles, stage {S}, bar hinge <= {BAR}")
print(f"{'readout':34s} {'q(true)':>8s} {'p(true)':>8s} {'1/|S|':>7s} "
      f"{'spread':>7s} {'H(q)':>7s} {'log|S|':>7s} {'CE':>7s} {'hinge':>7s}  verdict")
for name, _, _, _ in COMBO:
    T = agg[name]
    n = max(T["n"], 1)
    v = np.array(T["puz"])
    h = v.mean() if len(v) else -1.0
    print(f"{name:34s} {T['qt']/n:8.3f} {T['pt']/n:8.3f} {T['un']/n:7.3f} "
          f"{T['spr']/n:7.3f} {T['hq']/n:7.3f} {T['lk']/n:7.3f} "
          f"{T['ce']/n:7.3f} {h:7.4f}  "
          f"{'PASS' if 0 <= h <= BAR else 'FAIL'}")
print("\nhinge = mean over puzzles of max(0, mean CE(q||p) - mean log|S|)")