"""E5: comparison with a trained user-wise attention receiver.

A learned receiver representative of the end-to-end line (per-user query
attention over the matched-filter outputs, residual skip, unit
normalization) is trained at the operating SNR over uniformly random
affinities, then compared against the closed-form LMMSE and DR receivers
across the affinity sweep. The learned receiver receives no affinity side
information and must infer the coupling from data, which is the standard
setting of the learned line.

Outputs: data/e5_learned.csv, data/e5_train_log.csv
(figures come from replot_all.py only)
"""
import math
import os
import numpy as np
import torch

from semantic_mac import (affinity_matrix, matched_filter,
                          sample_latents_isotropic, demux_lmmse, demux_dr,
                          demux_sr, demux_sc, metrics)

HERE = os.path.dirname(os.path.abspath(__file__))
FIG = os.path.join(HERE, "..", "fig")
DATA = os.path.join(HERE, "..", "data")

U, D, DC = 4, 64, 16
SNR_DB = 10
RHO = 10 ** (SNR_DB / 10)
STEPS = 3000
BATCH = 64


class UserWiseAttention(torch.nn.Module):
    """Per-user query attention over the U matched-filter outputs."""

    def __init__(self, U, d, dk=16, heads=4):
        super().__init__()
        self.U, self.d, self.dk, self.H = U, d, dk, heads
        self.WK = torch.nn.Linear(d, dk * heads, bias=False)
        self.WV = torch.nn.Linear(d, dk * heads, bias=False)
        self.WO = torch.nn.Linear(dk * heads, d, bias=False)
        self.q = torch.nn.Parameter(torch.randn(U, heads, dk) * 0.1)
        self.log_eta = torch.nn.Parameter(torch.zeros(()))

    def forward(self, tilde):
        B, Uu, d = tilde.shape
        K = self.WK(tilde).view(B, Uu, self.H, self.dk)
        V = self.WV(tilde).view(B, Uu, self.H, self.dk)
        sc = torch.einsum('uhk,bihk->buih', self.q, K) / math.sqrt(self.dk)
        alpha = torch.softmax(torch.exp(self.log_eta) * sc, dim=2)
        ctx = torch.einsum('buih,bihk->buhk', alpha, V).reshape(B, Uu, -1)
        out = self.WO(ctx) + tilde
        return out / (out.norm(dim=2, keepdim=True) + 1e-12)


def gen_batch(rng, batch, beta):
    a = math.sqrt(beta) * np.ones(U)
    B = affinity_matrix(a)
    z = sample_latents_isotropic(batch, U, D, DC, a, rng)
    tilde, h = matched_filter(z, B, RHO, rng)
    return z, tilde, h, a, B


def main():
    torch.manual_seed(0)
    rng = np.random.default_rng(21)
    net = UserWiseAttention(U, D)
    n_par = sum(p.numel() for p in net.parameters())
    print(f"learned receiver parameters: {n_par}")
    opt = torch.optim.Adam(net.parameters(), lr=1e-3)
    train_log = []
    for it in range(STEPS):
        beta = float(rng.uniform(0.05, 0.9))
        z, tilde, h, a, B = gen_batch(rng, BATCH, beta)
        zt = torch.tensor(z, dtype=torch.float32)
        tt = torch.tensor(tilde, dtype=torch.float32)
        out = net(tt)
        loss = (1.0 - (out * zt).sum(-1)).mean()
        opt.zero_grad()
        loss.backward()
        opt.step()
        if (it + 1) % 50 == 0:
            train_log.append((it + 1, float(loss.item())))
        if (it + 1) % 500 == 0:
            print(f"  step {it+1}: loss={loss.item():.4f}")
    # convergence evidence for the fixed training budget
    with open(os.path.join(DATA, "e5_train_log.csv"), "w") as f:
        f.write("step,loss\n")
        for st, lo in train_log:
            f.write(f"{st},{lo}\n")

    # evaluation across the affinity sweep
    betas = [0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]
    Vc = np.eye(D)[:, :DC]
    rows = []
    net.eval()
    rng_e = np.random.default_rng(500)
    for beta in betas:
        z, tilde, h, a, B = gen_batch(rng_e, 4000, beta)
        with torch.no_grad():
            out_l = net(torch.tensor(tilde, dtype=torch.float32)).numpy()
        r = {}
        r["Learned"] = metrics(out_l.astype(np.float64), z)
        lm, _ = demux_lmmse(tilde, B, h, RHO)
        r["LMMSE"] = metrics(lm, z)
        r["DR"] = metrics(demux_dr(tilde, B, h, RHO, a, Vc), z)
        rows.append(r)
        print(f"beta={beta:.2f} " + " ".join(
            f"{k}:cos={v[0]:.3f},ser={v[2]:.3f}" for k, v in r.items()))
    with open(os.path.join(DATA, "e5_learned.csv"), "w") as f:
        keys = ["Learned", "LMMSE", "DR"]
        f.write("beta," + ",".join(f"{k}_cos,{k}_nmse,{k}_ser" for k in keys)
                + "\n")
        for b, r in zip(betas, rows):
            f.write(f"{b}," + ",".join(
                f"{r[k][0]},{r[k][1]},{r[k][2]}" for k in keys) + "\n")

    # figures are produced only by the canonical replot_all.py (uniform
    # geometry); experiment scripts write CSVs exclusively.
    print("E5 done. Run replot_all.py to regenerate the figures.")


if __name__ == "__main__":
    main()