src/amoe/testing/diffusion_invariants.py

"""CPU-testable diffusion invariants on tiny stubs (no downloads).

The contracts, each tracing to a paid-for receipt:
  toggle law (code-path skip, bit-exact)        — every r2 experiment
  P-INIT (zero-init W+B exactly inert)          — exp006 bias-leak law
  detach bit-exact via binding probe            — the text line's exp011 bar
  band windows sum-to-1, smooth, positional     — exp008
  flow x0 recovery exact and linear             — exp013 (the conditioning law)
  dtype law (declared, refuses meta)            — R0b catch, fork f4f1c46
  checkpoint round trips (.pt/.safetensors/legacy shapes) — burndown record
  blob-on-eps refusal                           — exp012, 2-seed
"""
from __future__ import annotations

import os
import tempfile

import torch
import torch.nn as nn


# ── tiny stubs ──────────────────────────────────────────────────────────

class _TinySite(nn.Module):
    """Marker block with a width witness; tensor output."""

    def __init__(self, d):
        super().__init__()
        self.d = d
        self.lin = nn.Linear(d, d)

    def forward(self, x):
        return x + torch.tanh(self.lin(x))


class _TinyTupleSite(_TinySite):
    def forward(self, x):
        return (x + torch.tanh(self.lin(x)), None)


class TinyUNetStub(nn.Module):
    """Scattered heterogeneous-width sites (down 8/8, mid 16, up 8) —
    the SD15 shape in miniature, one tuple-output site for the
    hybrid-safety path."""

    def __init__(self, tuple_site=False):
        super().__init__()
        torch.manual_seed(11)
        cls_last = _TinyTupleSite if tuple_site else _TinySite
        self.down = nn.ModuleList([_TinySite(8), _TinySite(8)])
        self.mid = _TinySite(16)
        self.up = nn.ModuleList([cls_last(8)])
        self.enc = nn.Linear(8, 16)
        self.dec = nn.Linear(16, 8)

    def forward(self, x):                       # x: [B, T, 8]
        for blk in self.down:
            out = blk(x)
            x = out[0] if isinstance(out, tuple) else out
        h = self.enc(x)
        out = self.mid(h)
        h = out[0] if isinstance(out, tuple) else out
        x = self.dec(h)
        out = self.up[0](x)
        return out[0] if isinstance(out, tuple) else out


class TinyBinding:
    name = "tiny_diffusion"
    expected_sites = 4

    def sites(self, model):
        from ..binding.diffusion import DiffusionBinding  # noqa: F401
        out = [(n, m, m.d) for n, m in model.named_modules()
               if isinstance(m, _TinySite)]
        assert len(out) == self.expected_sites, len(out)
        return out

    def replace(self, model, name, new):
        from ..binding.diffusion import walk_replace
        walk_replace(model, name, new)

    def declared_dtype(self, model):
        from ..binding.diffusion import _declared_dtype
        return _declared_dtype(model)

    def probe(self, model):
        g = torch.Generator().manual_seed(1400)
        dt = self.declared_dtype(model)
        x = torch.randn(2, 5, 8, generator=g).to(dt)
        with torch.no_grad():
            return model(x).detach().float().cpu()


# ── checkpoint builders ─────────────────────────────────────────────────

def _relay_ckpt(model, *, trained=True, name="t"):
    from ..diffusion.core.relay import RelayPatch2D
    from ..io.checkpoint import DiffusionAnchorCheckpoint
    b = TinyBinding()
    torch.manual_seed(23)
    adapters = {}
    widths = []
    for i, (_, _, d) in enumerate(b.sites(model)):
        r = RelayPatch2D(d, n_slots=4, K=8, hidden=12)
        if trained:
            nn.init.normal_(r.consume[-1].weight, std=0.05)
            nn.init.normal_(r.consume[-1].bias, std=0.05)
        for k, v in r.state_dict().items():
            adapters[f"{i}.{k}"] = v.clone()
        widths.append(d)
    return DiffusionAnchorCheckpoint(adapters, {
        "name": name, "adapter": {"kind": "relay"},
        "substrate": {"family": "tiny_diffusion", "n_sites": len(widths),
                      "site_names": [s[0] for s in b.sites(model)],
                      "widths": widths}})


def _mb_ckpt(model, *, trained=True):
    from ..diffusion.core.multiband import MultibandDelta
    from ..io.checkpoint import DiffusionAnchorCheckpoint
    b = TinyBinding()
    torch.manual_seed(29)
    adapters = {}
    for i, (_, _, d) in enumerate(b.sites(model)):
        m = MultibandDelta(d, r=3)
        if trained:
            for up in m.up:
                nn.init.normal_(up.weight, std=0.05)
        for k, v in m.state_dict().items():
            adapters[f"{i}.{k}"] = v.clone()
    return DiffusionAnchorCheckpoint(adapters, {
        "name": "mb", "adapter": {"kind": "multiband3"},
        "substrate": {"family": "tiny_diffusion", "n_sites": 4}})


# ── invariants ──────────────────────────────────────────────────────────

def assert_band_windows():
    from ..diffusion.core.multiband import band_of, band_weights
    s = torch.linspace(0, 1, 501)
    w = band_weights(s)
    assert (w.sum(-1) - 1).abs().max() < 1e-6, "windows do not sum to 1"
    assert (w >= 0).all() and (w <= 1).all()
    for v, b in ((0.1, 0), (0.5, 1), (0.9, 2)):
        assert band_of(v) == b
        assert int(band_weights(torch.tensor([v])).argmax()) == b
    # max slope of the cosine ramp is pi/(4*XFADE) ~= 13.1 per unit s01;
    # on a 1/500 grid that is ~0.0262 per step — assert the theory bound
    d = (w[1:] - w[:-1]).abs().max()
    assert d < 0.03, f"windows not smooth (max step {d} > theory ~0.0262)"


def assert_band_coordinate():
    """The eps band coordinate is s01 = t/1000 — the normalized DISCRETE
    timestep. Every certified bed trained on it (dexp008/011/012) and the
    proven controller gates on it (dexp010). The tempting substitute,
    1 - alphas_cumprod[t], puts t=500 in a different band entirely, which
    silently decouples training from inference. This test pins the
    convention and documents the divergence it guards against."""
    from ..diffusion.core.multiband import band_of

    # the REAL SD1.5 schedule (scaled_linear betas), not a stand-in: this
    # is the schedule the shipped eps stacks were trained against
    betas = torch.linspace(0.00085 ** 0.5, 0.012 ** 0.5, 1000,
                           dtype=torch.float64) ** 2
    acp = torch.cumprod(1.0 - betas, dim=0)

    disagreements = [t for t in range(1000)
                     if band_of(t / 1000.0) != band_of(float(1 - acp[t]))]
    assert len(disagreements) > 300, (
        f"only {len(disagreements)} timesteps disagree — the guard is not "
        "sharp enough to catch a regression; revisit it")
    # measured on this schedule: 316/1000 timesteps train a DIFFERENT band
    # under the proxy. t=300 -> trained LOW (0.300) vs proxy MID (0.410);
    # t=700 -> trained MID (0.700) vs proxy HIGH (0.918).
    assert band_of(0.300) == 0 and band_of(float(1 - acp[300])) == 1
    assert band_of(0.700) == 1 and band_of(float(1 - acp[700])) == 2

    # the trainer must use the trained coordinate
    import inspect
    from ..diffusion.train import trainer as _tr
    src = inspect.getsource(_tr.train)
    assert "t.float() / 1000.0" in src, \
        "trainer eps band coordinate is not t/1000 (band coordinate law)"
    assert "1 - acp[t]" not in src, \
        "trainer still uses the alphas_cumprod proxy for band windows"

    # and the sampler must agree with the trainer
    src_s = inspect.getsource(
        __import__("amoe.diffusion.runtime.sampler", fromlist=["x"]))
    assert "float(t) / 1000.0" in src_s, \
        "StepGatedSampler eps gate diverged from the training coordinate"


def assert_kind_classification():
    """Dispatch banks carry BOTH key_proj and down.0.weight; the bank test
    must win, or exp007/014/015 stacks load as multiband3 and blow up at
    module construction. Flat single-module cond adapters (exp002) must
    classify too."""
    from ..io.checkpoint import _sd_kind
    bank = {"key_proj.weight": None, "down.0.weight": None,
            "gates": None, "codebook": None}
    assert _sd_kind(bank) == "bank", "bank misclassified (ordering bug)"
    assert _sd_kind({"down.0.weight": None, "gates": None}) == "multiband3"
    assert _sd_kind({"down.weight": None, "gate": None}) == "mono"
    assert _sd_kind({"pos_table": None, "addr_proj.weight": None}) == "cond"
    assert _sd_kind({"addr.codebook": None, "proj.weight": None}) == "relay"
    assert _sd_kind({"nonsense": None}) == "unknown"


def assert_flow_x0_recovery():
    from ..diffusion.train.objectives import flow_pieces
    g = torch.Generator().manual_seed(5)
    lat = torch.randn(6, 4, 8, 8, generator=g)
    s = torch.rand(6, generator=g)
    noise = torch.randn(6, 4, 8, 8, generator=g)
    x_t, v = flow_pieces(lat, s, noise)
    x0 = x_t - s[:, None, None, None] * v
    assert (x0 - lat).abs().max() < 1e-6, "flow x0 recovery not linear/exact"


def assert_toggle_law_diffusion(tuple_site=False):
    from ..diffusion.runtime.attach import attach
    m = TinyUNetStub(tuple_site=tuple_site)
    b = TinyBinding()
    g = torch.Generator().manual_seed(31)
    x = torch.randn(2, 5, 8, generator=g)
    with torch.no_grad():
        base = m(x).clone()

    # relay: P-INIT (fresh zero-init anchors are EXACTLY inert)
    h = attach(m, _relay_ckpt(m, trained=False), binding=b)
    with torch.no_grad():
        assert torch.equal(m(x), base), "P-INIT broken (zero-init leak)"
    h.detach(verify=True)

    # relay: trained anchor moves the output; all_off is bit-exact
    h = attach(m, _relay_ckpt(m, trained=True), binding=b)
    with torch.no_grad():
        on = m(x)
    assert not torch.equal(on, base), "trained relay produced no delta"
    with h.all_off():
        with torch.no_grad():
            assert torch.equal(m(x), base), "toggle law violated (relay)"
    h.detach(verify=True)
    with torch.no_grad():
        assert torch.equal(m(x), base), "model altered after detach"

    # multiband: windows drive it; all_off + full lesion are bit-exact
    h = attach(m, _mb_ckpt(m, trained=True), binding=b)
    h.set_band_windows(torch.tensor([0.9, 0.1]))
    with torch.no_grad():
        on = m(x)
    assert not torch.equal(on, base), "trained multiband produced no delta"
    with h.all_off():
        with torch.no_grad():
            assert torch.equal(m(x), base), "toggle law violated (multiband)"
    with h.lesion_band(0, 1, 2):
        with torch.no_grad():
            assert torch.equal(m(x), base), "full band lesion not bit-exact"
    h.detach(verify=True)


def assert_lesion_guard():
    from ..diffusion.runtime.attach import attach
    m = TinyUNetStub()
    h = attach(m, _relay_ckpt(m), binding=TinyBinding())
    try:
        with h.lesion_band(0):
            pass
        raise AssertionError("lesion_band on a relay anchor must TypeError")
    except TypeError:
        pass
    h.detach(verify=True)


def assert_dtype_law():
    from ..binding.diffusion import _declared_dtype
    from ..diffusion.runtime.attach import attach
    m = TinyUNetStub().to(torch.bfloat16)
    b = TinyBinding()
    g = torch.Generator().manual_seed(31)
    x = torch.randn(2, 5, 8, generator=g).to(torch.bfloat16)
    with torch.no_grad():
        base = m(x).clone()
    ck = _relay_ckpt(m, trained=True)
    h = attach(m, ck, binding=b)                 # dtype=None -> declared
    for mod in h.modules():
        for p in mod.parameters():
            assert p.dtype == torch.bfloat16, "dtype law: adapter != trunk"
    with h.all_off():
        with torch.no_grad():
            assert torch.equal(m(x), base), \
                "toggle law must be dtype-independent (code-path skip)"
    h.detach(verify=True)

    class _Meta(nn.Module):
        def __init__(self):
            super().__init__()
            self.p = nn.Parameter(torch.empty(1, device="meta"))
    try:
        _declared_dtype(_Meta())
        raise AssertionError("declared_dtype must refuse meta params")
    except RuntimeError:
        pass


def assert_checkpoint_roundtrip():
    from ..io.checkpoint import load_diffusion_anchor
    from ..io.safetensors_io import (load_anchor_safetensors,
                                     save_anchor_safetensors)
    m = TinyUNetStub()
    ck = _relay_ckpt(m, trained=True)
    tmp = tempfile.mkdtemp()

    # .pt round trip
    p = os.path.join(tmp, "a.pt")
    ck.save(p)
    back = load_diffusion_anchor(p)
    assert back.kind == "relay" and set(back.adapters) == set(ck.adapters)
    for k in ck.adapters:
        assert torch.equal(back.adapters[k], ck.adapters[k])

    # safetensors round trip, canonical + comfy layouts, hash verified
    for layout in ("amoe", "comfy"):
        sp = os.path.join(tmp, f"a_{layout}.safetensors")
        save_anchor_safetensors(ck, sp, key_layout=layout)
        b2 = load_anchor_safetensors(sp)
        assert set(b2.adapters) == set(ck.adapters), layout
        for k in ck.adapters:
            assert torch.equal(b2.adapters[k], ck.adapters[k]), (layout, k)

    # legacy shapes: relays list / mods list / fork dict (home rebuilt)
    per = [ck.per_site(i) for i in range(4)]
    lp = os.path.join(tmp, "legacy_relays.pt")
    torch.save({"relays": per}, lp)
    assert load_diffusion_anchor(lp).kind == "relay"

    mb = _mb_ckpt(m)
    mp = os.path.join(tmp, "legacy_mods.pt")
    torch.save({"mods": [mb.per_site(i) for i in range(4)]}, mp)
    assert load_diffusion_anchor(mp).kind == "multiband3"

    fork = {str(i): {k: v for k, v in per[i].items() if k != "addr.home"}
            for i in range(4)}
    fp = os.path.join(tmp, "fork.pt")
    torch.save({"relays": fork}, fp)
    fk = load_diffusion_anchor(fp)
    assert fk.meta.get("home_reconstructed") is True, \
        "fork import must STAMP the reconstructed home"
    assert "0.addr.home" in fk.adapters


def assert_blob_on_eps_refusal():
    from ..diffusion.train.config import DiffusionTrainConfig
    from ..diffusion.train.trainer import train
    cfg = DiffusionTrainConfig(objective="eps", blob=True)
    try:
        train(None, {}, cfg)
        raise AssertionError("blob-on-eps must refuse (conditioning law)")
    except ValueError as e:
        assert "conditioning law" in str(e)


def assert_align_negative():
    from ..diffusion.train.aligner import align
    try:
        align()
        raise AssertionError("align must raise with the record")
    except NotImplementedError as e:
        assert "falsified" in str(e)


def run_all() -> None:
    assert_band_windows()
    assert_band_coordinate()
    assert_kind_classification()
    assert_flow_x0_recovery()
    assert_toggle_law_diffusion(tuple_site=False)
    assert_toggle_law_diffusion(tuple_site=True)
    assert_lesion_guard()
    assert_dtype_law()
    assert_checkpoint_roundtrip()
    assert_blob_on_eps_refusal()
    assert_align_negative()
    print("amoe.diffusion invariants: band windows, BAND COORDINATE "
          "(t/1000), kind classification, exact flow x0, toggle law "
          "(relay+multiband, tensor+tuple), P-INIT, bit-exact detach, "
          "dtype law, checkpoint round trips (.pt/.safetensors/legacy/"
          "fork), blob-on-eps refusal, align negative — ALL GREEN")


if __name__ == "__main__":
    run_all()