"""Architecture validation before any real training.

Checks:
  1. Dataset integrity: answer marginals ~uniform per k (no majority-class
     shortcut); LookupSolver scores 100% (pipeline exactness).
  2. Gradient flow through recurrence: a 1-block model unrolled R=8 gets
     non-zero gradients on every block parameter; loss at init is finite.
  3. Overfitting a tiny fixed training set: both architectures reach ~100%
     train accuracy on 256 fixed examples (gradients + capacity are real).
  4. Interface checks: n_iters variation works, smaller state spaces evaluate
     with the same weights (subset of vocab/positions), shapes correct.

Exit code 0 iff all checks pass. Run: python validate.py
"""

import sys

import time

import numpy as np
import torch

torch.set_num_threads(8)
import torch.nn.functional as F

from ptr.common import (LookupSolver, check_answer_marginals,
                        sample_eval_batch, sample_train_batch)
from ptr.models import PtrTransformer

DEV = "cpu"
S, D = 64, 128


def check_marginals_and_lookup():
    rng = np.random.default_rng(0)
    m = check_answer_marginals(rng, S=S, k_max=8, n=100_000)
    print(f"[1a] answer-marginal TV per k: "
          f"{ {k: round(v, 4) for k, v in m['tv_per_k'].items()} } -> "
          f"{'PASS' if m['pass'] else 'FAIL'}")
    assert m["pass"]

    solver = LookupSolver()
    b = sample_eval_batch(rng, 1000, S, k=8)
    pred = solver.solve(b["f"], b["x"], 8)
    acc = float((pred == b["answer"]).mean())
    print(f"[1b] LookupSolver exact-match on random data: {acc:.4f} -> "
          f"{'PASS' if acc == 1.0 else 'FAIL'}")
    assert acc == 1.0


def check_gradients():
    torch.manual_seed(0)
    model = PtrTransformer(S, D, 4, 512, depth=1, zero_init=True).to(DEV)
    rng = np.random.default_rng(1)
    b = sample_train_batch(rng, 64, S)
    f = torch.as_tensor(b["f"], device=DEV)
    x = torch.as_tensor(b["x"], device=DEV)
    k = torch.as_tensor(b["k"], device=DEV)
    y = torch.as_tensor(b["answer"], device=DEV)
    R = 8
    # Zero-initialised output projections make every block the identity at
    # step 0, so upstream parameters receive no gradient on the very first
    # backward. Warm up a few AdamW steps first, then check that the full
    # R=8 unroll propagates gradients to every block parameter.
    opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
    for _ in range(3):
        opt.zero_grad(set_to_none=True)
        F.cross_entropy(model(f, x, k, n_iters=R), y).backward()
        opt.step()
    opt.zero_grad(set_to_none=True)
    logits = model(f, x, k, n_iters=R)
    loss = F.cross_entropy(logits, y)
    loss.backward()
    bad = [n for n, p in model.named_parameters()
           if p.requires_grad and (p.grad is None or p.grad.abs().sum() == 0)]
    print(f"[2] recurrent unroll R={R}: init loss {loss.item():.4f} "
          f"(uniform would be {np.log(S):.4f}); "
          f"params with zero/missing grad: {bad or 'none'} -> "
          f"{'PASS' if not bad and torch.isfinite(loss) else 'FAIL'}")
    assert not bad and torch.isfinite(loss)


def overfit_tiny(depth, n_iters, steps=2500, lr=1e-3, target=0.99):
    torch.manual_seed(depth * 10 + n_iters)
    model = PtrTransformer(S, D, 4, 512, depth=depth, zero_init=True).to(DEV)
    rng = np.random.default_rng(7)
    fixed = sample_train_batch(rng, 256, S, k_max=8)
    f = torch.as_tensor(fixed["f"], device=DEV)
    x = torch.as_tensor(fixed["x"], device=DEV)
    k = torch.as_tensor(fixed["k"], device=DEV)
    y = torch.as_tensor(fixed["answer"], device=DEV)
    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
    t0 = time.time()
    acc = 0.0
    for i in range(steps):
        opt.zero_grad(set_to_none=True)
        logits = model(f, x, k, n_iters=n_iters)
        loss = F.cross_entropy(logits, y)
        loss.backward()
        opt.step()
        if i % 100 == 0 or i == steps - 1:
            acc = (logits.argmax(-1) == y).float().mean().item()
            print(f"    depth={depth} R={n_iters} step {i}: "
                  f"loss {loss.item():.4f} acc {acc:.3f} "
                  f"({time.time()-t0:.0f}s)", flush=True)
            if acc >= target:
                break
    acc = (logits.argmax(-1) == y).float().mean().item()
    ok = acc >= target
    print(f"[3] overfit tiny set (depth={depth}, R={n_iters}): "
          f"final acc {acc:.3f} -> {'PASS' if ok else 'FAIL'}")
    assert ok


def check_interface():
    torch.manual_seed(3)
    model = PtrTransformer(S, D, 4, 512, depth=4, zero_init=True).to(DEV)
    rng = np.random.default_rng(11)
    for n_iters in (1, 4, 12):
        b = sample_eval_batch(rng, 8, S, k=16)
        f = torch.as_tensor(b["f"], device=DEV)
        x = torch.as_tensor(b["x"], device=DEV)
        k = torch.as_tensor(b["k"], device=DEV)
        logits = model(f, x, k, n_iters=n_iters)
        assert logits.shape == (8, S)
    # smaller state space with the same weights (subset of vocab + positions)
    b = sample_eval_batch(rng, 8, 16, k=64)
    logits = model(torch.as_tensor(b["f"], device=DEV),
                   torch.as_tensor(b["x"], device=DEV),
                   torch.as_tensor(b["k"], device=DEV), n_iters=4)
    # head is sized to the trained S=64; a smaller eval state space S'=16
    # yields 64 logits and must be MASKED to the first S' classes
    assert logits.shape == (8, 64)
    assert logits[:, :16].argmax(-1).shape == (8,)
    print("[4] interface: n_iters 1/4/12 and S=16-subset eval (with head "
          f"masking) all OK (logits {tuple(logits.shape)}) -> PASS")


if __name__ == "__main__":
    check_marginals_and_lookup()
    check_gradients()
    overfit_tiny(depth=4, n_iters=1)
    overfit_tiny(depth=1, n_iters=4)
    overfit_tiny(depth=1, n_iters=8)
    check_interface()
    print("ALL VALIDATION CHECKS PASSED")
    sys.exit(0)