"""E1 — Degrees-of-freedom fairness study.

Adds full-dimensional receivers on the SAME received signal:
  - lmmse_blind : optimal linear receiver with cross-user correlation set to 0
                  (proves the d/U ceiling is fundamental to correlation-blind
                  processing, not an artifact of the OFDMA baseline)
  - lmmse_genie : optimal linear receiver given the true relevance matrix
                  (genie-aided upper reference; UWCA should approach it)
  - tdma_proj   : orthogonal scheme with an arbitrary orthonormal projection
                  (proves any orthogonal partition is statistically identical
                  to coordinate masking for isotropic embeddings)
Also trains the UWCA decoder per scenario under the single-signal model and
saves checkpoints for reuse (E6).
"""
import numpy as np
import torch

import lib
from lib import (SCENARIOS, SNR_GRID, UWCA, DEVICE, beta_matrix, block_masks,
                 eval_scheme, gen_embeddings, save_json, set_seed, train_multitask)

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

out = {"snr": SNR_GRID.tolist(), "scenarios": {}}
for scen_name in ["HIGH", "LOW", "MIX"]:
    scen = SCENARIOS[scen_name]
    B = beta_matrix(scen)

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

    model = UWCA(d, U, H).to(DEVICE)
    train_multitask(model, gen, tasks, epochs=300, tag=f"E1-{scen_name}")
    torch.save(model.state_dict(), lib.DATA / f"e1_uwca_{scen_name}.pt")

    res = {}
    rng_t = torch.Generator().manual_seed(1)
    for scheme in ["uwca", "ofdma", "sfdma", "noma", "lmmse_blind",
                   "lmmse_genie", "tdma_proj"]:
        sers, coss = [], []
        for snr in SNR_GRID:
            t = {"snr_db": float(snr)}
            s, c = eval_scheme(scheme, gen, t, n_mc=200, model=model, B=B,
                               masks=masks, rng_t=rng_t)
            sers.append(s); coss.append(c)
        res[scheme] = {"ser": sers, "cos": coss}
        print(f"[E1-{scen_name}] {scheme}: SER@10dB={sers[5]:.3f} "
              f"cos@10dB={coss[5]:.3f}", flush=True)
    out["scenarios"][scen_name] = res

save_json("e1_fair_baselines.json", out)