src/amoe/diffusion/runtime/attach.py

"""attach / DiffusionAttachHandle / detach — the diffusion runtime verbs.

attach(model, anchor) wraps every bound site (RelayPatch2D or
MultibandDelta by anchor kind) at the DECLARED trunk dtype (dtype law).
detach(verify=True) restores the retained modules and asserts BIT-EXACT
equality with a pre-attach probe where the binding provides one
(UNet bindings do; the DiT binding defers to the diffusion-pipe gates —
documented F4 — and degrades to structural restore with a warning).
"""
from __future__ import annotations

import contextlib
import warnings

import torch
import torch.nn as nn

from ...binding.diffusion import resolve_diffusion
from ...io.checkpoint import DiffusionAnchorCheckpoint, load_diffusion_anchor
from ..core.multiband import BandBlockWrap, MultibandDelta, band_weights
from ..core.relay import BlockWithRelay, RelayPatch2D
from ..laws import N_BANDS


class DiffusionAttachHandle:
    def __init__(self, model, binding, originals, names, wraps, fingerprint):
        self.model = model
        self.binding = binding
        self._originals = originals          # [(qualified_name, module)]
        self.names = list(names)
        self._wraps = wraps                  # BlockWithRelay | BandBlockWrap
        self._fingerprint = fingerprint

    @property
    def kind(self) -> str:
        return ("multiband3"
                if isinstance(self._wraps[0].mod, MultibandDelta) else "relay")

    def modules(self):
        return [w.mod for w in self._wraps]

    # -- masking / toggling (never renormalizes anything) ---------------
    def set_mask(self, mask: dict[str, bool]) -> None:
        for n, w in zip(self.names, self._wraps):
            w.mod.enabled = mask.get(n, True)

    @contextlib.contextmanager
    def only(self, *names):
        prev = [w.mod.enabled for w in self._wraps]
        self.set_mask({n: n in names for n in self.names})
        try:
            yield self
        finally:
            for w, e in zip(self._wraps, prev):
                w.mod.enabled = e

    @contextlib.contextmanager
    def all_off(self):
        prev = [w.mod.enabled for w in self._wraps]
        self.set_mask({n: False for n in self.names})
        try:
            yield self
        finally:
            for w, e in zip(self._wraps, prev):
                w.mod.enabled = e

    @contextlib.contextmanager
    def lesion_band(self, *bands: int):
        """Disable band expert(s) at every site (the exp008/exp010 lesion
        instrument). Multiband anchors only."""
        if self.kind != "multiband3":
            raise TypeError("lesion_band is a multiband instrument; this "
                            f"anchor is kind={self.kind}")
        bad = [b for b in bands if not 0 <= b < N_BANDS]
        if bad:
            raise ValueError(f"bands out of range: {bad}")
        prev = [list(w.mod.band_enabled) for w in self._wraps]
        for w in self._wraps:
            for b in bands:
                w.mod.band_enabled[b] = False
        try:
            yield self
        finally:
            for w, p in zip(self._wraps, prev):
                w.mod.band_enabled = p

    # -- band windows (the step gate) ------------------------------------
    def set_band_windows(self, s01: torch.Tensor) -> None:
        """s01: (B,) noise levels in [0,1] -> cosine crossfade windows on
        every wrap. No-op for relay anchors (windows unused)."""
        w = band_weights(s01)
        for wr in self._wraps:
            wr.w_bands = w

    # -- telemetry --------------------------------------------------------
    def telemetry(self, on: bool) -> None:
        for w in self._wraps:
            if hasattr(w.mod, "telemetry"):
                w.mod.telemetry = on

    def amplitude(self) -> dict[str, float]:
        out = {}
        for n, w in zip(self.names, self._wraps):
            r = getattr(w.mod, "last_delta_ratio", None)
            if r is not None:
                out[n] = round(float(r), 6)
        return out

    def gates(self) -> dict:
        vals = []
        for w in self._wraps:
            m = w.mod
            if hasattr(m, "gate"):
                vals.append(torch.sigmoid(m.gate).item())
            elif hasattr(m, "gates"):
                vals.extend(torch.sigmoid(m.gates).tolist())
        return {"gate_mean": round(sum(vals) / max(len(vals), 1), 5),
                "gate_min": round(min(vals), 5) if vals else None,
                "gate_max": round(max(vals), 5) if vals else None}

    # -- detach -----------------------------------------------------------
    def detach(self, *, verify: bool = True) -> nn.Module:
        for name, mod in self._originals:
            self.binding.replace(self.model, name, mod)
        if verify:
            if self._fingerprint is None:
                warnings.warn(
                    "detach: no pre-attach fingerprint for this binding "
                    "(probe deferred — F4); structural restore only. The "
                    "diffusion-pipe path carries its own parity gates.")
            else:
                post = self.binding.probe(self.model)
                if not torch.equal(post, self._fingerprint):
                    raise RuntimeError(
                        "detach verification FAILED: post-detach output is "
                        "not bit-exact with the pre-attach fingerprint")
        return self.model


def attach(model, anchor, *, binding=None, dtype: "torch.dtype | None" = None,
           strict: bool = True) -> DiffusionAttachHandle:
    b = resolve_diffusion(model, binding)
    sites = b.sites(model)

    ck = (anchor if isinstance(anchor, DiffusionAnchorCheckpoint)
          else load_diffusion_anchor(anchor))
    kind = ck.kind
    if kind not in ("relay", "multiband3"):
        raise ValueError(
            f"anchor kind '{kind}' is not attachable — 'mono'/'bank' are "
            "research-record controls (monolith falsifier / falsified "
            "dispatch bank), loadable for analysis via "
            "load_diffusion_anchor only. See "
            "amoe.diffusion.laws.ROUTING_NEGATIVES.")
    if ck.n_sites != len(sites):
        raise ValueError(
            f"anchor carries {ck.n_sites} sites, binding '{b.name}' "
            f"enumerates {len(sites)} — wrong substrate?")
    if strict:
        fam = ck.meta.get("substrate", {}).get("family")
        if fam and fam not in b.name:
            raise ValueError(
                f"anchor substrate family '{fam}' vs binding '{b.name}' "
                "(pass strict=False to override)")

    dt = dtype if dtype is not None else b.declared_dtype(model)   # dtype law
    fingerprint = b.probe(model)

    originals, wraps, names = [], [], []
    for i, (name, block, d) in enumerate(sites):
        sd = ck.per_site(i)
        if kind == "relay":
            mod = RelayPatch2D.from_state_dict(sd)
            assert mod.d == d, (f"site {i} ({name}): anchor width {mod.d} "
                                f"!= trunk width {d}")
            wrap = BlockWithRelay(block, mod)
        else:
            mod = MultibandDelta.from_state_dict(sd)
            assert mod.d == d, (f"site {i} ({name}): anchor width {mod.d} "
                                f"!= trunk width {d}")
            wrap = BandBlockWrap(block, mod)
        p0 = next(block.parameters(), None)
        dev = p0.device if p0 is not None else torch.device("cpu")
        mod.to(device=dev, dtype=dt)
        b.replace(model, name, wrap)
        originals.append((name, block))
        wraps.append(wrap)
        names.append(name)
    return DiffusionAttachHandle(model, b, originals, names, wraps,
                                 fingerprint)


def detach(handle: DiffusionAttachHandle, *,
           verify: bool = True) -> nn.Module:
    return handle.detach(verify=verify)