"""One validation puzzle: clues, unique-solution target, stage-1 S, and
what the current metrics mean for a couple of cells.

No model forward pass. The softmax rows are reconstructed to match the
latest eval numbers (out*~0.005, kl*~0.45, spread*~0.92) so the user can
see the shape, not a fake decode.
"""
import numpy as np

PUZ = ("/scratch/users/gatmiry/llm-reasoning-logic-puzzles/sudoku-code/"
       "datasets/test_sudoku_puzzles.npy")
CAND = "/tmp/sudoku_s12/test_cand_masks.npy"
INST = "/tmp/sudoku_superposition"


def digits(m):
    return [d for d in range(1, 10) if int(m) & (1 << (d - 1))]


def main():
    raw = np.load(PUZ, mmap_mode="r")
    row = np.asarray(raw[0]).astype(np.int64)
    si = int(row[0])
    rest = row[1:]
    seq = np.delete(rest, np.arange(81) * 4 + 3)
    tr = seq.reshape(81, 3)
    masks = np.load(CAND, mmap_mode="r")
    Sbits = np.asarray(masks[0, 0])
    starts = np.load(f"{INST}/test_starts.npy")
    counts = np.load(f"{INST}/test_counts.npy")
    A = np.load(f"{INST}/test_assignments.npy", mmap_mode="r")
    n = int(counts[0, 0])
    blk = np.asarray(A[starts[0, 0]:starts[0, 0] + n])

    print(f"val puzzle 0: {si} clues, {81 - si} empty, {n} stored instances")
    print("clues (input, fixed):")
    print(" ", " ".join(f"({r},{c})={v}" for r, c, v in tr[:si][:10]), "...")
    print("unique-solution output (NOT what val_acc scores):")
    print(" ", " ".join(f"({r},{c})={v}" for r, c, v in tr[si:si + 8]), "...")

    print("\nthree empty cells:")
    print(f"  {'cell':>8}  {'S':<16}{'unique':>7}  instance digits")
    shown = 0
    cells = []
    for r, c, v in tr[si:]:
        cell = int(r) * 9 + int(c)
        S = digits(Sbits[cell])
        if len(S) < 2:
            continue
        inst = [int(x) for x in blk[:, cell]]
        print(f"  ({int(r)},{int(c)})     {str(S):<16}{int(v):>7}  {inst}")
        cells.append((cell, S, int(v), inst))
        shown += 1
        if shown == 3:
            break

    # Reconstruct a typical current-run softmax: 0.5% outside S, KL~0.45 inside.
    print("\nwhat the THREE printed metrics look at (teacher-forced softmax,")
    print("not the greedy token). A cell with S={4,5}:")
    print("  unique-solution target would be one digit, say 4")
    print("  training instance target is also one digit (4 or 5)")
    print("  eval target for these metrics is the SET {4,5}")
    print()
    print("  typical model p at this cell right now:")
    print("    digit  1  2  3    4     5    6  7  8  9")
    print("    p      0  0  0  0.875 0.120  0  0  0  0.005   (tiny leak)")
    print()
    print("  mass*  = p(4)+p(5)           = 0.995")
    print("  val_acc = out* = 1-mass*     = 0.005   <-- printed val_acc")
    print("  q      = p restricted to S   = (0.88, 0.12)")
    print("  kl*    = KL(uniform || q)    = 0.45    <-- model prefers 4")
    print("  leak   = log(1/mass*)        = 0.005")
    print("  excess*= leak + kl*          = 0.455   <-- almost all from kl*")
    print()
    print("  if the model were uniform over S and leaked nothing:")
    print("    p(4)=p(5)=0.5, rest 0  ->  val_acc=0  kl*=0  excess*=0")


if __name__ == "__main__":
    main()