"""Gram-Whitened Directional Pooling: numerical MVP and mechanism checks."""
import json
import numpy as np


def unit_rows(x):
    x = np.asarray(x, dtype=float)
    return x / np.maximum(np.linalg.norm(x, axis=1, keepdims=True), 1e-12)


def softmax_weights(theta, atoms, epsilon):
    theta, atoms = unit_rows(theta), unit_rows(atoms)
    logits = theta @ atoms.T / epsilon
    logits -= logits.max(axis=1, keepdims=True)
    e = np.exp(logits)
    return e / e.sum(axis=1, keepdims=True)


def gram_and_pool(theta, q, atoms, epsilon, features):
    """Finite quadrature implementation; features has shape [K,C]."""
    W = softmax_weights(theta, atoms, epsilon)
    q = np.asarray(q, dtype=float); q = q / q.sum()
    c = W.T @ (q[:, None] * features)
    G = W.T @ (q[:, None] * W)
    return W, G, c


def whiten(c, G, lam=0.0):
    A = (G + G.T) / 2 + lam * np.eye(G.shape[0])
    val, vec = np.linalg.eigh(A)
    inv = vec @ np.diag(1.0 / np.maximum(val, 1e-12)) @ vec.T
    return inv @ c


def projected_energy(c, G, lam=0.0):
    z = whiten(c, G, lam)
    return float(np.sum(c * z))


def circle(K=512):
    a = np.arange(K) * 2 * np.pi / K
    return np.stack([np.cos(a), np.sin(a)], axis=1)


def mechanism_checks(seed=7):
    rng = np.random.default_rng(seed)
    theta, q = circle(), np.ones(512) / 512
    angles = np.array([0.0, .18, 1.7, 3.2])
    atoms = np.stack([np.cos(angles), np.sin(angles)], axis=1)
    ang = np.arctan2(theta[:, 1], theta[:, 0])
    f = (1.2*np.cos(2*ang) + .5*np.sin(5*ang))[:, None]
    W, G, c = gram_and_pool(theta, q, atoms, .22, f)
    signal_energy = float(np.sum(q*f[:, 0]**2))
    # Prediction 1: PoU is exact up to floating point.
    pou_err = float(np.max(np.abs(W.sum(1)-1)))
    # Prediction 2: lambda=0 gives an orthogonal projection, hence energy <= ||f||^2.
    z = whiten(c, G)
    reconstruction = W @ z[:, 0]
    reconstruction_energy = float(np.sum(q*reconstruction**2))
    projection_identity_error = abs(reconstruction_energy-projected_energy(c,G))
    # Prediction 3: ridge energy is decreasing; at large lambda, lambda*E -> c^T c.
    lambdas = np.array([0, 1e-5, 1e-3, 1e-2, 1e-1, 1, 10, 100, 1000.])
    energies = np.array([projected_energy(c,G,x) for x in lambdas])
    asymptotic_target = float(np.sum(c*c))
    asymptotic_observed = float(lambdas[-1]*energies[-1])
    # Increasing a dense atom family eventually captures the smooth signal.
    counts, captured = [], []
    for n in [2,3,4,6,8,12,16,24,32]:
        a = np.linspace(0, 2*np.pi, n, endpoint=False)
        at = np.stack([np.cos(a),np.sin(a)],1)
        _, gg, cc = gram_and_pool(theta,q,at,.18,f)
        counts.append(n); captured.append(projected_energy(cc,gg))
    return {
        'prediction_pou_max_error': pou_err,
        'prediction_projection_energy_le_signal': bool(projected_energy(c,G) <= signal_energy + 2e-5),
        'signal_energy': signal_energy,
        'projection_energy': projected_energy(c,G),
        'projection_identity_abs_error': projection_identity_error,
        'lambdas': lambdas.tolist(), 'ridge_energies': energies.tolist(),
        'prediction_ridge_monotone': bool(np.all(np.diff(energies) <= 1e-8)),
        'ridge_asymptotic_predicted_lambdaE': asymptotic_target,
        'ridge_asymptotic_observed_lambdaE_at_1000': asymptotic_observed,
        'atom_counts': counts, 'captured_energy': captured,
        'gram_condition_number': float(np.linalg.cond(G)),
        'energy_fraction': projected_energy(c,G)/signal_energy,
    }


def classification_experiment(seed=11, N=1600):
    """Matched nearest-centroid comparison on rotated frequency signals."""
    rng=np.random.default_rng(seed); K=64; theta=circle(K); q=np.ones(K)/K
    atoms=np.stack([np.cos(np.arange(4)*np.pi/2),np.sin(np.arange(4)*np.pi/2)],1)
    raw=[]; white=[]; y=[]
    for label in range(2):
        for _ in range(N//2):
            phase=rng.uniform(0,2*np.pi); a=np.arctan2(theta[:,1],theta[:,0])+phase
            sig=(np.cos((2 if label==0 else 5)*a)+.15*rng.normal(size=K))[:,None]
            _,G,c=gram_and_pool(theta,q,atoms,.25,sig)
            raw.append(c[:,0]); white.append(whiten(c,G,.01)[:,0]); y.append(label)
    y=np.asarray(y); raw=np.asarray(raw); white=np.asarray(white)
    order=rng.permutation(len(y)); split=int(.7*len(y)); tr,te=order[:split],order[split:]
    def acc(A):
        mu=np.stack([A[tr][y[tr]==j].mean(0) for j in [0,1]])
        pred=((A[te,None]-mu[None])**2).sum(2).argmin(1)
        return float((pred==y[te]).mean())
    return {'raw_soft_pool_accuracy':acc(raw), 'gram_whitened_accuracy':acc(white), 'n_test':len(te)}


def main():
    out={'mechanism':mechanism_checks(), 'classification':classification_experiment()}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()