import json, math, time
import numpy as np
from scipy.special import ndtri

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

def rank_normalize(X):
    n, p = X.shape
    Z = np.empty_like(X, dtype=float)
    for j in range(p):
        order = np.argsort(X[:, j], kind='mergesort')
        probs = (np.arange(n) + 0.5) / n
        zcol = ndtri(probs)
        Z[order, j] = zcol
    Z -= Z.mean(0, keepdims=True)
    return Z

def cov(X):
    X = X - X.mean(0, keepdims=True)
    return X.T @ X / len(X)

def spectral_clean(S, preserve_top=1):
    # MVP nonlinear spectral map: retain isolated spikes, pull noisy bulk to
    # its median. This is a monotone eigenvalue-specific pilot for g_H,gamma.
    w, V = np.linalg.eigh(S)
    w = np.maximum(w, 1e-8)
    p = len(w)
    bulk = w[:-preserve_top] if preserve_top and p > preserve_top else w
    med = np.median(bulk)
    # A mild nonlinear map: bulk is contracted toward its robust center;
    # top isolated directions are retained, not globally rescaled.
    cleaned = w.copy()
    if len(bulk):
        lo, hi = np.quantile(bulk, [0.1, 0.9])
        mask = np.arange(p) < p-preserve_top if preserve_top else np.ones(p, bool)
        cleaned[mask] = med + 0.35 * (w[mask] - med)
        cleaned[mask] = np.maximum(cleaned[mask], 0.05 * med)
    return (V * cleaned) @ V.T, w, cleaned

def latent_cov(p, spike):
    Q, _ = np.linalg.qr(rng.normal(size=(p, p)))
    vals = np.ones(p); vals[0] = spike
    return (Q * vals) @ Q.T, Q, vals

def experiment_copula_convergence():
    # Monotone heavy-tailed marginals preserve ranks but distort Pearson covariance.
    p, rho, reps = 8, 0.55, 30
    R = np.full((p,p), rho); np.fill_diagonal(R, 1.)
    rows=[]
    for n in [64, 128, 256, 512, 1024]:
        raw_err=[]; rank_err=[]
        for _ in range(reps):
            G = rng.multivariate_normal(np.zeros(p), R, size=n)
            # strongly nonlinear monotone marginals, finite but very heavy-tailed
            X = np.sign(G) * (np.exp(np.minimum(np.abs(G), 5.0)) - 1.0)
            raw_err.append(np.linalg.norm(cov(X)-R, 'fro')/np.linalg.norm(R,'fro'))
            rank_err.append(np.linalg.norm(cov(rank_normalize(X))-R, 'fro')/np.linalg.norm(R,'fro'))
        rows.append({'n':n,'raw_relerr':float(np.mean(raw_err)), 'rank_relerr':float(np.mean(rank_err)),
                     'rank_sd':float(np.std(rank_err)/math.sqrt(reps))})
    # predicted sampling law: rank error ~ n^-1/2
    slope = np.polyfit(np.log([x['n'] for x in rows]), np.log([x['rank_relerr'] for x in rows]), 1)[0]
    return rows, float(slope)

def experiment_outliers():
    # One contaminated fraction with arbitrarily large positive shifts.
    p,n,reps,q = 8,256,40,0.05
    R=np.full((p,p),.35); np.fill_diagonal(R,1.)
    rows=[]
    for A in [0, 2, 5, 10, 20, 50]:
        raw=[]; rank=[]; raw_var=[]; rank_var=[]
        for _ in range(reps):
            G=rng.multivariate_normal(np.zeros(p),R,size=n)
            X=G.copy()
            m=max(1,int(q*n)); ix=rng.choice(n,m,replace=False)
            X[ix,0] += A
            raw.append(np.linalg.norm(cov(X)-R,'fro')/np.linalg.norm(R,'fro'))
            Z=rank_normalize(X)
            rank.append(np.linalg.norm(cov(Z)-R,'fro')/np.linalg.norm(R,'fro'))
            raw_var.append(cov(X)[0,0]); rank_var.append(cov(Z)[0,0])
        rows.append({'A':A,'raw_relerr':float(np.mean(raw)),'rank_relerr':float(np.mean(rank)),
                     'raw_var0':float(np.mean(raw_var)),'rank_var0':float(np.mean(rank_var))})
    # predicted raw variance asymptote is q(1-q) A^2; fit coefficient at large A.
    As=np.array([r['A'] for r in rows[-3:]],float)
    vs=np.array([r['raw_var0'] for r in rows[-3:]])
    coef=float(np.polyfit(As**2,vs,1)[0])
    return rows, coef, q*(1-q)

def experiment_spike():
    # Rank scores plus eigenvalue-specific pilot cleaning: bulk spread shrinks,
    # while a separated top direction is intentionally retained.
    n,p,reps=128,12,30
    rows=[]
    for spike in [1,2,4,8]:
        rawe=[]; cle=[]; topraw=[]; topclean=[]
        C,Q,vals=latent_cov(p,spike)
        for _ in range(reps):
            G=rng.multivariate_normal(np.zeros(p), C, size=n)
            # marginal nonlinearities make Pearson S especially misleading
            X=np.sign(G)*(np.exp(np.minimum(np.abs(G),4))-1)
            S=cov(rank_normalize(X))
            Cc,w,wc=spectral_clean(S,1)
            rawe.append(np.linalg.norm(S-C,'fro'))
            cle.append(np.linalg.norm(Cc-C,'fro'))
            topraw.append(w[-1]); topclean.append(wc[-1])
        rows.append({'spike':spike,'raw_cov_error':float(np.mean(rawe)),
                     'clean_cov_error':float(np.mean(cle)), 'top_raw':float(np.mean(topraw)),
                     'top_clean':float(np.mean(topclean)), 'clean_win':float(np.mean(cle)<np.mean(rawe))})
    return rows



def invsqrt_psd(S, eps=1e-3):
    w, V = np.linalg.eigh(S)
    scale = max(float(np.mean(w)), 1e-8)
    w = np.maximum(w, eps * scale)
    return (V * (1.0 / np.sqrt(w))) @ V.T

def preprocess_batch(X, mode):
    if mode == 'raw':
        return X
    if mode == 'diag':
        d = np.sqrt(np.mean((X-X.mean(0))**2, axis=0) + 1e-3)
        return (X-X.mean(0)) / d
    Z = rank_normalize(X)
    S = cov(Z)
    Cclean, _, _ = spectral_clean(S, preserve_top=1)
    # Use the specified cleaned covariance to whiten activations.
    return (X-X.mean(0)) @ invsqrt_psd(Cclean, eps=1e-2)

def experiment_optimization():
    # Online linear regression with heavy-tailed, contaminated activations.
    # Same pre-generated batches are used by every method.
    p, batch, steps, reps = 16, 64, 350, 8
    lr = 0.025
    checkpoints = [50, 100, 200, 350]
    allres = {m: [] for m in ['raw','diag','rank_spectral']}
    for rep in range(reps):
        local = np.random.default_rng(SEED + 1000 + rep)
        beta = local.normal(size=p); beta /= np.linalg.norm(beta)
        batches=[]
        for _ in range(steps):
            G = local.normal(size=(batch,p))
            # Correlated latent directions, monotone heavy tails, plus rare shifts.
            G[:,1] = .65*G[:,0] + .76*G[:,1]
            X = np.sign(G) * (np.exp(np.minimum(np.abs(G), 4.0))-1.0)
            bad = local.random(batch) < .05
            X[bad,0] += local.choice([-1,1], size=bad.sum()) * 15.0
            y = X @ beta + .05*local.normal(size=batch)
            batches.append((X,y))
        for mode in allres:
            w=np.zeros(p); losses=[]
            for step,(X,y) in enumerate(batches,1):
                Xt=preprocess_batch(X, mode)
                pred=Xt@w
                err=pred-y
                losses.append(float(np.mean(err*err)))
                w -= lr * (Xt.T@err)/batch
                if step in checkpoints:
                    allres[mode].append((step, losses[-1]))
    summary={}
    for mode, vals in allres.items():
        summary[mode]={str(k): float(np.mean([v for s,v in vals if s==k])) for k in checkpoints}
    return summary

# Preserve the original verification output and add the optimization result.
if __name__ == '__main__':
    t=time.time()
    conv,slope=experiment_copula_convergence()
    out,coef,pred=experiment_outliers()
    spike=experiment_spike()
    opt=experiment_optimization()
    result={'seed':SEED,'predictions':{
      'P1_rank_sampling_scaling':'rank covariance error should scale as n^-1/2 after monotone marginal distortion',
      'P2_outlier_influence':'raw variance coefficient versus A^2 should approach q(1-q)=%.4f; rank variance should remain bounded'%pred,
      'P3_eigen_specific_cleaning':'bulk contraction should reduce covariance error while retaining top spike'},
      'observed':{'P1_rows':conv,'P1_loglog_slope':slope,'P2_rows':out,'P2_raw_A2_coefficient':coef,'P2_predicted_coefficient':pred,'P3_rows':spike,'optimization_mse':opt},
      'elapsed_sec':time.time()-t}
    print(json.dumps(result, indent=2))