import json, math, os
import numpy as np

SEED = 945
rng = np.random.default_rng(SEED)

def topk_residual(z, d):
    p = z.shape[-1]
    # Keep exactly d largest magnitudes, without sorting all coordinates.
    idx = np.argpartition(z*z, p-d, axis=-1)[..., :p-d]
    # residual is the sum of the p-d smallest squared coordinates
    return np.take_along_axis(z*z, idx, axis=-1).sum(axis=-1)

def orthogonal(p, r):
    a = r.normal(size=(p,p))
    q, rr = np.linalg.qr(a)
    q *= np.sign(np.diag(rr))[None, :]
    return q

def sample_diag(lam, n, r):
    return r.normal(size=(n,len(lam))) * np.sqrt(lam)[None,:]

def estimate(lam, d, U, n=250000, seed=0):
    r = np.random.default_rng(seed)
    x = sample_diag(np.asarray(lam), n, r)
    # U acts on column vectors; rows therefore multiply U.T.
    return float(topk_residual(x @ U.T, d).mean()), float(topk_residual(x, d).mean())

def ci_se(lam, d, U, n=100000, seed=1):
    r = np.random.default_rng(seed)
    x = sample_diag(np.asarray(lam), n, r)
    a = topk_residual(x @ U.T, d)
    return float(a.mean()), float(a.std(ddof=1)/math.sqrt(n))

def main():
    out = {'seed': SEED, 'predictions': {}, 'sweeps': {}}
    # Prediction 1: multiplying covariance by c multiplies expected residual by c.
    lam = np.array([9., 4., 1., .25, .09, .01])
    d = 2
    I = np.eye(len(lam))
    vals = []
    base = None
    for c in [0.25, 0.5, 1., 2., 4.]:
        v,se = ci_se(c*lam, d, I, 160000, 10)
        if c == 1.0: base = v
        vals.append({'scale':c, 'residual':v, 'ratio_to_unit_scale':v/base if base is not None else None, 'predicted_ratio':c})
    unit = next(a['residual'] for a in vals if a['scale'] == 1.0)
    for a in vals: a['ratio_to_unit_scale'] = a['residual'] / unit
    out['predictions']['homogeneity_in_covariance_scale'] = vals
    # Prediction 2: isotropic Gaussian is rotation invariant.
    p=8; di=3; iso=np.ones(p)
    rots=[np.eye(p)] + [orthogonal(p, np.random.default_rng(100+i)) for i in range(5)]
    iso_vals=[]
    for i,u in enumerate(rots):
        v,se=ci_se(iso,di,u,220000,20+i)
        iso_vals.append({'rotation':i,'residual':v,'se':se})
    out['predictions']['isotropic_rotation_invariance'] = iso_vals
    # Prediction 3: PCA basis should be no worse than tested rotations; inspect gap vs d.
    lam=np.array([16.,9.,4.,2.,1.,.5,.2,.05])
    rotation_bank=[orthogonal(len(lam), np.random.default_rng(500+i)) for i in range(12)]
    gap=[]
    for dd in [1,2,3,4,5,6,7]:
        pca,se=ci_se(lam,dd,np.eye(len(lam)),220000,300+dd)
        rs=[]
        for j,u in enumerate(rotation_bank):
            v,_=ci_se(lam,dd,u,70000,700+10*dd+j)
            rs.append(v)
        gap.append({'d':dd,'pca_residual':pca,'random_rotation_mean':float(np.mean(rs)),
                    'random_rotation_min':float(np.min(rs)), 'gap_mean':float(np.mean(rs)-pca),
                    'gap_fraction_of_pca':float((np.mean(rs)-pca)/pca)})
    out['predictions']['anisotropic_rotation_gap'] = gap
    # Mini experiment: finite calibration PCA versus random rotation on held-out activations.
    p=32; dlist=[3,8,16]; ncal=3000; ntest=100000
    lam=np.geomspace(12,.08,p)
    rr=np.random.default_rng(808)
    cal=sample_diag(lam,ncal,rr); test=sample_diag(lam,ntest,rr)
    # Covariance eigensystem from centered calibration; rows use H Q.
    C=np.cov(cal,rowvar=False,bias=True)
    ev,Q=np.linalg.eigh(C); Q=Q[:,np.argsort(ev)[::-1]]
    # Centering is zero in this controlled model, but use estimated mean as implementation does.
    mu=cal.mean(0); R=orthogonal(p,np.random.default_rng(809))
    mini=[]
    for dd in dlist:
        z_p=(test-mu)@Q
        z_r=(test-mu)@R.T
        vp=float(topk_residual(z_p,dd).mean()); vr=float(topk_residual(z_r,dd).mean())
        total=float(np.sum(lam)); mini.append({'d':dd,'d_over_p':dd/p,'pca_mse':vp,'random_rotation_mse':vr,
            'pca_retained_fraction':1-vp/total,'random_retained_fraction':1-vr/total,
            'stored_value_fraction':dd/p})
    out['mini_experiment']={'covariance_eigenbasis_vs_random_rotation':mini,
      'note':'MSE is exact top-d reconstruction residual; stored value fraction is d/p (indices/metadata excluded).'}
    # Numeric implementation sanity: reconstruct and compare direct residual.
    z=sample_diag(lam,1000,np.random.default_rng(901))@Q
    direct=topk_residual(z,8); mask=np.zeros_like(z); ii=np.argpartition(z*z,p-8,axis=1)[:,p-8:]
    # retain the d largest coordinates
    mask[np.arange(len(z))[:,None],ii]=z[np.arange(len(z))[:,None],ii]
    impl=np.sum((z-mask)**2,axis=1)
    out['math_sanity']={'max_abs_residual_formula_error':float(np.max(np.abs(direct-impl))),
                        'orthogonality_error':float(np.max(np.abs(Q.T@Q-np.eye(p))))}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()