"""Reproduce the original100-rack base sweeps and common signed-ridge shifts.

Source policy: MisfitLecture generators at revision a6c0c1f7cfce99745c17aaf4f5caa46f0bf66789.
Writes base-reproduced.json and base-samples.npz beside this script; never edits
base.json or the blog. Retains all original draws and the52 candidate paths.
"""
from pathlib import Path
import json, math, time
import numpy as np
HERE=Path(__file__).resolve().parent
CANDIDATES=np.r_[-np.linspace(0,.9,19),-np.linspace(.92,.995,16),np.logspace(-2,2,17)]
def sig7(a):
    return [float(f"{x:.7g}") for x in a]

def gram_rules(X, y, a, variance, fractions=CANDIDATES):
    """Independent eigensystem contractions of population and training risks."""
    n, p = X.shape
    if p == n:
        # Exactly square racks can be arbitrarily ill-conditioned. Use SVD here.
        u, s, vt = np.linalg.svd(X, full_matrices=False)
        d = s*s
        z = u.T @ y
        survival_basis = vt @ a
        coeff = s[:, None] * z[:, None] / (d[:, None] + d.min()*fractions)
        risk = (survival_basis @ coeff - 1)**2 + variance*np.sum(coeff*coeff, axis=0)
        residual = (X @ (vt.T @ coeff)) - y[:, None]
        train = np.mean(residual*residual, axis=0)
    elif p > n:
        G = np.einsum("ik,jk->ij", X, X)
        d, u = np.linalg.eigh(G)
        z = u.T @ y
        t = u.T @ (X @ a)
        alpha = z[:, None] / (d[:, None] + d.min()*fractions)
        risk = (t @ alpha - 1)**2 + variance*np.sum(d[:, None]*alpha*alpha, axis=0)
        train = np.sum((d.min()*fractions[None, :]*alpha)**2, axis=0)/n
    else:
        G = np.einsum("ik,il->kl", X, X)
        d, u = np.linalg.eigh(G)
        z = u.T @ (X.T @ y)
        t = u.T @ a
        coeff = z[:, None] / (d[:, None] + d.min()*fractions)
        risk = (t @ coeff - 1)**2 + variance*np.sum(coeff*coeff, axis=0)
        residual = X @ (u @ coeff) - y[:, None]
        train = np.mean(residual*residual, axis=0)
    if not np.isfinite(risk).all() or not np.isfinite(train).all() or d.min() <= 0:
        raise ArithmeticError("Invalid rack; abort instead of filtering/replacing it")
    return risk, train, float(d.min())

def reconstruct_original(n, variance, widths, seed, draws=100):
    """Original lecture RNG policy, independently calculated rules.

    rng.choice +/-1 sensor responses, Gaussian labels, Gaussian noise; consume
    original SGD rng.permutation(n) once after each rack to preserve the stream.
    This is a known source seed/policy, not a seed search to fit existing values.
    """
    rng = np.random.default_rng(seed)
    risk, train, smallest = [], [], []
    for p in widths:
        risks, trains, ds = [], [], []
        for draw in range(draws):
            a = rng.choice([-1., 1.], p)
            y = rng.standard_normal(n)
            X = y[:, None]*a + math.sqrt(variance)*rng.standard_normal((n, p))
            r, t, d = gram_rules(X, y, a, variance)
            risks.append(r); trains.append(t); ds.append(d)
            rng.permutation(n)
        risk.append(risks); train.append(trains); smallest.append(ds)
    return np.array(risk), np.array(train), np.array(smallest)

def summarize(widths, risk, train):
    medians = np.median(risk, axis=1)
    indices = np.argmin(medians, axis=1)
    chosen_risk = np.stack([risk[i, :, j] for i, j in enumerate(indices)])
    chosen_train = np.stack([train[i, :, j] for i, j in enumerate(indices)])
    interp_train = np.median(train[:, :, 0], axis=1)
    return dict(p=list(widths), lo=sig7(np.quantile(risk[:, :, 0], .25, axis=1)),
                interp=sig7(np.median(risk[:, :, 0], axis=1)), hi=sig7(np.quantile(risk[:, :, 0], .75, axis=1)),
                ridge=sig7(np.median(chosen_risk, axis=1)), lambda_fraction=sig7(CANDIDATES[indices]),
                negativeFraction=np.mean(CANDIDATES[np.argmin(risk, axis=2)] < 0, axis=1).tolist(),
                interpTrain=sig7(interp_train), ridgeTrain=sig7(np.median(chosen_train, axis=1)))

def main():
    expected=json.loads((HERE/'base.json').read_text()); output={};raw={};checked=0;start=time.perf_counter()
    for key,c in expected.items():
        seed=4000+100*c['n']+10*c['variance']
        r,t,d=reconstruct_original(c['n'],c['variance'],c['p'],seed)
        s=summarize(c['p'],r,t)
        for i,p in enumerate(c['p']):
            if p>=c['n']: s['interpTrain'][i]=0.
        for field in ['lo','interp','hi','ridge','lambda','negativeFraction','interpTrain','ridgeTrain']:
            for old,new in zip(c[field],s['lambda_fraction' if field=='lambda' else field]):
                if not math.isclose(old,new,rel_tol=1.1e-6,abs_tol=2e-9):
                    raise ArithmeticError(f"Fixture mismatch {key} {field}: {old} vs {new}")
                checked+=1
        output[key]=dict(n=c['n'],variance=c['variance'],seed=seed,racks=100,summary=s)
        raw[key+'_risk']=r;raw[key+'_train']=t;raw[key+'_minimum_eigenvalue']=d
    (HERE/'base-reproduced.json').write_text(json.dumps(output,indent=2)+'\n')
    np.savez_compressed(HERE/'base-samples.npz',**raw)
    print(f"{checked}/{checked} displayed base values reproduce at printed precision; {time.perf_counter()-start:.2f}s")
if __name__=='__main__': main()