import json
import numpy as np
from spectral_experiment import recovery_sweep

# For y=rho*(z^2-1)+eps, E[(y-Ey)(z^2-1)] = 2*rho exactly,
# since E[(z^2-1)^2]=2 and eps is independent.  This is the rank-one
# population spike predicted by Stein's identity.
def population_scaling(rhos=(.125,.25,.5,1.,2.), N=500000, noise=.35):
    rng=np.random.default_rng(9917)
    z=rng.normal(size=N); eps=rng.normal(size=N)
    rows=[]
    for rho in rhos:
        y=rho*(z*z-1)+noise*eps
        # No clipping here: this isolates the population identity.
        alpha=np.mean((y-y.mean())*(z*z-1))
        rows.append({'rho':rho,'predicted_2rho':2*rho,'measured_alpha':float(alpha),
                     'relative_error':float(abs(alpha-2*rho)/(2*rho))})
    # Independent labels imply zero signal coefficient.
    y0=rng.normal(size=N)
    null=np.mean((y0-y0.mean())*(z*z-1))
    return rows,float(null)

def transition_prediction(d=96, rho=.5, noise=.35):
    # Leading-order signed-Wishart heuristic: alpha=2rho, weighted-noise
    # variance s2=E[(y-Ey)^2]. The predicted outlier boundary is
    # gamma=d/n < alpha^2/s2, i.e. n/d > s2/alpha^2.
    alpha=2*rho
    s2=2*rho*rho+noise*noise
    critical=s2/(alpha*alpha)
    return {'rho':rho,'alpha_predicted':alpha,'weight_variance_predicted':s2,
            'predicted_critical_n_over_d':critical}

def main():
    scaling,null=population_scaling()
    trans=[]
    for rho in (.5,1.): trans.append(transition_prediction(rho=rho))
    # Reuse the actual estimator sweep and identify first ratio with overlap
    # clearly above the null 1/d baseline (a conservative finite-size onset).
    rec=recovery_sweep(d=96,rhos=(.5,1.),ratios=(.5,1,2,4,8),trials=8)
    onset=[]
    for rho in (.5,1.):
        rr=[x for x in rec if x['rho']==rho]
        hit=next((x['n_over_d'] for x in rr if x['overlap_mean']>4/96),None)
        onset.append({'rho':rho,'observed_onset_n_over_d':hit,
                      'null_overlap_prediction':1/96})
    out={'population_scaling':scaling,'null_alpha':null,
         'transition_predictions':trans,'observed_recovery_onsets':onset,
         'recovery_sweep':rec}
    with open('mechanism_results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()