"""Mandatory invariant suite (WP31 brief, Phase 3). Each test maps to one numbered invariant."""
import copy

import pytest
import torch

from prisme.workspace import action_table, rms
from tests.conftest import tiny_config
from prisme.model import PrismeModel


def random_actions(cfg, B, T, R, seed=0, include_stop=True):
    g = torch.Generator().manual_seed(seed)
    hi = cfg.n_actions + (1 if include_stop else 0)
    return torch.randint(0, hi, (B, T, R), generator=g)


# 1. shapes, dtypes, gradients, finiteness
def test_01_shapes_dtypes_grads_finite(model_nonzero, cfg, tokens):
    m = model_nonzero
    m.train()
    acts = random_actions(cfg, 3, 12, 4)
    out = m(tokens, acts)
    assert out["logits"].shape == (3, 12, cfg.vocab_size)
    assert out["logits"].dtype == torch.float32
    assert torch.isfinite(out["logits"]).all()
    out["logits"].sum().backward()
    for n, p in m.named_parameters():
        if n.startswith("controller"):
            continue
        assert p.grad is not None and torch.isfinite(p.grad).all(), n


# 2. exact backbone identity with budget 0
def test_02_budget_zero_identity(model_nonzero, tokens):
    m = model_nonzero
    ref = m.backbone.head(m.backbone(tokens))
    out = m(tokens, actions=None)["logits"]
    assert torch.equal(out, ref)
    out0 = m(tokens, actions=torch.empty(3, 12, 0, dtype=torch.long))["logits"]
    assert torch.equal(out0, ref)


# 3. exact identity with immediate STOP (no edits)
def test_03_stop_identity(model_nonzero, cfg, tokens):
    m = model_nonzero
    ref = m.backbone.head(m.backbone(tokens))
    stop = torch.full((3, 12, 4), cfg.stop_action)
    assert torch.equal(m(tokens, stop)["logits"], ref)


# 4. zero-initialised readout adapter preserves identity even with edits
def test_04_zero_readout_identity(model, cfg, tokens):
    assert torch.count_nonzero(model.readout.U.weight) == 0
    ref = model.backbone.head(model.backbone(tokens))
    acts = random_actions(cfg, 3, 12, 8, include_stop=False)
    out = model(tokens, acts, return_state=True)
    assert (out["S_final"] != out["ctx"].S0).any()  # edits happened
    assert torch.equal(out["logits"], ref)


# 5. future suffix changes do not alter preceding logits (with edits, nonzero readout)
def test_05_causality(model_nonzero, cfg, tokens):
    m = model_nonzero
    acts = random_actions(cfg, 3, 12, 4, include_stop=False)
    a = m(tokens, acts)["logits"]
    t2 = tokens.clone()
    t2[:, 8:] = (t2[:, 8:] + 5) % cfg.vocab_size
    b = m(t2, acts)["logits"]
    assert torch.allclose(a[:, :8], b[:, :8], atol=1e-5, rtol=1e-5)
    assert not torch.allclose(a[:, 8:], b[:, 8:])


# 6. inactive slots exactly unchanged after singles and pairs
def test_06_locality(model_nonzero, cfg, tokens):
    m = model_nonzero
    ctx = m.encode(tokens)
    tbl = action_table(cfg)
    for a in range(cfg.n_actions + 1):
        act = torch.full((3, 12), a)
        S1, info = m.apply_action(ctx.S0, ctx, act, 0, 3)
        slots = [s for s in tbl[a].tolist() if s >= 0]
        for j in range(cfg.n_slots):
            if j in slots:
                assert not torch.equal(S1[:, :, j], ctx.S0[:, :, j]), (a, j)
            else:
                assert torch.equal(S1[:, :, j], ctx.S0[:, :, j]), (a, j)


# 7. backbone / source-bank states unchanged by workspace editing
def test_07_bank_unchanged(model_nonzero, cfg, tokens):
    m = model_nonzero
    ctx = m.encode(tokens)
    snap = {k: getattr(ctx, k).clone() for k in ("H", "K", "V", "z", "S0")}
    S = ctx.S0
    for r in range(4):
        S, _ = m.apply_action(S, ctx, random_actions(cfg, 3, 12, 1, seed=r)[:, :, 0], r, 3 - r)
    for k, v in snap.items():
        assert torch.equal(getattr(ctx, k), v), k


# 8. RMS displacement bound: RMS(S_j^n - S_j^0) <= n_j * eta_max * c
def test_08_rms_displacement_bound(model_nonzero, cfg, tokens):
    m = model_nonzero
    ctx = m.encode(tokens)
    S = ctx.S0
    counts = torch.zeros(3, 12, cfg.n_slots)
    for r in range(8):
        act = random_actions(cfg, 3, 12, 1, seed=100 + r)[:, :, 0]
        S, info = m.apply_action(S, ctx, act, r, 7 - r)
        counts += info["write_mask"].float()
    disp = rms(S - ctx.S0)  # [B,T,M]
    bound = counts * cfg.eta_max * cfg.clip_c
    assert (disp <= bound + 1e-5).all()


# 9. differential correction is zero for identical initial/final states
def test_09_zero_delta_identical_states(model_nonzero, tokens):
    m = model_nonzero
    ctx = m.encode(tokens)
    d = m.readout(ctx.S0, ctx.S0, ctx.z)
    assert torch.equal(d, torch.zeros_like(d))


# 10. the hard action limit cannot be exceeded
def test_10_hard_cap(model_nonzero, cfg, tokens):
    with pytest.raises(AssertionError):
        model_nonzero(tokens, random_actions(cfg, 3, 12, cfg.r_max + 1))
    cost = torch.full((cfg.n_actions,), -1.0)  # negative cost -> controller always wants to edit
    with torch.no_grad():
        model_nonzero.controller.heads.bias[0] = 5.0
    out = model_nonzero.forward_policy(tokens, max_budget=50, cost=cost)
    assert out["actions"].shape[-1] == cfg.r_max
    assert int(out["n_edits"].max()) <= cfg.r_max


# 11. simultaneous pair edits do not depend on lane/iteration order
def test_11_pair_order_independence(model_nonzero, cfg, tokens):
    m = model_nonzero
    ctx = m.encode(tokens)
    pair = torch.full((3, 12), cfg.n_slots)  # PAIR_0_1
    S_a, _ = m.apply_action(ctx.S0, ctx, pair, 0, 0)
    tbl = m.editor.actions.clone()
    tbl[cfg.n_slots] = torch.tensor([1, 0])  # swap lanes
    m.editor.actions = tbl
    S_b, _ = m.apply_action(ctx.S0, ctx, pair, 0, 0)
    assert torch.equal(S_a, S_b)


# 12. valid gradients reach the readout, and then the editor after bootstrap (U != 0)
def test_12_gradient_flow(cfg, tokens):
    torch.manual_seed(0)
    m = PrismeModel(cfg)
    acts = random_actions(cfg, 3, 12, 2, include_stop=False)
    loss = m(tokens, acts)["logits"].logsumexp(-1).mean()
    loss.backward()
    assert m.readout.U.weight.grad is not None and m.readout.U.weight.grad.abs().sum() > 0
    ed = [p.grad for n, p in m.named_parameters() if n.startswith("editor.block")]
    assert all(g is None or g.abs().sum() == 0 for g in ed)  # U = 0 blocks the path, as predicted
    m.zero_grad()
    with torch.no_grad():
        m.readout.U.weight.normal_(std=0.05)
    m(tokens, acts)["logits"].logsumexp(-1).mean().backward()
    ed = [p.grad for n, p in m.named_parameters() if n.startswith("editor.block")]
    assert any(g is not None and g.abs().sum() > 0 for g in ed)


# 13. batched teacher forcing == per-sequence forward (masks are per query)
def test_13_batched_equals_single(model_nonzero, cfg, tokens):
    m = model_nonzero
    acts = random_actions(cfg, 3, 12, 3, include_stop=False)
    full = m(tokens, acts)["logits"]
    for b in range(3):
        single = m(tokens[b:b + 1], acts[b:b + 1])["logits"]
        assert torch.allclose(full[b:b + 1], single, atol=1e-5, rtol=1e-5)


# 14. no cross-request contamination
def test_14_no_cross_request_contamination(model_nonzero, cfg, tokens):
    m = model_nonzero
    acts = random_actions(cfg, 3, 12, 3, include_stop=False)
    a = m(tokens, acts, return_state=True)
    t2 = tokens.clone()
    t2[1] = (t2[1] + 3) % cfg.vocab_size
    b = m(t2, acts, return_state=True)
    assert torch.allclose(a["logits"][0], b["logits"][0], atol=1e-5)
    assert torch.allclose(a["logits"][2], b["logits"][2], atol=1e-5)
    assert torch.equal(a["S_final"][0], b["S_final"][0])


# 15. determinism under fixed seed
def test_15_determinism(cfg, tokens):
    outs = []
    for _ in range(2):
        torch.manual_seed(99)
        m = PrismeModel(cfg)
        with torch.no_grad():
            m.readout.U.weight.normal_(std=0.05)
        outs.append(m(tokens, random_actions(cfg, 3, 12, 4))["logits"])
    assert torch.equal(outs[0], outs[1])


# controller features are target-free and never execute edits
def test_controller_target_free(model_nonzero, cfg, tokens):
    m = model_nonzero
    ctx = m.encode(tokens)
    step = torch.zeros(3, 12, dtype=torch.long)
    p1 = m.controller(ctx.z, ctx.S0, step, step + 3)
    t2 = tokens.clone(); t2[:, -1] = (t2[:, -1] + 1) % cfg.vocab_size   # change only the last token
    ctx2 = m.encode(t2)
    p2 = m.controller(ctx2.z, ctx2.S0, step, step + 3)
    assert torch.allclose(p1["g_hat"][:, :-1], p2["g_hat"][:, :-1], atol=1e-5)
    assert p1["g_hat"].shape == (3, 12, cfg.n_actions)


def test_dense_writes_ablation_touches_all_slots(tokens):
    cfg = tiny_config(dense_writes=True)
    torch.manual_seed(0)
    m = PrismeModel(cfg)
    ctx = m.encode(tokens)
    S1, info = m.apply_action(ctx.S0, ctx, torch.zeros(3, 12, dtype=torch.long), 0, 0)
    assert (S1 != ctx.S0).any(-1).all()
    S2, _ = m.apply_action(ctx.S0, ctx, torch.full((3, 12), cfg.stop_action), 0, 0)
    assert torch.equal(S2, ctx.S0)


def test_controller_h_entropy_features_target_free(tokens):
    cfg = tiny_config(ctrl_input="h", ctrl_use_entropy=True)
    torch.manual_seed(0)
    m = PrismeModel(cfg).eval()
    ctx = m.encode(tokens)
    assert ctx.ctrl_in.shape[-1] == cfg.d_model and ctx.entropy.shape == tokens.shape
    step = torch.zeros(3, 12, dtype=torch.long)
    p1 = m.controller(ctx.ctrl_in, ctx.S0, step, step + 3, ctx.entropy)
    t2 = tokens.clone(); t2[:, -1] = (t2[:, -1] + 1) % cfg.vocab_size
    ctx2 = m.encode(t2)
    p2 = m.controller(ctx2.ctrl_in, ctx2.S0, step, step + 3, ctx2.entropy)
    assert torch.allclose(p1["g_hat"][:, :-1], p2["g_hat"][:, :-1], atol=1e-5)
    cost = torch.full((cfg.n_actions,), 0.05)
    out = m.forward_policy(tokens, 4, cost)
    assert out["logits"].shape == (3, 12, cfg.vocab_size)


def test_ablation_configs_locality_and_identity(tokens):
    for kw in [dict(n_slots=4), dict(n_slots=8), dict(n_slots=32), dict(n_editor_blocks=1), dict(n_editor_blocks=4), dict(pair_scheme="offset")]:
        cfg = tiny_config(**kw)
        torch.manual_seed(0)
        m = PrismeModel(cfg).eval()
        with torch.no_grad():
            m.readout.U.weight.normal_(std=0.05)
        ref = m.backbone.head(m.backbone(tokens))
        assert torch.equal(m(tokens, None)["logits"], ref)
        ctx = m.encode(tokens)
        tbl = action_table(cfg)
        pair = cfg.n_slots  # first pair action
        S1, info = m.apply_action(ctx.S0, ctx, torch.full((3, 12), pair), 0, 3)
        slots = [s for s in tbl[pair].tolist() if s >= 0]
        for j in range(cfg.n_slots):
            assert torch.equal(S1[:, :, j], ctx.S0[:, :, j]) == (j not in slots)
        if kw.get("pair_scheme") == "offset":
            assert slots == [0, cfg.n_slots // 2]