"""dexp014_te_dispatch.py — exp014: TEXT-ADDRESSED DISPATCH (Phil's "router
solidifier": the MoE accepts a trained text-encoder adapter as an addressing
conditioning router).

The junction of three findings: diffusion-side codes compress caption
structure to ~0 (exp003) | the text-address channel is live where not
redundant (exp002) | state+sigma dispatch alone finds nothing to route on
(exp007). Here the TEXT side enters the DISPATCH KEY, not the content path.

Per site: AmoeBank-style expert bank (A=4, rank-8 zero-init deltas, dense
signed aleph dispatch, no selectors). Dispatch-key arms:
  sigma_key      — key = key_proj(x) + sig_proj(fourier(sigma))  [exp007 ctrl]
  addr_key       — + txt_proj(frozen byte-trigram aleph address of the
                   caption, flattened [32*128]) — the CANONICAL address as
                   router (frozen address, trained projection).
  te_key         — + te_adapter(CLIP text pooler_output [768] -> D) — THE
                   TRAINED TE ADAPTER AS ROUTER SOLIDIFIER (trained jointly
                   through the diffusion loss).
Substrate: base_lune flow + the fused cache (exp013's regime — largest
adapter headroom; blob gauge carried as the role-aligned instrument).

Judged (all arms, in-bed): common flow-MSE per band; HIGH-band blob gauge;
ROUTER SIGNATURE = per-prompt usage variance vs a SHUFFLED-KEY null (val
pass rerun with deranged text keys; text-specific routing must beat the
null); toggle bit-exact.
Prereg: P1 te_key >= sigma_key on common val; P2 router signature: usage-
by-prompt variance(te_key) > 3x shuffled null; P3 toggle; P4 addr_key vs
te_key ordering = which text representation routes better (no prediction —
first measurement). 1-seed CANDIDATE.

Pod: bash pod2/run_exp014.sh   [DEXP14_SEED=0]
"""
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 MEM_FRACTION, enumerate_sd15_sites
from aleph_diffusion_core import derangement
from dexp009_bandroles import lp
from dexp013_blob_flow import load_lune, blob_lp_err

SHIFT = 2.5
N_TRAIN, N_VAL = 4096, 256
A_EXPERTS, RANK = 4, 8
BATCH = int(os.environ.get("DEXP14_BATCH", "16"))
STEPS = int(os.environ.get("DEXP14_STEPS", "3000"))
SEED = int(os.environ.get("DEXP14_SEED", "0"))
LR, CFG_DROPOUT = 1e-3, 0.1
D11 = ("/workspace/data/dexp011" if os.path.isdir("/workspace")
       else "./data/dexp011")
DATA_DIR = ("/workspace/data/dexp014" if os.path.isdir("/workspace")
            else os.path.join(os.environ.get("GEOLIP_DATA", "./data"),
                              "dexp014"))
CKPT_DIR = ("/workspace/ckpts2/dexp014" if os.path.isdir("/workspace")
            else DATA_DIR)


def fourier_sigma(s, dim=8):
    freqs = torch.pow(2.0, torch.arange(dim // 2, device=s.device))
    ang = 2 * math.pi * s[:, None] * freqs[None]
    return torch.cat([torch.sin(ang), torch.cos(ang)], dim=-1)


class TeDispatchBank(nn.Module):
    """Expert bank with a text-conditioned dense signed dispatch key."""

    def __init__(self, d: int, txt_dim: int = 0, A: int = A_EXPERTS,
                 r: int = RANK, D: int = 4, tau: float = 0.1):
        super().__init__()
        self.d, self.A, self.txt_dim = d, A, txt_dim
        self.tau = tau
        self.down = nn.ModuleList(nn.Linear(d, r, bias=False)
                                  for _ in range(A))
        self.up = nn.ModuleList(nn.Linear(r, d) for _ in range(A))
        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.key_proj = nn.Linear(d, D, bias=False)
        self.sig_proj = nn.Linear(8, D, bias=False)
        nn.init.orthogonal_(self.key_proj.weight)
        nn.init.orthogonal_(self.sig_proj.weight)
        if txt_dim > 0:
            self.txt_proj = nn.Linear(txt_dim, D, bias=False)
            nn.init.orthogonal_(self.txt_proj.weight)
        self.codebook = nn.Parameter(F.normalize(torch.randn(A, D), dim=-1))
        self.gates = nn.Parameter(torch.full((A,), -3.0))
        self.enabled = True
        self.last_usage = None

    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 dispatch(self, x, sig_feat, txt_feat):
        key = self.key_proj(x)
        sig = self.sig_proj(sig_feat)
        sig = sig.view(sig.shape[0], *([1] * (x.ndim - 2)), sig.shape[-1])
        key = key + sig
        if self.txt_dim > 0 and txt_feat is not None:
            tx = self.txt_proj(txt_feat)
            tx = tx.view(tx.shape[0], *([1] * (x.ndim - 2)), tx.shape[-1])
            key = key + tx
        Ax = F.normalize(self.codebook, dim=-1)
        u = (F.normalize(key, dim=-1) @ Ax.transpose(-1, -2)) / self.tau
        m = u.abs().amax(dim=-1, keepdim=True)
        ep, en = torch.exp(u - m), torch.exp(-u - m)
        return (ep - en) / (ep + en).sum(dim=-1, keepdim=True)

    def forward(self, x, sig_feat, txt_feat=None):
        if not self.enabled:
            return x
        w = self.dispatch(x, sig_feat, txt_feat)
        with torch.no_grad():
            self.last_usage = w.abs().mean(dim=tuple(
                range(1, w.ndim - 1))).detach().cpu()   # (B, A) per prompt
        g = torch.sigmoid(self.gates)
        delta = 0
        for k in range(self.A):
            delta = delta + g[k] * w[..., k:k + 1] * self.up[k](
                self.down[k](x))
        return x + delta


class Wrap(nn.Module):
    def __init__(self, block, bank):
        super().__init__()
        self.block = block
        self.bank = bank
        self.sig_feat = None
        self.txt_feat = None

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


def attach_banks(unet, txt_dim):
    sites = enumerate_sd15_sites(unet)
    assert len(sites) == 16
    mods, wraps = nn.ModuleList(), []
    for name, block, d in sites:
        p0 = next(block.parameters())
        m = TeDispatchBank(d, txt_dim).to(device=p0.device, dtype=p0.dtype)
        w = Wrap(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 build_text_feats(device):
    """addr: frozen byte-trigram aleph address of each fused caption
    (canonical recipe, flattened 4096); pooler: CLIP text pooler_output."""
    f = os.path.join(DATA_DIR, "textfeats.pt")
    if os.path.exists(f):
        return torch.load(f, map_location="cpu", weights_only=True)
    os.makedirs(DATA_DIR, exist_ok=True)
    from dexp011_fused_multiband import _iter_fused
    from transformers import CLIPTextModel, CLIPTokenizer
    caps = []
    for r in _iter_fused(N_TRAIN + N_VAL):
        caps.append(str(r.get("caption_joycaption")
                        or r.get("prompt_fused") or ""))
    from dexp002_sd15_addrcond import Addr
    enc = Addr(device)
    addr = enc.encode(caps).reshape(len(caps), -1)      # (N, 4096)
    del enc
    torch.cuda.empty_cache()
    tok = CLIPTokenizer.from_pretrained(
        "stable-diffusion-v1-5/stable-diffusion-v1-5", subfolder="tokenizer")
    te = CLIPTextModel.from_pretrained(
        "stable-diffusion-v1-5/stable-diffusion-v1-5",
        subfolder="text_encoder", torch_dtype=torch.float32).to(device).eval()
    pools = []
    with torch.no_grad():
        for i in range(0, len(caps), 64):
            ids = tok(caps[i:i + 64], padding="max_length", max_length=77,
                      truncation=True, return_tensors="pt").input_ids
            pools.append(te(ids.to(device)).pooler_output.cpu())
    del te
    torch.cuda.empty_cache()
    out = {"addr": addr, "pooler": torch.cat(pools)}
    torch.save(out, f)
    print(f"[textfeats] built: addr {out['addr'].shape}, "
          f"pooler {out['pooler'].shape}", flush=True)
    return out


def run(device="cuda"):
    torch.cuda.set_per_process_memory_fraction(MEM_FRACTION, 0)
    os.makedirs(CKPT_DIR, exist_ok=True)
    cache = torch.load(os.path.join(D11, "cache.pt"), map_location="cpu",
                       weights_only=True)
    with ledger_run("dexp014 text feats", budget_h=0.5):
        tf = build_text_feats(device)
    gv = torch.Generator().manual_seed(SEED + 11)
    u = torch.linspace(0.02, 0.98, N_VAL)
    val_sigma = (SHIFT * u) / (1 + (SHIFT - 1) * u)
    ARMS = {"sigma_key": (0, None),
            "addr_key": (4096, "addr"),
            "te_key": (768, "pooler")}

    def txt_of(kind, sel):
        if kind is None:
            return None
        return tf[kind][sel].float().to("cuda")

    def set_feats(wraps, s01, txt):
        sf = fourier_sigma(s01)
        for wr in wraps:
            wr.sig_feat = sf
            wr.txt_feat = txt

    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
        s = torch.rand(bsz, generator=gen, device=device)
        s = (SHIFT * s) / (1 + (SHIFT - 1) * s)
        noise = torch.randn(lat.shape, generator=gen, device=device)
        s4 = s[:, None, None, None]
        x_t, v = noise * s4 + lat * (1 - s4), noise - lat
        for wr in wraps:
            wr.sig_feat = fourier_sigma(s)
        pred = unet(x_t, s * 1000, ehs, return_dict=False)[0]
        return F.mse_loss(pred, v)

    @torch.no_grad()
    def val(unet, wraps, kind, shuffle_txt=False):
        tot, blob_high, usages = [], [], []
        perm = derangement(N_VAL, seed=SEED + 3)
        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)
            blob = cache["val_blob"][i:i + 32].float().to(device)
            s = val_sigma[i:i + 32].to(device)
            sel = torch.arange(N_TRAIN + i, N_TRAIN + i + lat.shape[0])
            if shuffle_txt:
                sel = N_TRAIN + perm[i:i + lat.shape[0]]
            set_feats(wraps, s, txt_of(kind, sel))
            s4 = s[:, None, None, None]
            x_t, v = noise * s4 + lat * (1 - s4), noise - lat
            pred = unet(x_t, s * 1000, ehs, return_dict=False)[0]
            tot += ((pred - v) ** 2).mean(dim=(1, 2, 3)).tolist()
            x0_hat = x_t - s4 * pred
            bg = blob_lp_err(x0_hat, lat, blob)
            for j, sv in enumerate(s.tolist()):
                if sv > 0.75:
                    blob_high.append(bg[j].item())
            for wr in wraps[8:9]:
                if wr.bank.last_usage is not None:
                    usages.append(wr.bank.last_usage)
        usage = torch.cat(usages) if usages else None
        usage_var = float(usage.var(dim=0).mean()) if usage is not None else 0
        return (sum(tot) / len(tot),
                round(sum(blob_high) / max(len(blob_high), 1), 6),
                round(usage_var, 8))

    results = {"config": {"steps": STEPS, "batch": BATCH, "seed": SEED,
                          "trunk": "base_lune flow", "A": A_EXPERTS,
                          "r": RANK}}
    with ledger_run(f"dexp014 frozen s{SEED}", budget_h=0.2) as h:
        unet = load_lune(device)
        v, bg, _ = val(unet, [], None)
        results["frozen"] = {"val": v, "blob_high": bg}
        del unet
        torch.cuda.empty_cache()
        h["verdict"] = f"val {v:.5f}"

    for arm, (txt_dim, kind) in ARMS.items():
        with ledger_run(f"dexp014 {arm} s{SEED}", budget_h=2.2) as h:
            unet = load_lune(device)
            mods, wraps = attach_banks(unet, txt_dim)
            for m in mods:
                m.assert_zero_init()
            # bind txt feats into the TRAIN loss via a wrapper closure
            opt = torch.optim.Adam(mods.parameters(), lr=LR, weight_decay=0.0)
            gen = torch.Generator(device=device).manual_seed(SEED + 42)
            idx = torch.Generator().manual_seed(SEED + 7)
            t0 = time.time()
            for step in range(1, STEPS + 1):
                sel = torch.randint(0, N_TRAIN, (BATCH,), generator=idx)
                for wr in wraps:
                    wr.txt_feat = txt_of(kind, sel)
                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"[{arm}] step {step}: loss {loss.item():.4f} | "
                          f"{(time.time() - t0) / step:.2f}s/step",
                          flush=True)
            v, bg, uv = val(unet, wraps, kind)
            _, _, uv_null = val(unet, wraps, kind, shuffle_txt=True) \
                if kind else (0, 0, 0)
            for m in mods:
                m.enabled = False
            v_off, _, _ = val(unet, wraps, kind)
            d = abs(v_off - results["frozen"]["val"])
            assert d < 1e-9, f"toggle parity broken: {d}"
            for m in mods:
                m.enabled = True
            torch.save({"mods": [m.state_dict() for m in mods]},
                       os.path.join(CKPT_DIR, f"{arm}_s{SEED}.pt"))
            results[arm] = {"val": v, "blob_high": bg,
                            "usage_var_by_prompt": uv,
                            "usage_var_shuffled_null": uv_null,
                            "val_toggled_off": v_off}
            del unet, mods
            torch.cuda.empty_cache()
            h["verdict"] = f"val {v:.5f} uvar {uv:.2e}"

    sk, ak, tk = (results[a] for a in ("sigma_key", "addr_key", "te_key"))
    results["verdict"] = {
        "P1_te_vs_sigma": "te" if tk["val"] <= sk["val"] else "sigma",
        "P2_router_signature": {
            "te_var": tk["usage_var_by_prompt"],
            "te_null": tk["usage_var_shuffled_null"],
            "hit_3x": tk["usage_var_by_prompt"]
            > 3 * max(tk["usage_var_shuffled_null"], 1e-12)},
        "P4_text_rep_ordering": {
            "addr_key_val": ak["val"], "te_key_val": tk["val"],
            "addr_key_var": ak["usage_var_by_prompt"],
            "te_key_var": tk["usage_var_by_prompt"]},
        "common": {a: results[a]["val"] for a in ARMS},
        "note": "router-solidifier v1; 1-seed CANDIDATE",
    }
    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"dexp014: {json.dumps(results['verdict']['P2_router_signature'])}")
    print(json.dumps(results["verdict"], indent=2))
    burn_down()
    return results


def smoke():
    b = TeDispatchBank(320, txt_dim=768)
    b.assert_zero_init()
    x = torch.randn(2, 9, 320)
    sf = fourier_sigma(torch.tensor([0.3, 0.8]))
    tx = torch.randn(2, 768)
    assert torch.equal(b(x, sf, tx), x), "zero-init must be exact"
    w = b.dispatch(x, sf, tx)
    w2 = b.dispatch(x, sf, torch.randn(2, 768))
    assert w.shape == (2, 9, 4) and not torch.equal(w, w2), \
        "text key must move the dispatch"
    b0 = TeDispatchBank(320, txt_dim=0)
    assert torch.equal(b0(x, sf, None), x)
    print("dexp014 smoke PASSED (zero-init, text-sensitive dispatch, "
          "no-text arm)")


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