"""E6 — Residual (error) orthogonality vs. content preservation.

Resolves the claimed contradiction: the decoded embeddings PRESERVE the shared
scene correlation (rho(e_hat_u, e_hat_v) tracks beta_uv), while the decoding
RESIDUALS r_u = e_hat_u - e_u decorrelate (rho(r_u, r_v) -> 0), which is the
interference-suppression property.  Measured per-sample across dimensions on
the E1 HIGH-trained decoder, vs SNR, together with the scene-component cosine.
"""
import numpy as np
import torch

import lib
from lib import (SCENARIOS, SNR_GRID, UWCA, DEVICE, beta_matrix, block_masks,
                 channel, gen_embeddings, ofdma_decode, sample_corr, save_json,
                 set_seed)

rng = set_seed(42)
d, U, H = 64, 4, 4
masks = block_masks(U, d)
scen = SCENARIOS["HIGH"]
B = beta_matrix(scen)

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

pairs = [(u, v) for u in range(U) for v in range(u + 1, U)]
out = {"snr": SNR_GRID.tolist(), "beta_uv_mean": float(np.mean(
    [B[u, v] for u, v in pairs])), "uwca": {}, "ofdma": {}}


@torch.no_grad()
def measure(decode_fn):
    rho_in, rho_out, rho_res = [], [], []
    for snr in SNR_GRID:
        a_in = a_out = a_res = 0.0
        n_mc = 100
        for _ in range(n_mc):
            E = gen_embeddings(64, d, U, rng, scen).to(DEVICE)
            ch = channel(E, snr_db=float(snr))
            Eh = decode_fn(ch)
            R = Eh - E
            pi = po = pr = 0.0
            for u, v in pairs:
                pi += sample_corr(E[:, u], E[:, v])
                po += sample_corr(Eh[:, u], Eh[:, v])
                pr += sample_corr(R[:, u], R[:, v])
            a_in += pi / len(pairs)
            a_out += po / len(pairs)
            a_res += pr / len(pairs)
        rho_in.append(a_in / n_mc)
        rho_out.append(a_out / n_mc)
        rho_res.append(a_res / n_mc)
    return rho_in, rho_out, rho_res


ri, ro, rr = measure(lambda ch: model(ch["yI"], ch["yQ"]))
out["uwca"] = {"rho_input": ri, "rho_decoded": ro, "rho_residual": rr}
print(f"[E6] UWCA rho_in={ri[5]:.3f} rho_dec={ro[5]:.3f} rho_res={rr[5]:.3f} @10dB",
      flush=True)

ri, ro, rr = measure(lambda ch: ofdma_decode(ch["yI"], masks))
out["ofdma"] = {"rho_input": ri, "rho_decoded": ro, "rho_residual": rr}

save_json("e6_residual_orth.json", out)