"""d1_substrate.py — runner-2 diffusion substrate gates (Day 0; no science on
an uncertified substrate; SUBSTRATE FROZEN AT THE GATE).

Gates (charter: history/plans/2026-07-16_runner2_diffusion_charter.md):
  G1 env      — torch/cuda/dtype facts; Blackwell sm_120 riders (torch>=2.7/
                cu128; flash-attn must NOT be installed — SDPA is the path).
  G2 memory   — set_per_process_memory_fraction(R2_MEM_FRACTION) so overruns
                fail LOUDLY; card + VRAM printed and ledgered.
  G3 hf       — HF_TOKEN env-only + whoami.
  G4 sd15     — ckpt-2500 UNet loads; site enumeration ASSERTS 16
                BasicTransformerBlocks; relay attach; P-INIT/P-TOGGLE bit-exact
                parity; P-FIRE; peak_mem + s/step printed.
  G5 mask     — encoder_attention_mask bit-exact probe (Tier-A toggle path);
                on failure the masked-append design falls back to length-toggle
                + measured presence-offset (ledgered, never silent).
  G6 adam     — plain torch.optim.Adam constructs and IS Adam (never AdamW).

Local:  python pod2/d1_substrate.py --smoke   (parse/shape gates only)
Pod:    python pod2/d1_substrate.py --gate    (all gates, GPU)
"""
from __future__ import annotations

import os
import sys
import time

sys.path[:0] = ["pod2", "."]

import torch

from aleph_diffusion_core import (RelayPatch2D, BlockWithRelay, FireCounter,
                                  toggle_parity)

SD_REPO = "AbstractPhil/sd15-flow-lune-json-prompt"
SD_SUB = "checkpoint-00002500/unet"
SD_BASE = "stable-diffusion-v1-5/stable-diffusion-v1-5"
EXPECTED_SD15_SITES = 16
MEM_FRACTION = float(os.environ.get("R2_MEM_FRACTION", "0.92"))


def gate_env():
    info = {"torch": torch.__version__,
            "cuda_available": torch.cuda.is_available()}
    if torch.cuda.is_available():
        cap = torch.cuda.get_device_capability(0)
        info["device"] = torch.cuda.get_device_name(0)
        info["sm"] = f"sm_{cap[0]}{cap[1]}"
        info["vram_gb"] = round(
            torch.cuda.get_device_properties(0).total_memory / 2**30, 1)
        if cap >= (12, 0):
            try:
                import flash_attn  # noqa: F401
                raise AssertionError(
                    "flash-attn installed on sm_120 — BROKEN there; uninstall "
                    "(SDPA is the sanctioned path, repos/anima-trainer.md)")
            except ImportError:
                pass
            maj, mnr = torch.__version__.split(".")[:2]
            assert (int(maj), int(mnr)) >= (2, 7), \
                "Blackwell sm_120 needs torch>=2.7/cu128"
    print(f"[G1 env] {info}", flush=True)
    return info


def gate_memory():
    assert torch.cuda.is_available(), "G2 needs the pod GPU"
    torch.cuda.set_per_process_memory_fraction(MEM_FRACTION, 0)
    total = torch.cuda.get_device_properties(0).total_memory / 2**30
    print(f"[G2 mem] fraction {MEM_FRACTION} of {total:.1f}GB "
          f"(~{MEM_FRACTION * total:.1f}GB) — overruns now fail LOUDLY",
          flush=True)


def gate_hf():
    assert os.environ.get("HF_TOKEN"), "HF_TOKEN missing from env (env-only law)"
    from huggingface_hub import whoami
    print(f"[G3 hf] identity: {whoami()['name']}", flush=True)


def enumerate_sd15_sites(unet):
    """Walk named_modules; return [(qualified_name, block, width)] for every
    BasicTransformerBlock. NEVER trust a hardcoded count — assert it."""
    from diffusers.models.attention import BasicTransformerBlock
    sites = []
    for name, mod in unet.named_modules():
        if isinstance(mod, BasicTransformerBlock):
            sites.append((name, mod, mod.norm1.normalized_shape[0]))
    return sites


def attach_relays(unet, site_filter=None):
    """Wrap each BasicTransformerBlock with BlockWithRelay(RelayPatch2D(d)).
    Returns nn.ModuleList of relays (fp32) in site order."""
    sites = enumerate_sd15_sites(unet)
    if site_filter:
        sites = [s for s in sites if site_filter(s[0])]
    relays = torch.nn.ModuleList()
    for name, block, d in sites:
        p0 = next(block.parameters())
        # DTYPE LAW (Phil 2026-07-16): adapter dtype MATCHES the trunk dtype
        # (fp32 adapters on a low-precision trunk spin fp32 noise into the
        # environment). SD15 exp001 trunk is fp32 -> matched by construction.
        relay = RelayPatch2D(d).to(device=p0.device, dtype=p0.dtype)
        wrapped = BlockWithRelay(block, relay)
        parent = unet
        parts = name.split(".")
        for p in parts[:-1]:
            parent = getattr(parent, p) if not p.isdigit() else parent[int(p)]
        last = parts[-1]
        if last.isdigit():
            parent[int(last)] = wrapped
        else:
            setattr(parent, last, wrapped)
        relays.append(relay)
    return relays, [s[0] for s in sites]


def _probe_batch(device, n=2, seed=7, cond_len=227):
    g = torch.Generator(device="cpu").manual_seed(seed)
    x = torch.randn(n, 4, 64, 64, generator=g).to(device)
    t = torch.full((n,), 500.0, device=device)
    ehs = torch.randn(n, cond_len, 768, generator=g).to(device)
    return x, t, ehs


def gate_sd15(device="cuda"):
    from diffusers import UNet2DConditionModel
    t0 = time.time()
    unet = UNet2DConditionModel.from_pretrained(
        SD_REPO, subfolder=SD_SUB, torch_dtype=torch.float32).to(device)
    unet.eval()
    print(f"[G4 sd15] UNet loaded fp32 in {time.time() - t0:.1f}s", flush=True)

    sites = enumerate_sd15_sites(unet)
    widths = [w for _, _, w in sites]
    assert len(sites) == EXPECTED_SD15_SITES, \
        f"site map changed: {len(sites)} blocks (expected {EXPECTED_SD15_SITES})"
    print(f"[G4 sd15] {len(sites)} BasicTransformerBlocks, widths {widths}",
          flush=True)

    probes = [_probe_batch(device, seed=s) for s in (7, 11, 13, 17)]
    with torch.no_grad():
        frozen_out = [unet(x, t, ehs, return_dict=False)[0]
                      for x, t, ehs in probes]

    relays, names = attach_relays(unet)
    for r in relays:
        r.assert_zero_init()                            # P-INIT precondition

    def adapted(i):
        x, t, ehs = probes[i]
        with torch.no_grad():
            return unet(x, t, ehs, return_dict=False)[0]

    # P-INIT: zero-init enabled adapters == frozen, bit-exact
    worst = max((adapted(i) - frozen_out[i]).abs().max().item()
                for i in range(len(probes)))
    assert worst == 0.0, f"P-INIT broken: max|delta| {worst} (bias leak class?)"
    # P-TOGGLE: disabled == frozen, bit-exact
    for r in relays:
        r.enabled = False
    worst = max((adapted(i) - frozen_out[i]).abs().max().item()
                for i in range(len(probes)))
    assert worst == 0.0, f"P-TOGGLE broken: max|delta| {worst}"
    for r in relays:
        r.enabled = True
    # P-FIRE
    with FireCounter(relays) as fc:
        adapted(0)
    fc.assert_all_fired()

    torch.cuda.reset_peak_memory_stats()
    t0 = time.time()
    adapted(0)
    dt_ = time.time() - t0
    peak = torch.cuda.max_memory_allocated() / 2**30
    print(f"[G4 sd15] parity gates GREEN | forward {dt_:.2f}s | "
          f"peak_mem {peak:.2f}GB", flush=True)
    return unet, relays, names


def gate_mask(unet, device="cuda"):
    """G5: masked-append vs plain must match; bit-exact preferred, else the
    delta is MEASURED and the Tier-A design falls back (never silent)."""
    x, t, ehs = _probe_batch(device, seed=23)
    extra = torch.randn(ehs.shape[0], 32, 768, device=device)
    ehs_app = torch.cat([ehs, extra], dim=1)
    mask = torch.cat([torch.ones(ehs.shape[:2], device=device),
                      torch.zeros(ehs.shape[0], 32, device=device)],
                     dim=1).bool()
    with torch.no_grad():
        plain = unet(x, t, ehs, return_dict=False)[0]
        masked = unet(x, t, ehs_app, encoder_attention_mask=mask,
                      return_dict=False)[0]
    delta = (plain - masked).abs().max().item()
    verdict = ("BIT-EXACT" if delta == 0.0 else
               f"delta {delta:.3e} — masked-append NOT bit-exact; Tier-A "
               f"toggle falls back to length-toggle + measured offset")
    print(f"[G5 mask] {verdict}", flush=True)
    return delta


def gate_adam():
    opt = torch.optim.Adam([torch.nn.Parameter(torch.zeros(2))], lr=1e-3,
                           weight_decay=0.0)
    assert type(opt) is torch.optim.Adam and not isinstance(
        opt, torch.optim.AdamW), "pure-Adam law violated"
    print("[G6 adam] torch.optim.Adam constructs, wd=0, not AdamW", flush=True)


def smoke():
    """Local parse/shape gates (no GPU, no model downloads)."""
    gate_adam()
    r = RelayPatch2D(320)
    r.assert_zero_init()
    x = torch.randn(1, 4, 320)
    assert torch.equal(r(x), x)
    print("d1_substrate smoke PASSED (adam gate + relay parity, local)")


def gate():
    gate_env()
    gate_memory()
    gate_hf()
    gate_adam()
    unet, relays, names = gate_sd15()
    gate_mask(unet)
    print("[substrate] ALL GATES GREEN — substrate FROZEN", flush=True)


if __name__ == "__main__":
    if "--gate" in sys.argv:
        gate()
    else:
        smoke()