"""E5 — Nonlinear inter-user semantic structure.

Embeddings are produced by fixed random per-user nonlinear view networks
e_u = normalize(g_u([kappa*s ; p_u])), so the inter-user dependence is
nonlinear and NOT captured by any scalar coefficient or linear covariance.
NONLIN-HIGH: shared s (kappa=1); NONLIN-LOW: independent s per user.
Shows the trained UWCA decoder still exploits the shared structure while the
scalar-parameterized genie LMMSE (mis-specified here) cannot fully.
"""
import numpy as np
import torch

import lib
from lib import (UWCA, DEVICE, ViewNets, block_masks, eval_scheme, save_json,
                 set_seed, train_multitask, SNR_GRID)

rng = set_seed(42)
d, U, H = 64, 4, 4
masks = block_masks(U, d)
vnets = ViewNets(d, U).to(DEVICE)
tasks = [{"snr_db": float(s)} for s in np.arange(0, 21, 4)]

out = {"snr": SNR_GRID.tolist(), "cases": {}}
for case, (kappa, shared) in {"NONLIN-HIGH": (1.0, True),
                              "NONLIN-LOW": (1.0, False)}.items():
    def gen(n, kappa=kappa, shared=shared):
        return vnets.gen(n, d, U, rng, kappa, shared)

    # empirical mean pairwise cosine (the "effective" relevance)
    E = gen(2048)
    C = torch.einsum("nud,nvd->uv", E, E) / 2048
    off = C[~torch.eye(U, dtype=torch.bool, device=C.device)]
    beta_emp = float(off.mean())
    print(f"[E5-{case}] empirical mean pairwise cosine = {beta_emp:.3f}",
          flush=True)

    # mis-specified scalar-model LMMSE uses beta_emp for every pair
    B = np.full((U, U), beta_emp)
    np.fill_diagonal(B, 1.0)

    model = UWCA(d, U, H).to(DEVICE)
    train_multitask(model, gen, tasks, epochs=300, tag=f"E5-{case}")

    res = {"beta_emp": beta_emp}
    for scheme in ["uwca", "ofdma", "noma", "lmmse_genie"]:
        sers, coss = [], []
        for snr in SNR_GRID:
            s, c = eval_scheme(scheme, gen, {"snr_db": float(snr)}, n_mc=200,
                               model=model, B=B, masks=masks)
            sers.append(s); coss.append(c)
        res[scheme] = {"ser": sers, "cos": coss}
        print(f"[E5-{case}] {scheme}: SER@10dB={sers[5]:.3f}", flush=True)
    out["cases"][case] = res

save_json("e5_nonlinear.json", out)