"""dexp008_multiband.py — exp008: MULTIBAND STEPWISE AMoE, the band-assignment
mechanism test (stage plan: history/plans/2026-07-17_multiband_stage_plan.md).

Phil's band-role map (roles enter as OBJECTIVES in exp009; this bed certifies
the MECHANISM): HIGH (0.75,1] diversity/blobbing | MID (0.35,0.75]
continuity/semantics | LOW [0,0.35] fidelity/detail.

Per site (16): three rank-16 zero-init deltas (W+b, own gates -3.0), combined
by SMOOTH COSINE CROSSFADE WINDOWS on s01 (width 0.06, sum to 1 everywhere —
STRUCTURAL/positional gating, dense + differentiable, no selection event).
Per-band gradients arise from the windows themselves: a sample at s01=0.9
trains (almost) only the HIGH expert. Falsifier: matched monolith rank-48.

Prereg (stage plan):
  P1 multiband >= monolith on overall paired val;
  P2 SPECIALIZATION SIGNATURE — each band expert beats the monolith IN ITS
     OWN band (per-band paired val);
  P3 toggle bit-exact post-train;
  P4 BAND LESION — disabling band-b's experts damages band-b val >> other
     bands (surgical decoupling on the sigma axis).
Substrate: core epred SD1.5 (2-seed certified). Reuses the dexp006 cache.

Pod: bash pod2/run_exp008.sh   [DEXP8_SEED=0 DEXP8_STEPS=3000]
"""
from __future__ import annotations

import json
import math
import os
import sys
import time

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

import torch
import torch.nn as nn
import torch.nn.functional as F

from pod_ledger import ledger_run, note, burn_down
from d1_substrate import enumerate_sd15_sites, MEM_FRACTION
from dexp006_sd15core_relay import make_schedule, add_noise

SD_BASE = "stable-diffusion-v1-5/stable-diffusion-v1-5"
N_TRAIN, N_VAL = 4096, 256
RANK, N_BANDS = 16, 3
BAND_EDGES = (0.35, 0.75)          # LOW [0,.35] | MID (.35,.75] | HIGH (.75,1]
XFADE = 0.06                        # cosine crossfade half-width
BATCH = int(os.environ.get("DEXP8_BATCH", "16"))
STEPS = int(os.environ.get("DEXP8_STEPS", "3000"))
SEED = int(os.environ.get("DEXP8_SEED", "0"))
LR, CFG_DROPOUT = 1e-3, 0.1
DATA_DIR = ("/workspace/data/dexp008" if os.path.isdir("/workspace")
            else os.path.join(os.environ.get("GEOLIP_DATA", "./data"),
                              "dexp008"))
DEXP6_DIR = ("/workspace/data/dexp006" if os.path.isdir("/workspace")
             else os.path.join(os.environ.get("GEOLIP_DATA", "./data"),
                               "dexp006"))
CKPT_DIR = ("/workspace/ckpts2/dexp008" if os.path.isdir("/workspace")
            else DATA_DIR)


def band_weights(s01: torch.Tensor) -> torch.Tensor:
    """Smooth structural band windows: (B, 3), sum to 1 everywhere.
    Cosine crossfade of half-width XFADE around each edge — positional
    gating on the sigma axis, never comparative."""
    def ramp(x):                     # 0 below -XFADE, 1 above +XFADE, smooth
        t = ((x / XFADE).clamp(-1, 1) + 1) / 2
        return 0.5 - 0.5 * torch.cos(t * math.pi)
    e1, e2 = BAND_EDGES
    up1, up2 = ramp(s01 - e1), ramp(s01 - e2)
    low = 1 - up1
    mid = up1 * (1 - up2)
    high = up1 * up2
    return torch.stack([low, mid, high], dim=-1)


def band_of(s01: float) -> int:
    e1, e2 = BAND_EDGES
    return 0 if s01 <= e1 else (1 if s01 <= e2 else 2)


class MultibandDelta(nn.Module):
    """Three band experts per site, window-combined; per-expert enable flags
    (band lesions) + global enable (toggle law)."""

    def __init__(self, d: int, r: int = RANK):
        super().__init__()
        self.d, self.r = d, r
        self.down = nn.ModuleList(nn.Linear(d, r, bias=False)
                                  for _ in range(N_BANDS))
        self.up = nn.ModuleList(nn.Linear(r, d) for _ in range(N_BANDS))
        for dn, up in zip(self.down, self.up):
            nn.init.orthogonal_(dn.weight)
            nn.init.zeros_(up.weight)
            nn.init.zeros_(up.bias)
        self.gates = nn.Parameter(torch.full((N_BANDS,), -3.0))
        self.enabled = True
        self.band_enabled = [True] * N_BANDS

    def assert_zero_init(self):
        for up in self.up:
            assert up.weight.abs().max().item() == 0.0
            assert up.bias.abs().max().item() == 0.0

    def forward(self, x, w_bands):                  # w_bands: (B, 3)
        if not self.enabled:
            return x
        g = torch.sigmoid(self.gates)
        delta = 0
        w = w_bands.view(w_bands.shape[0],
                         *([1] * (x.ndim - 2)), N_BANDS)
        for b in range(N_BANDS):
            if not self.band_enabled[b]:
                continue
            delta = delta + g[b] * w[..., b:b + 1] * self.up[b](
                self.down[b](x))
        return x + delta if not isinstance(delta, int) else x


class MonoDelta(nn.Module):
    def __init__(self, d: int, r: int = RANK * N_BANDS):
        super().__init__()
        self.down = nn.Linear(d, r, bias=False)
        nn.init.orthogonal_(self.down.weight)
        self.up = nn.Linear(r, d)
        nn.init.zeros_(self.up.weight)
        nn.init.zeros_(self.up.bias)
        self.gate = nn.Parameter(torch.tensor(-3.0))
        self.enabled = True

    def forward(self, x, w_bands=None):
        if not self.enabled:
            return x
        return x + torch.sigmoid(self.gate) * self.up(self.down(x))


class BlockWrap(nn.Module):
    def __init__(self, block, mod):
        super().__init__()
        self.block = block
        self.mod = mod
        self.w_bands = None

    def forward(self, *args, **kwargs):
        out = self.block(*args, **kwargs)
        h = out[0] if isinstance(out, tuple) else out
        h = self.mod(h, self.w_bands)
        return (h,) + out[1:] if isinstance(out, tuple) else h


def load_unet(device):
    from diffusers import UNet2DConditionModel
    unet = UNet2DConditionModel.from_pretrained(
        SD_BASE, subfolder="unet", torch_dtype=torch.float32).to(device)
    unet.requires_grad_(False)
    unet.eval()
    unet.enable_gradient_checkpointing()
    return unet


def attach(unet, mk):
    sites = enumerate_sd15_sites(unet)
    assert len(sites) == 16
    mods, wraps = nn.ModuleList(), []
    for name, block, d in sites:
        p0 = next(block.parameters())
        m = mk(d).to(device=p0.device, dtype=p0.dtype)
        w = BlockWrap(block, m)
        parent = unet
        parts = name.split(".")
        for p in parts[:-1]:
            parent = getattr(parent, p) if not p.isdigit() else parent[int(p)]
        if parts[-1].isdigit():
            parent[int(parts[-1])] = w
        else:
            setattr(parent, parts[-1], w)
        mods.append(m)
        wraps.append(w)
    return mods, wraps


def run(device="cuda"):
    torch.cuda.set_per_process_memory_fraction(MEM_FRACTION, 0)
    os.makedirs(CKPT_DIR, exist_ok=True)
    os.makedirs(DATA_DIR, exist_ok=True)
    acp = make_schedule(device)
    cache_f = os.path.join(DEXP6_DIR, "cache.pt")
    assert os.path.exists(cache_f), "needs the dexp006 cache"
    cache = torch.load(cache_f, map_location="cpu", weights_only=True)

    def set_w(wraps, s01):
        w = band_weights(s01)
        for wr in wraps:
            wr.w_bands = w

    def loss_of(unet, wraps, lat, ehs, gen):
        bsz = lat.shape[0]
        drop = torch.rand(bsz, generator=gen, device=device) < CFG_DROPOUT
        ehs = ehs.clone()
        ehs[drop] = 0
        t = torch.randint(0, 1000, (bsz,), generator=gen, device=device)
        set_w(wraps, t.float() / 1000.0)
        noise = torch.randn(lat.shape, generator=gen, device=device)
        pred = unet(add_noise(lat, noise, t, acp), t, ehs,
                    return_dict=False)[0]
        return F.mse_loss(pred, noise)

    @torch.no_grad()
    def val(unet, wraps):
        tot, per_band = [], {0: [], 1: [], 2: []}
        for i in range(0, N_VAL, 32):
            lat = cache["val_lat"][i:i + 32].to(device)
            ehs = cache["val_ehs"][i:i + 32].to(device)
            noise = cache["val_noise"][i:i + 32].to(device)
            t = cache["val_t"][i:i + 32].to(device)
            set_w(wraps, t.float() / 1000.0)
            pred = unet(add_noise(lat, noise, t, acp), t, ehs,
                        return_dict=False)[0]
            mse = ((pred - noise) ** 2).mean(dim=(1, 2, 3))
            tot += mse.tolist()
            for j, tv in enumerate((t.float() / 1000.0).tolist()):
                per_band[band_of(tv)].append(mse[j].item())
        return (sum(tot) / len(tot),
                {f"band{b}": round(sum(v) / max(len(v), 1), 6)
                 for b, v in per_band.items()},
                {f"band{b}_n": len(v) for b, v in per_band.items()})

    def train(unet, wraps, params, label):
        opt = torch.optim.Adam(params, lr=LR, weight_decay=0.0)
        gen = torch.Generator(device=device).manual_seed(SEED + 42)
        idx = torch.Generator().manual_seed(SEED + 7)
        torch.cuda.reset_peak_memory_stats()
        t0 = time.time()
        for step in range(1, STEPS + 1):
            sel = torch.randint(0, N_TRAIN, (BATCH,), generator=idx)
            loss = loss_of(unet, wraps, cache["lat"][sel].to(device),
                           cache["ehs"][sel].to(device), gen)
            loss.backward()
            opt.step()
            opt.zero_grad(set_to_none=True)
            if step == 50 or step % 500 == 0:
                print(f"[{label}] step {step}: loss {loss.item():.4f} | "
                      f"{(time.time() - t0) / step:.2f}s/step | peak "
                      f"{torch.cuda.max_memory_allocated() / 2**30:.1f}GB",
                      flush=True)
        return round((time.time() - t0) / STEPS, 3)

    results = {"config": {"bands": 3, "edges": BAND_EDGES, "xfade": XFADE,
                          "rank": RANK, "steps": STEPS, "batch": BATCH,
                          "lr": LR, "seed": SEED, "substrate": "core epred"}}

    with ledger_run(f"dexp008 frozen s{SEED}", budget_h=0.2) as h:
        unet = load_unet(device)
        v, pb, pn = val(unet, [])
        results["frozen"] = {"val": v, "per_band": pb, "band_n": pn}
        del unet
        torch.cuda.empty_cache()
        h["verdict"] = f"val {v:.5f}"

    with ledger_run(f"dexp008 multiband3 s{SEED}", budget_h=2.5) as h:
        unet = load_unet(device)
        mods, wraps = attach(unet, lambda d: MultibandDelta(d))
        for m in mods:
            m.assert_zero_init()
        n_params = sum(p.numel() for p in mods.parameters())
        print(f"[mb3] 16 sites x 3 band experts, {n_params:,} trainable",
              flush=True)
        spd = train(unet, wraps, mods.parameters(), "mb3")
        v_on, pb_on, _ = val(unet, wraps)
        for m in mods:                                # P3 toggle
            m.enabled = False
        v_off, _, _ = val(unet, wraps)
        d = abs(v_off - results["frozen"]["val"])
        assert d < 1e-9, f"toggle parity broken: {d}"
        for m in mods:
            m.enabled = True
        # P4 band lesions: disable band b everywhere, re-val per band
        lesions = {}
        for b in range(N_BANDS):
            for m in mods:
                m.band_enabled[b] = False
            _, pb_les, _ = val(unet, wraps)
            lesions[f"lesion_band{b}"] = pb_les
            for m in mods:
                m.band_enabled[b] = True
        torch.save({"mods": [m.state_dict() for m in mods]},
                   os.path.join(CKPT_DIR, f"mb3_s{SEED}.pt"))
        results["multiband3"] = {
            "n_params": n_params, "val": v_on, "per_band": pb_on,
            "val_toggled_off": v_off, "s_per_step": spd,
            "lesions": lesions,
            "gates": [[round(torch.sigmoid(g).item(), 4)
                       for g in m.gates] for m in mods[8:9]][0]}
        del unet, mods
        torch.cuda.empty_cache()
        h["verdict"] = f"val {v_on:.5f}"

    with ledger_run(f"dexp008 monolith r48 s{SEED}", budget_h=2.5) as h:
        unet = load_unet(device)
        mods, wraps = attach(unet, lambda d: MonoDelta(d))
        n_params = sum(p.numel() for p in mods.parameters())
        print(f"[mono48] {n_params:,} trainable", flush=True)
        spd = train(unet, wraps, mods.parameters(), "mono48")
        v, pb, _ = val(unet, wraps)
        torch.save({"mods": [m.state_dict() for m in mods]},
                   os.path.join(CKPT_DIR, f"mono48_s{SEED}.pt"))
        results["monolith"] = {"n_params": n_params, "val": v,
                               "per_band": pb, "s_per_step": spd}
        del unet, mods
        torch.cuda.empty_cache()
        h["verdict"] = f"val {v:.5f}"

    mb, mo = results["multiband3"], results["monolith"]
    per_band_wins = {b: mb["per_band"][b] < mo["per_band"][b]
                     for b in ("band0", "band1", "band2")}
    lesion_asym = {}
    for b in range(N_BANDS):
        les = mb["lesions"][f"lesion_band{b}"]
        own = les[f"band{b}"] - mb["per_band"][f"band{b}"]
        other = sum(les[f"band{o}"] - mb["per_band"][f"band{o}"]
                    for o in range(N_BANDS) if o != b) / (N_BANDS - 1)
        lesion_asym[f"band{b}"] = {"own_damage": round(own, 6),
                                   "other_damage": round(other, 6),
                                   "surgical": own > 3 * abs(other) and own > 0}
    results["verdict"] = {
        "P1_multiband_vs_monolith": "multiband" if mb["val"] <= mo["val"]
        else "monolith",
        "P2_per_band_wins": per_band_wins,
        "P4_lesion_asymmetry": lesion_asym,
        "param_ratio": round(mb["n_params"] / mo["n_params"], 3),
        "note": f"seed {SEED}; mechanism test — roles enter in exp009",
    }
    with open(os.path.join(DATA_DIR, "results.json" if SEED == 0
                           else f"results_s{SEED}.json"), "w") as f:
        json.dump(results, f, indent=2)
    note(f"dexp008 verdict: {json.dumps(results['verdict'])}")
    print(json.dumps(results["verdict"], indent=2))
    burn_down()
    return results


def smoke():
    w = band_weights(torch.tensor([0.1, 0.35, 0.55, 0.75, 0.9, 0.72, 0.78]))
    assert w.shape == (7, 3)
    assert torch.allclose(w.sum(-1), torch.ones(7), atol=1e-6), "windows must sum to 1"
    assert w[0].argmax() == 0 and w[2].argmax() == 1 and w[4].argmax() == 2
    m = MultibandDelta(320)
    m.assert_zero_init()
    x = torch.randn(2, 9, 320)
    wb = band_weights(torch.tensor([0.2, 0.9]))
    assert torch.equal(m(x, wb), x), "zero-init must be exact"
    m.enabled = False
    assert m(x, wb) is x
    m.enabled = True
    m.band_enabled = [False, False, False]
    assert m(x, wb) is x, "all-band lesion must be code-path exact"
    m.band_enabled = [True] * 3
    nn.init.normal_(m.up[0].weight, std=0.02)
    assert not torch.equal(m(x, wb), x)
    mb = sum(p.numel() for p in MultibandDelta(320).parameters())
    mo = sum(p.numel() for p in MonoDelta(320).parameters())
    print(f"dexp008 smoke PASSED (windows sum-1 + band argmax, zero-init, "
          f"toggle, lesion path; params mb {mb:,} vs mono {mo:,} @d=320)")


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