"""
Real-data validation v2 (self-contained; no heavy import).
sklearn digits (8x8=64-dim REAL images) -> mean-centered, unit-normalized
embeddings.  Mean removal decorrelates the shared ink/DC structure so the LOW
scenario attains genuinely low inter-user relevance.

Adds the downstream task-accuracy metric (nearest class-prototype) alongside
SER and mean cosine, and reports the empirical beta_uv per scenario.
"""
import numpy as np
from sklearn.datasets import load_digits

RNG = np.random.default_rng(7)
U, D = 4, 64
DPU = D // U
MASKS = np.zeros((U, D))
for u in range(U):
    MASKS[u, u*DPU:(u+1)*DPU] = 1.0
NOMA_POWER = np.array([0.40, 0.30, 0.20, 0.10])
TAU = 0.45

X, y = load_digits(return_X_y=True)
X = X.astype(np.float64)
X = X - X.mean(0, keepdims=True)                       # remove shared DC structure
X = X / (np.linalg.norm(X, axis=1, keepdims=True) + 1e-8)
by_class = {c: X[y == c] for c in range(10)}
# class prototypes (gallery) for the downstream nearest-prototype classifier
PROTO = np.stack([by_class[c].mean(0) for c in range(10)])
PROTO = PROTO / (np.linalg.norm(PROTO, axis=1, keepdims=True) + 1e-8)

def _norm(E): return E / (np.linalg.norm(E, axis=-1, keepdims=True) + 1e-8)
def cos_sim(Eh, Eg): return (Eh * Eg).sum(-1)

def se_channel(E, snr_db):
    n = E.shape[0]
    X_ = E * MASKS[None]; Ytx = X_.sum(1)
    h = np.sqrt(RNG.standard_normal((n, U, 1))**2 + RNG.standard_normal((n, U, 1))**2) * np.sqrt(0.5)
    nstd = np.sqrt(float(np.mean(Ytx**2)) / (10**(snr_db/10)))
    return h * Ytx[:, None, :] + RNG.standard_normal((n, U, D)) * nstd

def ofdma_decode(Yrx):
    return np.stack([_norm(Yrx[:, u, :] * MASKS[u]) for u in range(U)], 1)

def noma_channel(E, snr_db):
    n = E.shape[0]
    h = np.sqrt(RNG.standard_normal((n, U, 1))**2 + RNG.standard_normal((n, U, 1))**2) * np.sqrt(0.5)
    w = E * np.sqrt(NOMA_POWER)[None, :, None] * h
    yv = w.sum(1)
    nstd = np.sqrt(float(np.mean(yv**2)) / (10**(snr_db/10)))
    return yv + RNG.standard_normal((n, D)) * nstd, h

def noma_sic(yv, h):
    n = yv.shape[0]; Eh = np.zeros((n, U, D)); res = yv.copy()
    for u in range(U):
        z = res / (h[:, u, :] + 1e-8); Eh[:, u, :] = _norm(z)
        res -= h[:, u, :] * np.sqrt(NOMA_POWER[u]) * Eh[:, u, :]
    return Eh

def uwca_decode(Yrx, beta_mat):
    R = Yrx[:, :, None, :] * MASKS[None, None]
    a = beta_mat.copy(); np.fill_diagonal(a, 1.0); a /= a.sum(1, keepdims=True) + 1e-8
    ctx = np.einsum('ui,buid->bud', a, R)
    return np.stack([_norm(ctx[:, u, :]) for u in range(U)], 1)

SCEN = {'HIGH': [3, 3, 3, 3], 'LOW': [0, 1, 7, 4], 'MIX': [3, 3, 8, 1]}

def sample_users(n, ca):
    return np.stack([by_class[ca[u]][RNG.integers(0, len(by_class[ca[u]]), n)] for u in range(U)], 1)

def empirical_beta(ca, n=4000):
    E = sample_users(n, ca)
    B = np.einsum('nud,nvd->uv', E, E) / n
    return B

def downstream_acc(Eh, labels):
    # labels: (U,) true class per user; Eh: (n,U,D)
    sims = np.einsum('nud,cd->nuc', _norm(Eh), PROTO)   # (n,U,10)
    pred = sims.argmax(-1)                               # (n,U)
    return (pred == np.array(labels)[None, :]).mean()

def run(scen, snr, n_mc=400):
    ca = SCEN[scen]; bm = empirical_beta(ca); bmc = bm.copy(); np.fill_diagonal(bmc, 0.0); bmc = np.clip(bmc, 0, None)
    M = {'OFDMA': [0,0,0], 'NOMA-SIC': [0,0,0], 'UWCA': [0,0,0]}  # ser, cos, acc
    for _ in range(n_mc):
        E = sample_users(64, ca)
        Y = se_channel(E, snr); Eh = ofdma_decode(Y)
        M['OFDMA'][0]+=(cos_sim(Eh,E)<TAU).mean(); M['OFDMA'][1]+=cos_sim(Eh,E).mean(); M['OFDMA'][2]+=downstream_acc(Eh,ca)
        yv,h = noma_channel(E,snr); Eh = noma_sic(yv,h)
        M['NOMA-SIC'][0]+=(cos_sim(Eh,E)<TAU).mean(); M['NOMA-SIC'][1]+=cos_sim(Eh,E).mean(); M['NOMA-SIC'][2]+=downstream_acc(Eh,ca)
        Y = se_channel(E, snr); Eh = uwca_decode(Y, bmc)
        M['UWCA'][0]+=(cos_sim(Eh,E)<TAU).mean(); M['UWCA'][1]+=cos_sim(Eh,E).mean(); M['UWCA'][2]+=downstream_acc(Eh,ca)
    return {k:[v[i]/n_mc for i in range(3)] for k,v in M.items()}

if __name__ == '__main__':
    print("=== REAL-DATA v2 (mean-centered sklearn digits, d=64) ===")
    for scen in ['HIGH', 'LOW', 'MIX']:
        B = empirical_beta(SCEN[scen]); off = B[~np.eye(U, dtype=bool)]
        print(f"\n[{scen}] empirical beta_uv (mean off-diag cosine) = {off.mean():.3f} "
              f"(min {off.min():.3f}, max {off.max():.3f})")
        for snr in [0, 10, 20]:
            r = run(scen, snr)
            s = "  ".join([f"{m}: SER={r[m][0]:.3f} cos={r[m][1]:.3f} acc={r[m][2]:.3f}" for m in ['OFDMA','NOMA-SIC','UWCA']])
            print(f"   SNR={snr:3d}dB  {s}")
    print("\nDONE.")