"""Test 8: how far must λ_L1 be reduced before DAB and UNI become meaningful?

Same setup as test 7 (saved Adam state from C7 best_uni_step6500.pt, 4 batches
of 8 tiles). Sweep λ_L1 in {25, 15, 10, 5}. For each:
  * Compute influence P_L1, P_UNI, P_DAB, P_HF via leave-one-out counterfactual
    (using saved Adam moments)
  * Total update norm ||Δθ_all||_2
  * cos(Δθ_all(λ), Δθ_all(25))  — how much the update rotates as we tame L1
  * head_conv per-block breakdown

Other lambdas unchanged: λ_UNI=0.25, λ_DAB=1.0, λ_HF=0.01.
"""
import os, sys, glob, random, math, json, copy, time
sys.path.insert(0, "/home/ubuntu/work")
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from src_ds_her2_thunder import UniDwtC7, UNI2HSpatialLoss, StainDeconvDAB
from pytorch_wavelets import DWTForward

DEVICE = "cuda"
DATA = "/home/ubuntu/her2_data"
CKPT = "/home/ubuntu/ckpts/c7_best.pt"
BATCH = 8
K_BATCHES = 4
SEED = 42
LR = 2e-4
LAMBDAS_L1 = [25, 15, 10, 5]
LOSSES = ["L1", "UNI", "DAB", "HF"]

print("Loading UNI-2H + C7 + saved optimizer state ...", flush=True)
uni_loss = UNI2HSpatialLoss(target_size=224).to(DEVICE).eval()
for p in uni_loss.uni.parameters(): p.requires_grad_(False)
G = UniDwtC7(
    uni_module=uni_loss.uni, wavelet="haar", taps=(13, 19),
    dec_ch=(512, 384, 256, 128, 64), skip_ch=(384, 256, 128),
    blocks_per_stage=(4, 3, 3, 2, 2), uni_target_size=224,
).to(DEVICE).train()
ck = torch.load(CKPT, map_location=DEVICE, weights_only=False)
G.load_state_dict(ck["model"], strict=False)
for n, p in G.named_parameters():
    p.requires_grad_(not n.startswith("uni."))
saved_opt_state = ck["optimizer"]
print(f"  model @ step={ck['step']}  optimizer states={len(saved_opt_state['state'])}", flush=True)

stain_dab = StainDeconvDAB().to(DEVICE).eval()
dwt = DWTForward(J=2, wave="haar", mode="periodization").to(DEVICE)

BLOCK_TARGETS = [
    "dwt_skips.enc_sk4", "dwt_skips.enc_sk3", "dwt_skips.enc_sk2",
    "decoder.stem", "decoder.stage1", "decoder.up1", "decoder.fuse2",
    "decoder.stage2", "decoder.stage3", "decoder.stage4", "decoder.stage5",
    "decoder.head_conv",
]
named = dict(G.named_modules())
def _block_of(pname):
    for tgt in BLOCK_TARGETS:
        if pname.startswith(tgt + "."):
            return tgt
    return None

trainable = [(n, p) for n, p in G.named_parameters() if p.requires_grad]


def compute_loss_terms(gens, ihcs, lam_L1=25.0):
    l1 = F.l1_loss(gens, ihcs) * lam_L1
    unic = uni_loss(gens, target=ihcs) * 0.25
    dab = F.l1_loss(stain_dab(gens), stain_dab(ihcs)) * 1.0
    ylg, yhg = dwt((gens + 1) * 0.5)
    ylt, yht = dwt((ihcs + 1) * 0.5)
    hf = F.l1_loss(yhg[0], yht[0]) * 0.01 + F.l1_loss(yhg[1], yht[1]) * 0.01
    return dict(L1=l1, UNI=unic, DAB=dab, HF=hf)


def _load_pair(p):
    he = Image.open(p.replace("/IHC/", "/HE/")).convert("RGB").resize((256, 256), Image.BICUBIC)
    ic = Image.open(p).convert("RGB").resize((256, 256), Image.BICUBIC)
    return (torch.from_numpy(np.asarray(he, dtype=np.float32) / 255.).permute(2, 0, 1) * 2 - 1,
            torch.from_numpy(np.asarray(ic, dtype=np.float32) / 255.).permute(2, 0, 1) * 2 - 1)

def _bf(pth):
    im = Image.open(pth).convert("RGB").resize((256, 256), Image.BICUBIC)
    x = torch.from_numpy(np.asarray(im, dtype=np.float32) / 255.).permute(2, 0, 1).unsqueeze(0) * 2 - 1
    with torch.no_grad():
        return float((stain_dab(x.to(DEVICE)) > 0.3).float().mean().item())
def _slide(p): return Path(p).name.split("_")[0]

random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
ihc = sorted(glob.glob(f"{DATA}/IHC/val/*.jpg"))
pool = random.sample(ihc, min(len(ihc), 1500))
scan = [(p, _bf(p)) for p in pool]
by_p, by_n = {}, {}
for p, s in scan:
    sid = _slide(p)
    if s > 0.1: by_p.setdefault(sid, []).append(p)
    elif s < 0.01: by_n.setdefault(sid, []).append(p)
def _bal(bs, N):
    out = []; slides = list(bs.keys()); random.shuffle(slides)
    while len(out) < N and slides:
        for s in list(slides):
            if not bs[s]: slides.remove(s); continue
            out.append(bs[s].pop())
            if len(out) >= N: break
    return out
her2p = _bal(by_p, 64); her2n = _bal(by_n, 64)
picks = her2p + her2n
random.shuffle(picks)
picks = picks[:K_BATCHES * BATCH]
print(f"tiles: {len(picks)}", flush=True)


def _snapshot_model():
    return {n: p.detach().clone() for n, p in G.named_parameters() if p.requires_grad}

def _restore_model(snap):
    with torch.no_grad():
        for n, p in G.named_parameters():
            if n in snap:
                p.data.copy_(snap[n])

def _mk_opt_with_saved_state():
    opt = torch.optim.AdamW([p for _, p in trainable], lr=LR)
    sd = copy.deepcopy(saved_opt_state)
    try:
        opt.load_state_dict(sd)
    except Exception:
        saved_ids = sd["param_groups"][0]["params"]
        new_ids = opt.state_dict()["param_groups"][0]["params"]
        new_sd = {"param_groups": opt.state_dict()["param_groups"], "state": {}}
        for old_id, new_id in zip(saved_ids, new_ids):
            if old_id in sd["state"]:
                new_sd["state"][new_id] = sd["state"][old_id]
        opt.load_state_dict(new_sd)
    return opt

def _flatten_delta(pre, post):
    return {n: (post[n] - pre[n]).flatten() for n in pre}


# =============================================================================
# Sweep
# =============================================================================
# For cosine-vs-original, we need delta_all(λ=25) as reference for each batch.
per_lambda_results = {lam: {"per_batch": []} for lam in LAMBDAS_L1}
ref_deltas_batch = []   # delta_all at λ_L1=25, one per batch

t0 = time.time()
for bi in range(K_BATCHES):
    print(f"\n--- batch {bi+1}/{K_BATCHES} ---", flush=True)
    chunk = picks[bi * BATCH:(bi + 1) * BATCH]
    pairs = [_load_pair(p) for p in chunk]
    hes = torch.stack([p[0] for p in pairs]).to(DEVICE)
    ihcs = torch.stack([p[1] for p in pairs]).to(DEVICE)
    pre_params = _snapshot_model()

    for lam_L1 in LAMBDAS_L1:
        # delta_all with this lambda
        opt = _mk_opt_with_saved_state()
        G.zero_grad(set_to_none=True)
        gens = G(hes)
        terms = compute_loss_terms(gens, ihcs, lam_L1=lam_L1)
        total = sum(terms.values())
        total.backward()
        opt.step()
        delta_all = _flatten_delta(pre_params, {n: p.detach().clone() for n, p in G.named_parameters() if p.requires_grad})
        _restore_model(pre_params)

        # Leave-one-out per loss
        per_loss_delta = {}
        for skip in LOSSES:
            _restore_model(pre_params)
            opt = _mk_opt_with_saved_state()
            G.zero_grad(set_to_none=True)
            gens = G(hes)
            terms = compute_loss_terms(gens, ihcs, lam_L1=lam_L1)
            total_minus = sum(v for k, v in terms.items() if k != skip)
            total_minus.backward()
            opt.step()
            per_loss_delta[skip] = _flatten_delta(pre_params, {n: p.detach().clone() for n, p in G.named_parameters() if p.requires_grad})
            _restore_model(pre_params)

        # Influence per loss + per block
        overall_inf = {loss: 0.0 for loss in LOSSES}
        per_block_inf = {loss: {tgt: 0.0 for tgt in BLOCK_TARGETS} for loss in LOSSES}
        with torch.no_grad():
            for pname, _ in trainable:
                d_all = delta_all[pname]
                for loss in LOSSES:
                    diff = (d_all - per_loss_delta[loss][pname]).norm().pow(2).item()
                    overall_inf[loss] += diff
                    blk = _block_of(pname)
                    if blk is not None:
                        per_block_inf[loss][blk] += diff
        for loss in LOSSES:
            overall_inf[loss] = math.sqrt(overall_inf[loss])
            for blk in BLOCK_TARGETS:
                per_block_inf[loss][blk] = math.sqrt(per_block_inf[loss][blk])

        # Total update norm ||Δθ_all||_2
        with torch.no_grad():
            update_norm_sq = 0.0
            flat_all_pieces = []
            for pname, _ in trainable:
                v = delta_all[pname]
                update_norm_sq += float(v.pow(2).sum().item())
                flat_all_pieces.append(v)
            update_norm = math.sqrt(update_norm_sq)
            flat_all = torch.cat(flat_all_pieces)

        # Save reference at λ=25 for this batch
        if lam_L1 == 25:
            ref_deltas_batch.append(flat_all.clone())
            cos_vs_ref = 1.0
        else:
            ref = ref_deltas_batch[bi]
            cos_vs_ref = float((flat_all @ ref) / (flat_all.norm() * ref.norm() + 1e-12))

        tot_i = sum(overall_inf.values()) + 1e-30
        entry = dict(
            overall=overall_inf, per_block=per_block_inf,
            update_norm=update_norm, cos_vs_ref_lambda25=cos_vs_ref,
            overall_pct={loss: overall_inf[loss] / tot_i * 100 for loss in LOSSES},
        )
        per_lambda_results[lam_L1]["per_batch"].append(entry)
        print(f"  λ_L1={lam_L1:>3}  P(L1)={entry['overall_pct']['L1']:>5.1f}% "
              f"P(UNI)={entry['overall_pct']['UNI']:>4.1f}% "
              f"P(DAB)={entry['overall_pct']['DAB']:>4.1f}% "
              f"P(HF)={entry['overall_pct']['HF']:>4.1f}% "
              f"||Δθ||={update_norm:.3f}  cos={cos_vs_ref:.4f}", flush=True)

# =============================================================================
# Aggregate + report
# =============================================================================
print("\n" + "=" * 80)
print(f"{'λ_L1':>5}  {'P(L1)':>7}  {'P(UNI)':>7}  {'P(DAB)':>7}  {'P(HF)':>6}"
       f"  {'||Δθ||':>10}  {'cos vs λ25':>10}")
print("-" * 80)
final = {}
for lam in LAMBDAS_L1:
    batches = per_lambda_results[lam]["per_batch"]
    avg_pct = {loss: float(np.mean([b["overall_pct"][loss] for b in batches])) for loss in LOSSES}
    avg_norm = float(np.mean([b["update_norm"] for b in batches]))
    avg_cos = float(np.mean([b["cos_vs_ref_lambda25"] for b in batches]))
    per_block_pct_avg = {}
    for tgt in BLOCK_TARGETS:
        row_avg = {}
        for loss in LOSSES:
            row_avg[loss] = float(np.mean([b["per_block"][loss][tgt] for b in batches]))
        row_tot = sum(row_avg.values()) + 1e-30
        per_block_pct_avg[tgt] = {loss: row_avg[loss] / row_tot * 100 for loss in LOSSES}
    final[lam] = dict(overall_pct=avg_pct, update_norm=avg_norm,
                         cos_vs_ref_lambda25=avg_cos,
                         per_block_pct=per_block_pct_avg)
    print(f"{lam:>5}  {avg_pct['L1']:>6.2f}%  {avg_pct['UNI']:>6.2f}%  {avg_pct['DAB']:>6.2f}% "
          f"  {avg_pct['HF']:>5.2f}%  {avg_norm:>10.4f}  {avg_cos:>10.4f}")

print("\nHead_conv per-loss share by lambda:")
print(f"{'λ_L1':>5}  {'P(L1)':>7}  {'P(UNI)':>7}  {'P(DAB)':>7}  {'P(HF)':>6}")
for lam in LAMBDAS_L1:
    row = final[lam]["per_block_pct"]["decoder.head_conv"]
    print(f"{lam:>5}  {row['L1']:>6.2f}%  {row['UNI']:>6.2f}%  {row['DAB']:>6.2f}%  {row['HF']:>5.2f}%")

outp = Path("/home/ubuntu/work/test8_lambda_L1_sweep.json")
outp.write_text(json.dumps({"config": {"lambdas": LAMBDAS_L1, "K_batches": K_BATCHES,
                                          "batch": BATCH, "lr": LR},
                              "results": final}, indent=2))
print(f"\nSaved: {outp}")
print(f"Wall: {time.time() - t0:.0f}s")