"""E4 v3 — Timing offsets with BLOCK-WISE receiver realignment (R1.8).

The BS knows the per-user timing estimates (pilot-based) and realigns each
user's block region individually inside the single received frame:
    y_al[b_w + t] = y[b_w + t + delta_w],  t = 0..d/U-1
Energy that crossed block boundaries is lost or appears as residual
interference, exactly as in a real system with per-user timing advance error.
"""
import numpy as np
import torch
import torch.nn.functional as F

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

rng = set_seed(42)
d, U, H = 64, 4, 4
dpu = d // U
masks = block_masks(U, d)
scen = SCENARIOS["HIGH"]
off_rng = np.random.default_rng(5)


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


model = UWCA(d, U, H).to(DEVICE)
model.load_state_dict(torch.load(lib.DATA / "e1_uwca_HIGH.pt",
                                 map_location=DEVICE))
model.eval()


def realign(y, offs_hat):
    """Block-wise realignment: shift each user's block back by its offset."""
    n, dd = y.shape
    out = torch.zeros_like(y)
    for u in range(U):
        b0 = u * dpu
        for o in offs_hat[:, u].unique():
            o = int(o.item())
            idx = offs_hat[:, u] == o
            src_end = min(b0 + dpu + o, dd)
            ln = src_end - (b0 + o)
            if ln > 0:
                out[idx, b0:b0 + ln] = y[idx, b0 + o:src_end]
    return out


@torch.no_grad()
def evaluate(mode, dmax, snr, err_p=0.0, n_mc=200):
    s_acc = c_acc = 0.0
    for _ in range(n_mc):
        E = gen(64)
        offs = None
        if dmax > 0:
            offs = torch.from_numpy(
                off_rng.integers(0, dmax + 1, size=(64, U))).to(DEVICE)
        ch = channel(E, snr_db=snr, offsets=offs)
        y = ch["yI"]
        if dmax > 0 and mode.endswith("cor"):
            offs_hat = offs.clone()
            if err_p > 0:
                flip = torch.from_numpy(
                    off_rng.random((64, U)) < err_p).to(DEVICE)
                pm = torch.from_numpy(
                    off_rng.choice([-1, 1], size=(64, U))).to(DEVICE)
                offs_hat = (offs_hat + flip.long() * pm).clamp(min=0)
            y = realign(y, offs_hat)
        if mode.startswith("uwca"):
            Eh = model(y, ch["yQ"])
        else:
            Eh = ofdma_decode(y, masks)
        s_acc += ser(Eh, E)
        c_acc += mean_cos(Eh, E)
    return s_acc / n_mc, c_acc / n_mc


dgrid = [0, 1, 2, 4, 8]
out = {"dmax": dgrid, "snr_eval": [10.0, 20.0], "curves": {}}
for label, mode, ep in [("uwca_uncorrected", "uwca_unc", 0.0),
                        ("ofdma_uncorrected", "ofdma_unc", 0.0),
                        ("uwca_corrected", "uwca_cor", 0.0),
                        ("ofdma_corrected", "ofdma_cor", 0.0),
                        ("uwca_corrected_err20", "uwca_cor", 0.2)]:
    cur = {}
    for snr in out["snr_eval"]:
        cur[str(snr)] = [evaluate(mode, dm, snr, ep) for dm in dgrid]
    out["curves"][label] = cur
    print(f"[E4v3] {label} 10dB SER: "
          f"{[round(a[0],3) for a in cur['10.0']]}", flush=True)

save_json("e4_v3_async.json", out)