#!/usr/bin/env python3
"""Validation gates for the colony harness — run before trusting any colony result.

1. K=1, identity standpoint, mean aggregation must match TRM *exactly* (same
   weights, same inputs -> same logits). If this fails, the colony is not a
   generalization of TRM and no comparison against the 84.5% anchor is valid.
2. K>1 must run forward and backward in every blackboard/aggregate/halt mode.
3. Clonal colonies must not grow parameters with K (only role embeddings may).
4. Standpoint transforms must round-trip: rho^-1(rho(x)) == x.

Usage: DSEM_DEVICE=cpu python scripts/check_colony.py
"""

import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
os.environ.setdefault("DSEM_DEVICE", "cpu")

import torch  # noqa: E402

from dsem.models.recursive_reasoning.colony import (  # noqa: E402
    ColonyModel_ACTV1, build_standpoints, invert_perm,
)
from dsem.models.recursive_reasoning.trm import TinyRecursiveReasoningModel_ACTV1  # noqa: E402

BATCH, SEQ, VOCAB = 4, 81, 11


def base_cfg(**over):
    cfg = dict(
        batch_size=BATCH, seq_len=SEQ, vocab_size=VOCAB, num_puzzle_identifiers=1,
        puzzle_emb_ndim=64, puzzle_emb_len=16,
        H_cycles=2, L_cycles=2, H_layers=0, L_layers=2,
        hidden_size=64, expansion=4, num_heads=1, pos_encodings="none",
        halt_max_steps=2, halt_exploration_prob=0.1, forward_dtype="float32",
        mlp_t=True, no_ACT_continue=True, causal=False,
    )
    cfg.update(over)
    return cfg


def make_batch():
    g = torch.Generator().manual_seed(0)
    return {
        "inputs": torch.randint(1, VOCAB, (BATCH, SEQ), generator=g),
        "labels": torch.randint(1, VOCAB, (BATCH, SEQ), generator=g),
        "puzzle_identifiers": torch.zeros(BATCH, dtype=torch.int32),
    }


def gate_1_equivalence():
    batch = make_batch()
    torch.manual_seed(0)
    trm = TinyRecursiveReasoningModel_ACTV1(base_cfg())
    torch.manual_seed(0)
    col = ColonyModel_ACTV1(base_cfg(num_agents=1, blackboard="latent",
                                     standpoints="none", aggregate="mean", halt_rule="boss"))

    # transplant weights: colony's `level` corresponds to TRM's `L_level`
    mapping = {k.replace("inner.L_level", "inner.level"): v for k, v in trm.state_dict().items()}
    mapping = {k.replace("inner.H_init", "inner.y_init").replace("inner.L_init", "inner.z_init"): v
               for k, v in mapping.items()}
    missing, unexpected = col.load_state_dict(mapping, strict=False)
    # colony-only buffers (views/masks) have no TRM counterpart, so they stay at init
    colony_only = ("inner.standpoint", "inner.standpoint_inv", "inner.view_mask")
    unexpected = [k for k in unexpected if not k.startswith(colony_only)]
    assert not unexpected, f"unexpected keys: {unexpected}"
    assert all(k.startswith(colony_only) for k in missing), f"missing: {missing}"

    trm.eval(); col.eval()
    with torch.no_grad():
        c1 = trm.initial_carry(batch); c1.halted[:] = True
        c2 = col.initial_carry(batch); c2.halted[:] = True
        _, o1 = trm(carry=c1, batch=batch)
        _, o2 = col(carry=c2, batch=batch)
    diff = (o1["logits"] - o2["logits"]).abs().max().item()
    assert diff < 1e-4, f"GATE 1 FAILED: K=1 colony differs from TRM by {diff}"
    print(f"gate 1 OK — K=1 colony reproduces TRM (max logit diff {diff:.2e})")


def gate_2_modes():
    batch = make_batch()
    for blackboard in ("latent", "tokens"):
        for aggregate in ("mean", "vote", "confidence"):
            for halt_rule in ("quorum", "any", "boss"):
                model = ColonyModel_ACTV1(base_cfg(
                    num_agents=3, blackboard=blackboard, standpoints="sudoku",
                    aggregate=aggregate, halt_rule=halt_rule, role_embeddings=True))
                model.train()
                carry = model.initial_carry(batch)
                carry.halted[:] = True
                carry, out = model(carry=carry, batch=batch)
                assert out["logits"].shape == (BATCH, SEQ, VOCAB), out["logits"].shape
                assert out["q_halt_logits"].shape == (BATCH,), out["q_halt_logits"].shape
                # touch every head so a missing gradient means a real disconnection
                loss = (out["logits"].square().mean()
                        + out["q_halt_logits"].square().mean()
                        + out["q_continue_logits"].square().mean())
                loss.backward()
                bad = [n for n, p in model.named_parameters()
                       if p.requires_grad and (p.grad is None or not torch.isfinite(p.grad).all())]
                assert not bad, f"bad grads in {blackboard}/{aggregate}/{halt_rule}: {bad}"
    print("gate 2 OK — all blackboard x aggregate x halt modes run forward+backward")


def gate_3_clonal_params():
    counts = []
    for k in (1, 2, 4, 8):
        m = ColonyModel_ACTV1(base_cfg(num_agents=k, standpoints="sudoku"))
        counts.append(sum(p.numel() for p in m.parameters()))
    assert len(set(counts)) == 1, f"GATE 3 FAILED: params vary with K: {counts}"
    print(f"gate 3 OK — clonal colony keeps {counts[0]:,} params for K=1,2,4,8")


def gate_6_distributed_evidence():
    """view_fraction<1 must give units genuinely different questions, cover every
    cell, and leave the K=1/full-evidence path untouched."""
    from dsem.models.recursive_reasoning.colony import build_view_masks
    m = build_view_masks(4, SEQ, fraction=0.5, seed=0)
    assert m.shape == (4, SEQ)
    assert m.any(0).all(), "some cell is invisible to every unit — evidence lost, not distributed"
    frac = m.float().mean().item()
    assert 0.4 < frac < 0.65, f"unexpected visible fraction {frac}"
    pairwise = (m[0] == m[1]).float().mean().item()
    assert pairwise < 0.95, "units see near-identical evidence"

    full = build_view_masks(4, SEQ, fraction=1.0, seed=0)
    assert full.all(), "full evidence must mean everyone sees everything"

    batch = make_batch()
    model = ColonyModel_ACTV1(base_cfg(num_agents=4, view_fraction=0.5, standpoints="none"))
    model.eval()
    with torch.no_grad():
        carry = model.initial_carry(batch); carry.halted[:] = True
        inner = model.inner.reset_carry(carry.halted, carry.inner_carry)
        views = model.inner._question_views(batch)
    spread = (views[:, 0] - views[:, 1]).abs().max().item()
    assert spread > 0, "units received identical question embeddings despite masking"
    _ = inner
    print(f"gate 6 OK — distributed evidence: {frac:.0%} visible per unit, all cells covered, "
          f"units see different questions (max emb diff {spread:.2f})")


def gate_7_instrumentation_matches():
    """forward_instrumented must compute exactly what forward computes.

    It once didn't: it bypassed the view masks, gave every unit the full question,
    and so reported perfect consensus and redundant members that were artefacts of
    the measurement rather than properties of the model.
    """
    batch = make_batch()
    for frac in (1.0, 0.5):
        torch.manual_seed(0)
        model = ColonyModel_ACTV1(base_cfg(num_agents=3, view_fraction=frac, standpoints="sudoku"))
        model.eval()
        with torch.no_grad():
            carry = model.initial_carry(batch); carry.halted[:] = True
            ic = model.inner.reset_carry(carry.halted, carry.inner_carry)
            _, logits, _ = model.inner(ic, batch)

            ic2 = model.inner.reset_carry(carry.halted, carry.inner_carry)
            out = model.inner.forward_instrumented(ic2, batch)
        agg = out["rounds"][-1]["aggregate"]
        diff = (logits.argmax(-1) != agg).float().mean().item()
        assert diff < 1e-6, (
            f"GATE 7 FAILED at view_fraction={frac}: instrumented path disagrees with "
            f"forward on {diff:.1%} of cells")
    print("gate 7 OK — instrumented path reproduces forward() at full and partial evidence")


def gate_5_device_context():
    """Buffers must land on the ambient device — pretrain.py builds models inside
    `with torch.device(...)`, and a CPU-pinned buffer only fails there, not here."""
    from dsem.device import get_device
    dev = get_device()
    if dev.type == "cpu":
        print("gate 5 skipped (no accelerator on this host)")
        return
    with torch.device(dev):
        m = ColonyModel_ACTV1(base_cfg(num_agents=2, standpoints="sudoku"))
    bad = [n for n, b in m.named_buffers() if b.device.type != dev.type]
    assert not bad, f"GATE 5 FAILED: buffers not on {dev}: {bad}"
    print(f"gate 5 OK — all buffers land on {dev} when built in a device context")


def gate_4_standpoints():
    perms = build_standpoints("sudoku", 8, SEQ, seed=0)
    inv = invert_perm(perms)
    x = torch.arange(SEQ).view(1, 1, SEQ).expand(2, 8, SEQ)
    fwd = torch.gather(x, 2, perms.unsqueeze(0).expand(2, -1, -1))
    back = torch.gather(fwd, 2, inv.unsqueeze(0).expand(2, -1, -1))
    assert torch.equal(back, x), "GATE 4 FAILED: standpoint round-trip broken"
    assert torch.equal(perms[0], torch.arange(SEQ)), "agent 0 must be the identity view"
    uniq = len({tuple(p.tolist()) for p in perms})
    print(f"gate 4 OK — standpoints round-trip; {uniq}/8 distinct views; agent 0 = identity")


if __name__ == "__main__":
    gate_4_standpoints()
    gate_3_clonal_params()
    gate_1_equivalence()
    gate_2_modes()
    gate_5_device_context()
    gate_6_distributed_evidence()
    gate_7_instrumentation_matches()
    print("\nall gates passed")