"""
Ablation / sensitivity studies for the TWC revision.
  Part 1 (torch): H (heads), K (number of SNR meta-tasks), S (inner steps)
                  sensitivity -> SER@10dB, cos@10dB, |rho_off|@10dB on HIGH.
  Part 2 (numpy): refined scalability -> mean cosine + top-k retention vs U.
"""
import types
import numpy as np
import torch
import maml_semantic as M

torch.manual_seed(0)

class _Sim:                      # minimal stand-ins (avoid importing the heavy sim module)
    BATCH = 64
    @staticmethod
    def _norm(E):
        return E / (np.linalg.norm(E, axis=-1, keepdims=True) + 1e-8)
    @staticmethod
    def cos_sim(Eh, Eg):
        return (Eh * Eg).sum(-1)
sim = _Sim()

def make_cfg(**kw):
    base = dict(d=64, U=4, H=4, tau=0.45, lam=0.1, snr_min=0.0, snr_max=20.0,
                snr_step=2.0, inner_lr=0.01, inner_steps=5, outer_lr=1e-3,
                meta_epochs=70, joint_epochs=70, batch=64, n_mc=60, seed=42,
                scenario='HIGH', decoder_only=True)
    base.update(kw)
    return types.SimpleNamespace(**base)

def train_eval(cfg):
    device = 'cpu'
    rng = np.random.default_rng(cfg.seed)
    scen = M.SCENARIO_CONFIGS[cfg.scenario]
    model = M.SemanticCommSystem(cfg.d, cfg.U, cfg.H, decoder_only=cfg.decoder_only).to(device)
    tr = M.MAMLTrainer(model, cfg, device, rng, scen)
    tr.train()
    res = M.evaluate_model(model, cfg, device, rng, "rayleigh", scen)
    snr = np.arange(cfg.snr_min, cfg.snr_max + 1e-6, cfg.snr_step)
    i10 = int(np.argmin(np.abs(snr - 10)))
    rho = res['rho'][i10]
    mask = ~np.eye(cfg.U, dtype=bool)
    return res['ser'][i10], res['cos'][i10], float(np.abs(rho[mask]).mean()), len(snr)

def part1():
    print("\n=== ABLATION (HIGH scenario, decoder-only, SNR=10 dB) ===")
    print("\n-- Attention heads H (d=64) --")
    for H in [1, 2, 4, 8]:
        ser, cos, rho, _ = train_eval(make_cfg(H=H))
        print(f"  H={H}: dk={64//H:2d}  SER={ser:.3f}  cos={cos:.3f}  |rho_off|={rho:.3f}")
    print("\n-- Number of SNR meta-tasks K (via snr_step) --")
    for step in [20.0, 10.0, 4.0, 2.0, 1.0]:
        ser, cos, rho, K = train_eval(make_cfg(snr_step=step))
        print(f"  K={K:2d} (step={step:>4}): SER={ser:.3f}  cos={cos:.3f}  |rho_off|={rho:.3f}")
    print("\n-- Inner-loop steps S --")
    for S in [1, 3, 5, 10]:
        ser, cos, rho, _ = train_eval(make_cfg(inner_steps=S))
        print(f"  S={S:2d}: SER={ser:.3f}  cos={cos:.3f}  |rho_off|={rho:.3f}")

# ---- Part 2: refined scalability (cosine + top-k retention) ----
RNG = np.random.default_rng(11)

def gen_clustered(n, U, D, g, beta):
    nc = U // g
    scenes = [RNG.standard_normal(D) for _ in range(nc)]
    scenes = [s / np.linalg.norm(s) for s in scenes]
    embs = []
    for u in range(U):
        s = scenes[u // g]
        priv = RNG.standard_normal((n, D)); priv /= np.linalg.norm(priv, axis=-1, keepdims=True) + 1e-8
        e = np.sqrt(1 - beta**2) * priv + beta * s[None, :]
        e /= np.linalg.norm(e, axis=-1, keepdims=True) + 1e-8
        embs.append(e)
    return np.stack(embs, 1)

def masks_for(U, D):
    dpu = D // U; Mk = np.zeros((U, D))
    for u in range(U):
        Mk[u, u*dpu:(u+1)*dpu] = 1.0
    return Mk

def se_chan(E, snr_db, Mk):
    n, U, D = E.shape
    Ytx = (E * Mk[None]).sum(1)
    h = np.sqrt(RNG.standard_normal((n, U, 1))**2 + RNG.standard_normal((n, U, 1))**2) * np.sqrt(0.5)
    nstd = np.sqrt(float(np.mean(Ytx**2)) / (10**(snr_db/10)))
    return h * Ytx[:, None, :] + RNG.standard_normal((n, U, D)) * nstd

def attn(Yrx, Mk, bm, topk=None):
    n, U, D = Yrx.shape
    R = Yrx[:, :, None, :] * Mk[None, None]
    a = bm.copy(); np.fill_diagonal(a, 1.0)
    if topk is not None and topk < U:
        for u in range(U):
            order = np.argsort(-a[u]); keep = set(order[:topk]) | {u}
            for v in range(U):
                if v not in keep: a[u, v] = 0.0
    a /= a.sum(1, keepdims=True) + 1e-8
    ctx = np.einsum('ui,buid->bud', a, R)
    return np.stack([sim._norm(ctx[:, u, :]) for u in range(U)], 1)

def part2():
    print("\n=== SCALABILITY: cosine + top-k retention (g=4, beta=0.65, SNR=20 dB) ===")
    g, beta, snr = 4, 0.65, 20.0
    for U in [4, 8, 16, 32]:
        D = 16 * U; Mk = masks_for(U, D)
        bm = np.zeros((U, U))
        for i in range(U):
            for j in range(U):
                if i // g == j // g: bm[i, j] = beta*beta
        nmc = 150; cf = ct = 0.0
        for _ in range(nmc):
            E = gen_clustered(sim.BATCH, U, D, g, beta)
            cf += sim.cos_sim(attn(se_chan(E, snr, Mk), Mk, bm), E).mean()
            ct += sim.cos_sim(attn(se_chan(E, snr, Mk), Mk, bm, topk=g), E).mean()
        cf /= nmc; ct /= nmc
        print(f"  U={U:3d}  cos_full={cf:.3f}  cos_topk(k={g})={ct:.3f}  "
              f"retention={100*ct/cf:5.1f}%  ops full={U*U} topk={U*g} ({U//g}x)")

if __name__ == '__main__':
    part1()
    part2()
    print("\nDONE.")