"""E3 — Dynamic user arrivals and departures.

U_max = 8 mask slots, d = 64.  Activity-aware meta-training samples a random
active subset each batch; at inference the attention softmax is restricted to
the active set announced by the scheduler (no retraining).
Compared against: (i) a model trained with all 8 users always active
(mismatch), and (ii) oracle models retrained for each fixed active count.
"""
import numpy as np
import torch

import lib
from lib import (SCENARIOS, UWCA, DEVICE, block_masks, channel, gen_embeddings,
                 mean_cos, save_json, semantic_loss, ser, set_seed)

rng = set_seed(42)
d, Umax, H = 64, 8, 4
masks = block_masks(Umax, d)
scen = SCENARIOS["HIGH8"]
snrs = [0.0, 4.0, 8.0, 12.0, 16.0, 20.0]
act_rng = np.random.default_rng(11)


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


def sample_active(n, k=None):
    """(n, Umax) bool with k active users (random subset per sample)."""
    A = np.zeros((n, Umax), dtype=bool)
    for i in range(n):
        kk = k if k is not None else int(act_rng.integers(2, Umax + 1))
        A[i, act_rng.choice(Umax, size=kk, replace=False)] = True
    return torch.from_numpy(A).to(DEVICE)


def train(model, epochs=250, k=None, tag=""):
    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]
    opt = torch.optim.Adam([{"params": other, "lr": 1e-3},
                            {"params": mask_p, "lr": 0.1}])
    for ep in range(1, epochs + 1):
        opt.zero_grad()
        loss = 0.0
        for s in snrs:
            E = gen(64)
            act = sample_active(64, k)
            Ez = E * act[:, :, None]              # inactive users transmit nothing
            ch = channel(Ez, snr_db=s)
            Eh = model(ch["yI"], ch["yQ"], active=act)
            loss = loss + semantic_loss(Eh, E, 0.1, active=act.float())
        (loss / len(snrs)).backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
        opt.step()
        if ep % 50 == 0:
            print(f"  [E3-{tag}] ep {ep}/{epochs} loss={float(loss)/len(snrs):.4f}",
                  flush=True)


@torch.no_grad()
def evaluate(model, k, snr, n_mc=200):
    s_acc = c_acc = 0.0
    for _ in range(n_mc):
        E = gen(64)
        act = sample_active(64, k)
        Ez = E * act[:, :, None]
        ch = channel(Ez, snr_db=snr)
        Eh = model(ch["yI"], ch["yQ"], active=act)
        s_acc += ser(Eh, E, active=act.float())
        c_acc += mean_cos(Eh, E, active=act.float())
    return s_acc / n_mc, c_acc / n_mc


print("[E3] training activity-aware model", flush=True)
m_act = UWCA(d, Umax, H).to(DEVICE)
train(m_act, tag="act")
torch.save(m_act.state_dict(), lib.DATA / "e3_uwca_act.pt")

print("[E3] training fixed-U8 model", flush=True)
m_fix = UWCA(d, Umax, H).to(DEVICE)
train(m_fix, k=Umax, tag="fix8")

ks = [2, 3, 4, 5, 6, 7, 8]
out = {"k": ks, "snr_eval": [10.0, 20.0], "activity": {}, "fixed8": {},
       "oracle": {}}
for snr in out["snr_eval"]:
    out["activity"][str(snr)] = [evaluate(m_act, k, snr) for k in ks]
    out["fixed8"][str(snr)] = [evaluate(m_fix, k, snr) for k in ks]
    print(f"[E3] snr={snr} activity={[f'{a[0]:.3f}' for a in out['activity'][str(snr)]]}",
          flush=True)

for k in [2, 4, 6, 8]:
    m_o = UWCA(d, Umax, H).to(DEVICE)
    train(m_o, k=k, epochs=250, tag=f"oracle{k}")
    out["oracle"][str(k)] = {str(snr): evaluate(m_o, k, snr)
                             for snr in out["snr_eval"]}
    print(f"[E3] oracle k={k}: {out['oracle'][str(k)]}", flush=True)

save_json("e3_dynamic_users.json", out)