import json
import numpy as np
from pathlib import Path

EPS = 1e-10

def corr_eigh(X):
    X = X - X.mean(0, keepdims=True)
    s = X.std(0, ddof=1) + 1e-10
    C = (X / s).T @ (X / s) / X.shape[0]
    w, U = np.linalg.eigh(C)
    order = np.argsort(w)[::-1]
    return w[order], U[:, order]

def band_statistics(X, bands):
    w, U = corr_eigh(X)
    Z = (X - X.mean(0)) @ U
    vals, us = [], []
    for a, b in bands:
        z = Z[:, a:b]
        v = np.mean(z*z, axis=0).mean()
        # aggregate fourth cumulant over modes in the band
        k4 = (np.mean(z**4, axis=0) - 3*np.mean(z*z, axis=0)**2).mean()
        vals.append(np.mean(w[a:b]))
        us.append(-k4/(v*v + EPS))
    vals, us = np.asarray(vals), np.asarray(us)
    d = np.diff(np.log(np.abs(us)+1e-7)) / np.diff(np.log(vals+1e-7))
    return w, U, vals, us, d

def rg_mask(w, us, bands, retain):
    # Practical RG proxy: a persistent change in log quartic coupling identifies
    # the signal side; rank bands by slope magnitude, then keep whole bands.
    scale = np.asarray([w[a:b].mean() for a,b in bands])
    d = np.diff(np.log(np.abs(us)+1e-7)) / np.diff(np.log(scale+1e-7))
    score = np.r_[np.abs(d), np.abs(d[-1]) if len(d) else 0]
    # retain contiguous/whole bands with the strongest scale dependence
    chosen = np.argsort(score)[::-1][:max(1, int(np.ceil(retain / (bands[0][1]-bands[0][0]))))]
    mask = np.zeros(len(w), dtype=bool)
    for r in chosen:
        a,b = bands[r]; mask[a:b] = True
    # exact width for fair comparison
    ids = np.flatnonzero(mask)
    if len(ids) > retain: mask[ids[retain:]] = False
    return mask, d, score

def make_controlled(n=12000, p=40, seed=0, s=1.5, signal_bands=(4,7), strength=1.0):
    rng=np.random.default_rng(seed)
    lam=np.linspace(2.0,0.25,p)
    X=rng.normal(size=(n,p))*np.sqrt(lam)
    # Non-Gaussian modes distributed in the middle/bulk. Laplace has excess kurtosis 3.
    a,b=signal_bands
    for j in range(a,b):
        z=rng.laplace(size=n)/np.sqrt(2) # variance 1, excess kurtosis 3
        # preserve covariance eigenvalue and impose u4 approximately -strength*lambda^s
        X[:,j]=np.sqrt(lam[j])*((1-strength)*rng.normal(size=n)+strength*z)
    return X, lam

def verification():
    # Prediction 1: the Gaussian excess-kurtosis estimator has zero mean and
    # standard deviation proportional to N^-1/2.
    ns=[500,2000,8000]
    noise=[]
    for n in ns:
        vals=[]
        for seed in range(30):
            x=np.random.default_rng(seed).normal(size=(n,8))
            k=np.mean(x**4,0)-3*np.mean(x*x,0)**2
            vals.append(np.mean(k))
        noise.append(float(np.std(vals)))
    scaling=float(np.polyfit(np.log(ns),np.log(noise),1)[0])

    # Prediction 2: if excess kurtosis kappa4 = c*lambda^(2+s), then
    # u4=-kappa4/variance^2 has logarithmic slope s.  Use known independent
    # modes so eigendecomposition noise cannot obscure this algebraic check.
    slopes=[]
    rng=np.random.default_rng(4)
    lam=np.linspace(0.5,2.0,40)
    for s in [0.5,1.0,2.0]:
        us=[]
        for l in lam:
            alpha=0.18*l**s/(2.0**s) # mixture excess kurtosis = 3 alpha
            choose=rng.random(12000)<alpha
            z=rng.normal(size=12000)
            lap=rng.laplace(size=12000)/np.sqrt(2)
            z=np.where(choose,lap,z)*np.sqrt(l)
            v=np.mean(z*z); k=np.mean(z**4)-3*v*v
            us.append(abs(-k/(v*v+EPS)))
        slopes.append(float(np.polyfit(np.log(lam),np.log(np.asarray(us)+1e-7),1)[0]))

    # Prediction 3: for a fixed variance mixture, excess kurtosis (and |u4|)
    # grows approximately linearly with the non-Gaussian mixture strength.
    strengths=np.linspace(0,1,5)
    observed=[]
    for st in strengths:
        vals=[]
        for seed in range(10):
            rr=np.random.default_rng(seed); choose=rr.random(10000)<(0.22*st)
            z=np.where(choose,rr.laplace(size=10000)/np.sqrt(2),rr.normal(size=10000))
            v=np.mean(z*z); vals.append(abs((np.mean(z**4)-3*v*v)/(v*v+EPS)))
        observed.append(float(np.mean(vals)))
    linear=float(np.polyfit(strengths,observed,1)[0])
    return {'noise_std_by_N':dict(zip(ns,noise)), 'noise_loglog_slope_observed':scaling,
            'noise_prediction':-0.5, 'controlled_slope_prediction':[.5,1.,2.],
            'controlled_slope_observed':slopes, 'strengths':strengths.tolist(),
            'abs_u4_by_strength':observed, 'u4_strength_linear_slope':linear,
            'u4_strength_prediction':'linear'}

def classification(seed=20):
    rng=np.random.default_rng(seed); p=40; ntr=3000; nte=1500
    # Signal is extensive-rank: classes differ in several middle eigenmodes,
    # while leading modes are high-variance nuisance directions.
    lam=np.linspace(2.5,.3,p); sig=np.arange(12,28)
    def gen(n):
        y=rng.integers(0,2,n); X=rng.normal(size=(n,p))*np.sqrt(lam)
        X[:,sig] += (2*y[:,None]-1)*0.42*np.sqrt(lam[sig])[None,:]
        return X,y
    cal,_=gen(12000); Xtr,ytr=gen(ntr); Xte,yte=gen(nte)
    bands=[(i,i+4) for i in range(0,p,4)]
    w,U,bv,u,d=band_statistics(cal,bands)
    retain=16
    rmask,_,score=rg_mask(w,u,bands,retain)
    top=np.zeros(p,bool); top[:retain]=1
    def acc(mask):
        # Linear discriminant in retained PCA coordinates, fit without sklearn.
        ztr=(Xtr-Xtr.mean(0))@U[:,mask]; zte=(Xte-Xtr.mean(0))@U[:,mask]
        a=ztr[ytr==1].mean(0)-ztr[ytr==0].mean(0)
        c=.5*(ztr[ytr==1].mean(0)+ztr[ytr==0].mean(0))
        pred=( (zte-c)@a > 0 )
        return float(np.mean(pred==yte))
    return {'retained_width':retain,'rg_selected_bands':np.flatnonzero(rmask.reshape(-1,4).any(1)).tolist(),
            'top_selected_bands':list(range(4)), 'rg_accuracy':acc(rmask),
            'top_pca_accuracy':acc(top), 'quartic_u4':u.tolist(), 'slopes':d.tolist()}

def main():
    out={'verification':verification(),'classification':classification()}
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()