import json, math, hashlib
from pathlib import Path
import numpy as np

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

# Golden ratio is a Pisot number (its conjugate is -1/phi).
PHI = (1.0 + math.sqrt(5.0)) / 2.0


def beta_orbit(n, u0, beta=PHI):
    u = float(u0)
    out = np.empty(n, dtype=np.float64)
    for k in range(n):
        out[k] = u
        u = beta * u
        u -= math.floor(u)
    return out


def orbit_coefficients(n, u0, burn=10000):
    # Estimate empirical centering/scaling on a separate long orbit, as proposed.
    calib = beta_orbit(burn + n, u0, PHI)[burn:]
    mu, sigma = float(calib.mean()), float(calib.std())
    if sigma == 0:
        raise ValueError('degenerate orbit')
    vals = beta_orbit(n + burn, u0, PHI)[burn:]
    return (vals - mu) / sigma, mu, sigma


def make_projection(d, m, u0, gap=1, offset=0, coeff_n=300000):
    # Rows are z[offset + r*gap + i], exactly matching the proposed construction.
    need = offset + (m - 1) * gap + d + 1
    z, mu, sigma = orbit_coefficients(max(need, coeff_n), u0)
    P = np.empty((m, d), dtype=np.float64)
    for r in range(m):
        start = offset + r * gap
        P[r] = z[start:start+d] / math.sqrt(m)
    return P, mu, sigma


def pair_distortion(P, X):
    D = X[:, None, :] - X[None, :, :]
    orig = np.sum(D * D, axis=2)
    proj = np.einsum('nkd,md->nkm', D, P)
    proj = np.sum(proj * proj, axis=2)
    mask = orig > 1e-12
    ratios = proj[mask] / orig[mask]
    return float(np.max(np.abs(ratios - 1))), float(np.mean(np.abs(ratios - 1)))


def alpha_calibrate(P, X):
    a = np.sum(X * X)
    b = np.sum((X @ P.T) ** 2)
    return math.sqrt(a / b)


def corr_decay(z, maxlag=1000):
    z = z - z.mean()
    den = np.dot(z, z)
    cs = np.array([np.dot(z[:-k], z[k:]) / den for k in range(1, maxlag+1)])
    # Fit log envelope to lag 10..maxlag, robustly excluding zeros.
    lags = np.arange(1, maxlag+1)
    env = np.maximum.accumulate(np.abs(cs)[::-1])[::-1] # conservative not useful for fit
    sel = (lags >= 10) & (np.abs(cs) > 1e-5)
    slope = np.polyfit(lags[sel], np.log(np.abs(cs[sel])), 1)[0]
    return cs, float(math.exp(slope))


def variance_sweep(d=32, ms=(8,16,32,64,128), trials=3000):
    # Random isotropic vectors make the norm estimator's variance a direct JL diagnostic.
    out=[]
    for m in ms:
        vals=[]
        for t in range(trials):
            u0 = (0.017 + (t+1)*0.61803398875) % 1.0
            P,_,_ = make_projection(d,m,u0,gap=1,coeff_n=5000)
            x = rng.normal(size=d); x /= np.linalg.norm(x)
            vals.append(np.sum((P @ x)**2))
        vals=np.asarray(vals)
        out.append((m,float(vals.mean()),float(vals.var()),float(vals.var()*m)))
    return out


def main():
    d,m=24,16
    u0=0.3141592653
    z,mu,sigma=orbit_coefficients(120000,u0)
    cs,rho=corr_decay(z,500)
    # Prediction 1: gap-dependent correlation should decrease geometrically.
    gaps=[1,2,4,8,16,32]
    gap_corr=[]
    for g in gaps:
        # Correlation of same-column coefficients in adjacent rows.
        vals=np.array([np.corrcoef(z[:-g],z[g:])[0,1]])
        gap_corr.append((g,float(abs(vals[0])),float(abs(cs[g-1]))))

    # Calibration search prediction: optimize finite seed/gap and compare to fixed choice.
    X=rng.normal(size=(18,d)); X -= X.mean(0); X /= X.std(0)
    candidates=[]
    for si in range(30):
        seed=(0.0314159 + si*0.2718281) % 1.0
        for g in [1,2,4,8,16,32]:
            P,_,_=make_projection(d,m,seed,gap=g,coeff_n=3000)
            a=alpha_calibrate(P,X)
            md,ad=pair_distortion(a*P,X)
            candidates.append((md,ad,si,g,a))
    candidates.sort()
    best=candidates[0]
    fixed=[x for x in candidates if x[1] == min(y[1] for y in candidates)][:1]
    # fixed is not a fair baseline; use first seed gap=1 explicitly.
    P0,_,_=make_projection(d,m,0.0314159,gap=1,coeff_n=3000)
    a0=alpha_calibrate(P0,X); base=pair_distortion(a0*P0,X)

    var=variance_sweep()
    result={
      'seed':SEED,'beta':PHI,'orbit_mu':mu,'orbit_sigma':sigma,
      'prediction_correlation':{'claim':'|corr(z_t,z_{t+g})| decreases approximately geometrically with gap','fit_rho_per_step':rho,'observed':gap_corr},
      'prediction_variance':{'claim':'Var(||Px||^2) approximately proportional to 1/m','observed_m_mean_var_mvar':var,'loglog_slope':float(np.polyfit(np.log([x[0] for x in var]),np.log([x[2] for x in var]),1)[0])},
      'prediction_search':{'claim':'seed/gap calibration search lowers worst pairwise distortion','fixed_seed_gap1':{'max':base[0],'mean':base[1]},'best':{'max':best[0],'mean':best[1],'seed_index':best[2],'gap':best[3],'alpha':best[4]}},
      'reproducibility_hash':hashlib.sha256(np.asarray(make_projection(d,m,0.123456,gap=7,coeff_n=3000)[0],dtype=np.float64).tobytes()).hexdigest()
    }
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()