"""dexp013_blob_flow.py — exp013: BLOB SUPERVISION ON THE FLOW SUBSTRATE
(Phil's hypothesis 2026-07-18: "Blob likely will be less effective for eps
than it will be for flow matching").

Mechanism under test: eps-prediction recovers structure through
x0_hat = (x_t - sqrt(1-abar)*eps_hat)/sqrt(abar) — ILL-CONDITIONED at high
noise (divide by vanishing sqrt(abar)) exactly where blob supervision
applies; exp012 measured the resulting inertia (+0.03%). Rectified flow
recovers x0_hat = x_t - sigma*v_hat — EXACT AND LINEAR at every sigma.
Prediction: the same blob coupling moves the blob gauge substantially more
on the flow substrate.

Substrate: base_lune UNet (flow-trained; its strongest judged arm was nl77)
+ the dexp011 FUSED cache (same SD15 VAE, latents + nl77 ehs + blob targets).
Objective: trainer-verbatim flow (sigma~U -> SHIFT-2.5 warp, x_t =
noise*s + lat*(1-s), v = noise - lat), blob term = foreground-weighted
LP-x0 with x0_hat = x_t - s*v_hat, HIGH band only, lambda 0.5 (exp012-dosed
for direct comparison).

Arms (shared cache): frozen | uniform-mb3-flow | blob-mb3-flow. Gauges:
common flow-MSE per band + the role-aligned foreground-LP-x0 gauge (all
arms, in-bed).
Prereg: P1 blob-vs-uniform blob-gauge movement in the HIGH band >= 5x
exp012's relative +0.03% (i.e. >= +0.15%); P2 common gauge undegraded
(<0.5%); P3 toggle bit-exact; P5 (mechanistic) the blob-gauge SCALE is
better-conditioned: report gauge magnitudes per band vs exp012's.

Pod: bash pod2/run_exp013.sh   [DEXP13_SEED=0]
"""
from __future__ import annotations

import json
import os
import sys
import time

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

import torch
import torch.nn.functional as F

from pod_ledger import ledger_run, note, burn_down
from d1_substrate import MEM_FRACTION
from dexp008_multiband import (MultibandDelta, band_weights, band_of, attach,
                               N_BANDS)
from dexp009_bandroles import lp

LUNE_REPO = "AbstractPhil/sd15-flow-lune-flux"
LUNE_SUB = "flux_t2_6_pose_t4_6_port_t1_4/checkpoint-00018765/unet"
SHIFT = 2.5
N_TRAIN, N_VAL = 4096, 256
BATCH = int(os.environ.get("DEXP13_BATCH", "16"))
STEPS = int(os.environ.get("DEXP13_STEPS", "3000"))
SEED = int(os.environ.get("DEXP13_SEED", "0"))
LAM = float(os.environ.get("DEXP13_LAMBDA", "0.5"))
LR, CFG_DROPOUT = 1e-3, 0.1
D11 = ("/workspace/data/dexp011" if os.path.isdir("/workspace")
       else "./data/dexp011")
DATA_DIR = ("/workspace/data/dexp013" if os.path.isdir("/workspace")
            else os.path.join(os.environ.get("GEOLIP_DATA", "./data"),
                              "dexp013"))
CKPT_DIR = ("/workspace/ckpts2/dexp013" if os.path.isdir("/workspace")
            else DATA_DIR)


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


def blob_lp_err(x0_hat, x0, blob):
    d2 = (lp(x0_hat) - lp(x0)) ** 2
    m = blob[:, None]
    denom = m.sum(dim=(1, 2, 3)).clamp_min(1.0) * d2.shape[1]
    return (d2 * m).sum(dim=(1, 2, 3)) / denom


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)
    cache = torch.load(os.path.join(D11, "cache.pt"), map_location="cpu",
                       weights_only=True)
    assert int(cache["blob_empty_count"]) == 0, "need the rebuilt blob targets"
    # fixed val sigmas for the FLOW substrate (warped uniform grid, paired)
    gv = torch.Generator().manual_seed(SEED + 11)
    u = torch.linspace(0.02, 0.98, N_VAL)
    val_sigma = (SHIFT * u) / (1 + (SHIFT - 1) * u)

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

    def flow_pieces(lat, s, noise):
        s4 = s[:, None, None, None]
        return noise * s4 + lat * (1 - s4), noise - lat   # x_t, v

    def loss_of(unet, wraps, lat, ehs, blob, gen, use_blob):
        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)
        w = set_w(wraps, s)
        noise = torch.randn(lat.shape, generator=gen, device=device)
        x_t, v = flow_pieces(lat, s, noise)
        pred = unet(x_t, s * 1000, ehs, return_dict=False)[0]
        base = ((pred - v) ** 2).mean(dim=(1, 2, 3))
        if not use_blob:
            return base.mean()
        x0_hat = x_t - s[:, None, None, None] * pred      # EXACT linear
        blob_term = blob_lp_err(x0_hat, lat, blob)
        return (base + LAM * w[:, 2] * blob_term).mean()

    @torch.no_grad()
    def val(unet, wraps):
        tot = []
        per_band = {0: [], 1: [], 2: []}
        blob_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)
            blob = cache["val_blob"][i:i + 32].float().to(device)
            s = val_sigma[i:i + 32].to(device)
            w = set_w(wraps, s)
            x_t, v = flow_pieces(lat, s, noise)
            pred = unet(x_t, s * 1000, ehs, return_dict=False)[0]
            mse = ((pred - v) ** 2).mean(dim=(1, 2, 3))
            x0_hat = x_t - s[:, None, None, None] * pred
            bg = blob_lp_err(x0_hat, lat, blob)
            tot += mse.tolist()
            for j, sv in enumerate(s.tolist()):
                b = band_of(sv)
                per_band[b].append(mse[j].item())
                blob_band[b].append(bg[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}": round(sum(v_) / max(len(v_), 1), 6)
                 for b, v_ in blob_band.items()})

    def train(unet, wraps, params, label, use_blob):
        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)
        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),
                           cache["blob"][sel].float().to(device), gen,
                           use_blob)
            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", flush=True)

    results = {"config": {"lambda": LAM, "steps": STEPS, "batch": BATCH,
                          "seed": SEED, "trunk": "base_lune 18765 (flow)",
                          "x0_recovery": "x0_hat = x_t - sigma*v (exact)"}}
    with ledger_run(f"dexp013 frozen s{SEED}", budget_h=0.2) as h:
        unet = load_lune(device)
        v, pb, bb = val(unet, [])
        results["frozen"] = {"val": v, "per_band": pb, "blob_gauge": bb}
        del unet
        torch.cuda.empty_cache()
        h["verdict"] = f"val {v:.5f}"

    for label, use_blob in (("uniform_mb3", False), ("blob_mb3", True)):
        with ledger_run(f"dexp013 {label} s{SEED}", budget_h=2.5) as h:
            unet = load_lune(device)
            mods, wraps = attach(unet, lambda d: MultibandDelta(d))
            for m in mods:
                m.assert_zero_init()
            train(unet, wraps, mods.parameters(), label, use_blob)
            v, pb, bb = val(unet, wraps)
            for m in mods:
                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
            torch.save({"mods": [m.state_dict() for m in mods]},
                       os.path.join(CKPT_DIR, f"{label}_s{SEED}.pt"))
            results[label] = {"val": v, "per_band": pb, "blob_gauge": bb,
                              "val_toggled_off": v_off}
            del unet, mods
            torch.cuda.empty_cache()
            h["verdict"] = f"val {v:.5f} blobH {bb['band2']}"

    bm, um = results["blob_mb3"], results["uniform_mb3"]
    rel_move = (um["blob_gauge"]["band2"] - bm["blob_gauge"]["band2"]) \
        / max(um["blob_gauge"]["band2"], 1e-12)
    results["verdict"] = {
        "P1_blob_gauge_high_band": {
            "blob": bm["blob_gauge"]["band2"],
            "uniform": um["blob_gauge"]["band2"],
            "relative_move": round(rel_move, 5),
            "exp012_reference_move": 0.0003,
            "hit_5x": rel_move >= 0.0015},
        "P2_common_undegraded":
            bm["val"] <= um["val"] * 1.005,
        "common_val": {"blob": bm["val"], "uniform": um["val"],
                       "frozen": results["frozen"]["val"]},
        "note": "flow-substrate blob test (Phil's conditioning hypothesis); "
                "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"dexp013: {json.dumps(results['verdict']['P1_blob_gauge_high_band'])}")
    print(json.dumps(results["verdict"], indent=2))
    burn_down()
    return results


def smoke():
    lat = torch.randn(2, 4, 64, 64)
    noise = torch.randn_like(lat)
    s = torch.tensor([0.9, 0.2])
    s4 = s[:, None, None, None]
    x_t = noise * s4 + lat * (1 - s4)
    v = noise - lat
    x0 = x_t - s4 * v
    assert (x0 - lat).abs().max() < 1e-5, "flow x0 recovery must be exact"
    print("dexp013 smoke PASSED (exact linear x0 recovery on the flow path)")


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