"""Focused invariant tests for the PNS experiment (pragmatic set).

Run: python3 -m pytest tests/ -q
"""
import sys
from pathlib import Path

import numpy as np
import pytest
import torch

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from pns.common import EVAL_METADATA_KEYS, MODEL_INPUT_KEYS, SUPERVISION_KEYS
from pns.world.reducer import CAP, S_GOAL, S_OBS, S_SELF, ExactReducer, N_PTR_SLOTS
from pns.world.schema import ET, Mode
from pns.world.sim import generate_lifetime

SEED = 4242


@pytest.fixture(scope="module")
def lt():
    return generate_lifetime(SEED, 256)


def test_determinism(lt):
    assert generate_lifetime(SEED, 256)["digest"] == lt["digest"]


def test_gold_certified_live(lt):
    """Every pointer gold rid was live (in the bounded store) at its question."""
    red = ExactReducer()
    for ev in lt["events"]:
        red.step(ev)
        if ev.ptr_gold_rid >= 0:
            assert ev.ptr_gold_rid in red.records
        for rid in ev.op_arg_rids:
            assert rid in red.records


def test_bounded_stores(lt):
    red = ExactReducer()
    for ev in lt["events"]:
        red.step(ev)
        for s in (S_OBS, S_SELF, S_GOAL):
            assert len(red.by_store[s]) <= CAP[s]
        # questions never write records
        if ev.etype == ET.QUESTION:
            assert not red.last_adds


def test_slot_stability(lt):
    """A record keeps its slot for its whole live span."""
    red = ExactReducer()
    seen: dict[int, int] = {}
    for ev in lt["events"]:
        red.step(ev)
        for rid in red.records:
            slot = red.slot_index(rid)
            if rid in seen:
                assert seen[rid] == slot
            seen[rid] = slot


def test_pack_replay_parity(tmp_path, lt):
    """The packed per-event deltas reproduce the reducer's live sets and the
    packed ptr_gold_slot equals the reducer's slot index - the single-source-
    of-truth gate between generation and every training/eval consumer."""
    from pns.data import tokenize
    from pns.data.pack import pack_lifetimes

    tok = tokenize.load_tokenizer()
    out = tmp_path / "t.npz"
    pack_lifetimes([generate_lifetime(SEED, 256)], tok, out)
    z = np.load(out)
    red = ExactReducer()
    live = {}
    rlo = 0
    for e, ev in enumerate(lt["events"]):
        red.step(ev)
        for rid in z["evict_rid"][z["evict_off"][e]:z["evict_off"][e + 1]]:
            live.pop(int(z["rec_slot"][rlo + int(rid)]), None)
        for rid in z["add_rid"][z["add_off"][e]:z["add_off"][e + 1]]:
            live[int(z["rec_slot"][rlo + int(rid)])] = int(rid)
        true_live = {red.slot_index(r) for r in red.records}
        assert set(live) == true_live, f"live set mismatch at event {e}"
        if ev.ptr_gold_rid >= 0:
            assert int(z["ptr_gold_slot"][e]) == red.slot_index(ev.ptr_gold_rid)


def test_firewall_disjoint():
    assert not (MODEL_INPUT_KEYS & SUPERVISION_KEYS)
    assert not (MODEL_INPUT_KEYS & EVAL_METADATA_KEYS)
    assert not (SUPERVISION_KEYS & EVAL_METADATA_KEYS)


def test_deadline_balance():
    from collections import Counter
    g = Counter()
    for s in range(5000, 5120):
        for ev in generate_lifetime(s, 256)["events"]:
            if ev.etype == ET.QUESTION and ev.family == 7:
                g[ev.gold_value] += 1
    frac = max(g.values()) / sum(g.values())
    assert frac < 0.60, dict(g)


# ------------------------------------------------------------------- model
@pytest.fixture(scope="module")
def model():
    from pns.model.pnsr import PNSR, PNSRConfig
    torch.manual_seed(0)
    return PNSR(PNSRConfig()).eval()


def _dummy(B=2, T=48):
    tok = torch.randint(1, 8000, (B, T))
    tok[:, 40:] = 0
    rec = torch.randn(B, N_PTR_SLOTS, 384)
    live = torch.zeros(B, N_PTR_SLOTS, dtype=torch.bool)
    live[:, :17] = True
    return tok, rec, live


def test_noop_bitwise(model):
    tok, rec, live = _dummy()
    st = model.initial_state(2, "cpu") + 0.31
    s2, _ = model.step(st, tok, torch.tensor([int(ET.NOOP)] * 2),
                       torch.tensor([4, 4]), rec, live)
    assert torch.equal(st, s2)


def test_typed_eligibility(model):
    """A QUESTION may write discourse/workspace/temporal but must leave the
    entity/belief/goal banks bitwise untouched."""
    tok, rec, live = _dummy()
    st = model.initial_state(2, "cpu") + 0.11
    s2, _ = model.step(st, tok, torch.tensor([int(ET.QUESTION)] * 2),
                       torch.tensor([4, 4]), rec, live)
    ent_belief_goal = slice(0, 40 + 28 + 12)
    assert torch.equal(st[:, ent_belief_goal], s2[:, ent_belief_goal])
    assert not torch.equal(st, s2)  # something else did change


def test_pointer_dead_masked(model):
    tok, rec, live = _dummy()
    st = model.initial_state(2, "cpu")
    _, out = model.step(st, tok, torch.tensor([0, 0]), torch.tensor([4, 4]), rec, live)
    assert torch.isinf(out["ptr"][:, 17:]).all()
    assert torch.isfinite(out["ptr"][:, :17]).all()


def test_param_match():
    from pns.model.pnsr import PNSR, PNSRConfig
    from pns.model.tx import TX, TXConfig
    np_ = sum(p.numel() for p in PNSR(PNSRConfig()).parameters())
    nt = sum(p.numel() for p in TX(TXConfig()).parameters())
    assert abs(np_ - nt) / max(np_, nt) < 0.03


def test_loader_live_matches_pack(tmp_path):
    from pns.data import tokenize
    from pns.data.pack import pack_lifetimes
    from pns.data.view import Shard
    tok = tokenize.load_tokenizer()
    out = tmp_path / "t.npz"
    pack_lifetimes([generate_lifetime(SEED + 1, 256)], tok, out)
    sh = Shard(out)
    lv = sh.lifetime(0)
    red = ExactReducer()
    for ev_true, ev_view in zip(generate_lifetime(SEED + 1, 256)["events"], lv):
        red.step(ev_true)
        want = {red.slot_index(r) for r in red.records}
        got = {s for s in range(N_PTR_SLOTS) if ev_view.live_slots[s] >= 0}
        assert want == got
        if ev_view.mode_gold == int(Mode.ANSWER_POINTER):
            assert ev_view.ptr_gold_slot in got