import json, math, hashlib
from pathlib import Path
import numpy as np
from experiment import make_projection, alpha_calibrate, pair_distortion, variance_sweep, PHI, orbit_coefficients, corr_decay, rng, SEED


def gaussian_projection(d,m):
    return np.random.default_rng(SEED + 99).normal(size=(m,d))/math.sqrt(m)


def main():
    d,m=24,16
    X=rng.normal(size=(18,d)); X -= X.mean(0); X /= X.std(0)
    # Same calibration set and scalar norm calibration for deterministic and Gaussian matrices.
    P,_,_=make_projection(d,m,0.0314159,gap=1,coeff_n=4000)
    a=alpha_calibrate(P,X)
    det=(a,)+pair_distortion(a*P,X)
    G=gaussian_projection(d,m)
    ag=alpha_calibrate(G,X)
    gauss=(ag,)+pair_distortion(ag*G,X)
    # Search as in the proposal.
    best=None
    for si in range(30):
        seed=(0.0314159 + si*0.2718281)%1
        for g in [1,2,4,8,16,32]:
            Q,_,_=make_projection(d,m,seed,gap=g,coeff_n=4000)
            aq=alpha_calibrate(Q,X); md,ad=pair_distortion(aq*Q,X)
            item=(md,ad,si,g,aq)
            if best is None or item < best: best=item
    # Correlation fit only after the visible transient, and report all lags.
    z,mu,sig=orbit_coefficients(120000,0.3141592653)
    cs,_=corr_decay(z,500)
    lags=np.arange(4,33)
    fit=float(math.exp(np.polyfit(lags,np.log(np.maximum(np.abs(cs[lags-1]),1e-12)),1)[0]))
    var=variance_sweep(d=24,ms=(8,16,32,64,128),trials=700)
    slope=float(np.polyfit(np.log([v[0] for v in var]),np.log([v[2] for v in var]),1)[0])
    out={'seed':SEED,'dimensions':[d,m],
      'baseline_gaussian':{'alpha':gauss[0],'max_pair_distortion':gauss[1],'mean_pair_distortion':gauss[2]},
      'deterministic_fixed':{'alpha':det[0],'max_pair_distortion':det[1],'mean_pair_distortion':det[2]},
      'deterministic_searched':{'max_pair_distortion':best[0],'mean_pair_distortion':best[1],'seed_index':best[2],'gap':best[3],'alpha':best[4]},
      'correlation_prediction':{'global_fit_rho':float(corr_decay(z,500)[1]),'post_transient_fit_rho':fit,'lags_1_to_32':[(i,float(abs(cs[i-1]))) for i in range(1,33)]},
      'variance_prediction':{'m_var_values':var,'loglog_slope':slope},
      'reproducibility':hashlib.sha256(P.astype(np.float64).tobytes()).hexdigest()}
    Path('results.json').write_text(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))
if __name__=='__main__': main()