"""E8 — End-to-end joint encoder-decoder training, source-anchored metrics.

All fidelity metrics are measured against the SOURCE embedding normalize(x),
never against the trainable encoder output (a moving target that makes
collapse look like success).

Configs (HIGH):
  frozen        : identity encoder, decoder trained          (paper reference)
  e2e_moving    : trainable encoder, loss vs f_phi(x)        (collapse demo)
  e2e_anchored  : trainable encoder, loss vs normalize(x)    (fixed anchor)
  e2e_vicreg    : trainable encoder, loss vs f_phi(x) + VICReg anti-collapse
Metrics: SER/cos vs normalize(x); batch nearest-neighbor retrieval accuracy;
encoder-output effective rank and off-diagonal correlation.
"""
import numpy as np
import torch
import torch.nn.functional as F

import lib
from lib import (SCENARIOS, SNR_GRID, UWCA, DEVICE, Encoder, block_masks,
                 channel, gen_embeddings, save_json, semantic_loss, set_seed)

rng = set_seed(42)
d, U, H = 64, 4, 4
masks = block_masks(U, d)
scen = SCENARIOS["HIGH"]
snrs = [0.0, 4.0, 8.0, 12.0, 16.0, 20.0]


def gen(n):
    return gen_embeddings(n, d, U, rng, scen).to(DEVICE)


def vicreg_reg(Z):
    Zc = Z - Z.mean(0, keepdim=True)
    std = (Zc.var(0) + 1e-4).sqrt()
    v = F.relu(1.0 / d ** 0.5 - std).mean()
    C = (Zc.T @ Zc) / (Z.shape[0] - 1)
    off = C - torch.diag(torch.diag(C))
    c = (off ** 2).sum() / d
    return 25.0 * v + 100.0 * c


def train(mode, epochs=300):
    model = UWCA(d, U, H).to(DEVICE)
    enc = Encoder(d).to(DEVICE) if mode != "frozen" else None
    mask_p = [p for nm, p in model.named_parameters() if "mask_logits" in nm]
    other = [p for nm, p in model.named_parameters() if "mask_logits" not in nm]
    groups = [{"params": other, "lr": 1e-3}, {"params": mask_p, "lr": 0.1}]
    if enc is not None:
        groups.append({"params": enc.parameters(), "lr": 1e-3})
    opt = torch.optim.Adam(groups)
    for ep in range(1, epochs + 1):
        opt.zero_grad()
        loss = 0.0
        for s in snrs:
            X = gen(64)
            n = X.shape[0]
            E = enc(X.reshape(-1, d)).reshape(n, U, d) if enc is not None else X
            target = E if mode in ("frozen", "e2e_moving", "e2e_vicreg") else X
            ch = channel(E, snr_db=s)
            Eh = model(ch["yI"], ch["yQ"])
            L = semantic_loss(Eh, target.detach()
                              if mode == "e2e_moving_detach" else target, 0.1)
            if mode == "e2e_vicreg":
                L = L + vicreg_reg(E.reshape(-1, d))
            loss = loss + L
        (loss / len(snrs)).backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
        opt.step()
        if ep % 75 == 0:
            print(f"  [E8-{mode}] ep {ep}/{epochs} "
                  f"loss={float(loss)/len(snrs):.4f}", flush=True)
    return model, enc


@torch.no_grad()
def collapse_metrics(enc):
    X = gen(1024).reshape(-1, d)
    Z = enc(X) if enc is not None else F.normalize(X, dim=-1)
    Zc = Z - Z.mean(0, keepdim=True)
    C = (Zc.T @ Zc) / (Z.shape[0] - 1)
    ev = torch.linalg.eigvalsh(C).clamp(min=1e-12)
    p = ev / ev.sum()
    return float(torch.exp(-(p * p.log()).sum()))


@torch.no_grad()
def curves(model, enc):
    ss, cc, rr = [], [], []
    for snr in SNR_GRID:
        a = b = r = 0.0
        n_mc = 150
        for _ in range(n_mc):
            X = gen(64)
            n = X.shape[0]
            E = enc(X.reshape(-1, d)).reshape(n, U, d) if enc is not None else X
            ch = channel(E, snr_db=float(snr))
            Eh = model(ch["yI"], ch["yQ"])
            # all metrics vs the SOURCE
            cos = (Eh * X).sum(-1)
            a += float((cos < 0.45).float().mean())
            b += float(cos.mean())
            # batch retrieval: nearest ENCODED gallery entry (collapse makes
            # the gallery indistinguishable and drives accuracy to chance)
            q = Eh.reshape(-1, d)
            g = E.reshape(-1, d)
            sim = q @ g.T
            r += float((sim.argmax(1) == torch.arange(q.shape[0],
                                                      device=q.device))
                       .float().mean())
        ss.append(a / n_mc); cc.append(b / n_mc); rr.append(r / n_mc)
    return ss, cc, rr


out = {"snr": SNR_GRID.tolist(), "configs": {}}
for mode in ["frozen", "e2e_moving", "e2e_anchored", "e2e_vicreg"]:
    model, enc = train(mode)
    erank = collapse_metrics(enc)
    ss, cc, rr = curves(model, enc)
    out["configs"][mode] = {"ser": ss, "cos": cc, "retr": rr, "erank": erank}
    print(f"[E8] {mode}: erank={erank:.1f} SER@10={ss[5]:.3f} "
          f"cos@10={cc[5]:.3f} retr@10={rr[5]:.3f}", flush=True)

save_json("e8_e2e.json", out)