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

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

# Quadratic target regularizer R(x)=lambda*x^2/2. Its proximal map is a*v,
# where a=1/(1+gamma*lambda). PnP-PGD uses A=I and f=.5||x-y||^2.
def prox_coeff(gamma, lam):
    return 1.0 / (1.0 + gamma * lam)

def pnp_run(a, gamma, lam, y, steps=60):
    x = 0.0
    hist = []
    for _ in range(steps):
        v = x - gamma * (x - y)
        x = a * v
        F = .5 * (x-y)**2 + .5 * lam*x*x
        station = abs((x-y) + lam*x)
        hist.append((x, F, station))
    return np.asarray(hist)

def main():
    out = {"seed": SEED}

    # ---------------- Stage 1: analytic mechanism verification ----------------
    # Prediction 1: prox coefficient is exactly 1/(1+gamma*lambda).
    coeff_rows=[]
    for gamma in [0.1, 0.3, 0.7, 1.0]:
        for lam in [0.0, 0.5, 2.0, 8.0]:
            v = rng.normal(size=20000)
            # empirical least-squares fit to teacher action
            t = prox_coeff(gamma, lam) * v
            fitted = float(np.dot(v,t)/np.dot(v,v))
            pred = prox_coeff(gamma, lam)
            coeff_rows.append({"gamma":gamma,"lambda":lam,"pred":pred,"fit":fitted,"abs_err":abs(fitted-pred)})
    max_coeff_err=max(r["abs_err"] for r in coeff_rows)

    # Prediction 2: squared mismatch is quadratic in coefficient error, slope 2 on log-log.
    v = rng.normal(size=100000)
    target_a = prox_coeff(.6, 3.0)
    errors=np.logspace(-4,-1,10)
    mismatch=[]
    for e in errors:
        mismatch.append(float(np.mean(((target_a+e)*v-target_a*v)**2)))
    slope=float(np.polyfit(np.log(errors),np.log(mismatch),1)[0])
    # Also verify lambda=0 gives identity teacher and identity coefficient has zero mismatch.
    zero_mismatch=float(np.mean((1.0*v-1.0*v)**2))

    # Prediction 3: PnP linear contraction factor q=a(1-gamma), boundary |q|=1.
    # With gamma=.5 and positive a, predicted boundary a=2.0. Measure whether
    # trajectory grows (unstable) or contracts from x0=1,y=0.
    gamma=.5; y=0.0
    boundary=1.0/(1.0-gamma)
    stability=[]
    for a in [0.5, 1.0, 1.8, 1.99, 2.01, 2.2, 3.0]:
        h=pnp_run(a,gamma,0.0,y,steps=25)
        ratio=abs(h[-1,0])/max(abs(h[0,0]),1e-12)
        q=a*(1-gamma)
        stability.append({"a":a,"predicted_q":q,"observed_abs_final_over_first":float(ratio),"stable_observed":bool(ratio<1)})
    # pnp_run starts x=0, so explicitly use recurrence for stability test from x0=1
    stability=[]
    for a in [0.5,1.0,1.8,1.99,2.01,2.2,3.0]:
        q=a*(1-gamma); vals=np.array([q**k for k in range(25)])
        stability.append({"a":a,"predicted_q":q,"observed_final_abs":float(abs(vals[-1])),"stable_observed":bool(abs(vals[-1])<1)})
    stable_boundary=max(r["a"] for r in stability if r["stable_observed"])
    unstable_boundary=min(r["a"] for r in stability if not r["stable_observed"])
    out["math_check"]={"coefficient_max_abs_error":max_coeff_err,
        "coefficient_rows":coeff_rows,"mismatch_loglog_slope":slope,
        "mismatch_quadratic_prediction":2.0,"lambda_zero_mismatch":zero_mismatch,
        "stability_predicted_boundary":boundary,"largest_stable_grid_a":stable_boundary,
        "smallest_unstable_grid_a":unstable_boundary,"stability_rows":stability}

    # ---------------- Stage 2: tiny practical adaptation experiment ----------------
    # Equal-capacity denoiser D_a(v)=a*v. Target regularizer is known to teacher,
    # while ordinary clean-MSE sees noisy clean-image pairs at another noise level.
    rng2=np.random.default_rng(SEED+1)
    n=4000; lam=2.0; gamma=.6; a_star=prox_coeff(gamma,lam)
    target_var=.25; sigma_mse=.35
    clean=rng2.normal(0,np.sqrt(target_var),n)
    mse_inputs=clean+rng2.normal(0,sigma_mse,n)
    mse_a=float(np.dot(mse_inputs,clean)/np.dot(mse_inputs,mse_inputs))
    # Generate actual PnP intermediate states from random target observations.
    ys=clean+rng2.normal(0,.45,n)
    states=[]; teachers=[]
    for y0 in ys:
        x0=0.0
        for k in range(12):
            v0=x0-gamma*(x0-y0)
            states.append(v0); teachers.append(a_star*v0)
            x0=a_star*v0
    states=np.asarray(states); teachers=np.asarray(teachers)
    prox_a=float(np.dot(states,teachers)/np.dot(states,states))
    # A mixed objective with equal weighting in normalized scalar regression.
    mixed_a=float(np.dot(states,teachers)+np.dot(mse_inputs,clean))/(
        np.dot(states,states)+np.dot(mse_inputs,mse_inputs))
    # Evaluate reconstruction of held-out y. Exact target optimum is y/(1+lambda).
    test_y=rng2.normal(0,1.0,1000)
    rows=[]
    for name,a in [("clean_mse",mse_a),("proximal_match",prox_a),("mixed",mixed_a),("ideal",a_star)]:
        losses=[]; errs=[]; stations=[]
        for y0 in test_y:
            h=pnp_run(a,gamma,lam,float(y0),steps=30)
            losses.append(h[-1,1]); stations.append(h[-1,2]); errs.append((h[-1,0]-y0/(1+lam))**2)
        rows.append({"method":name,"a":a,"final_F":float(np.mean(losses)),
                     "final_stationarity":float(np.mean(stations)),
                     "reconstruction_MSE_to_optimum":float(np.mean(errs)),
                     "prox_mismatch_on_states":float(np.mean((a*states-teachers)**2))})
    out["adaptation"]={"a_star":a_star,"clean_mse_coefficient":mse_a,
        "proximal_coefficient":prox_a,"mixed_coefficient":mixed_a,"results":rows}
    out["interpretation"]={
      "mechanism_confirmed": bool(max_coeff_err<2e-3 and abs(slope-2)<0.03 and zero_mismatch<1e-20 and stable_boundary < boundary < unstable_boundary),
      "practical_prox_better_than_clean": bool(rows[1]["final_F"] < rows[0]["final_F"] and rows[1]["final_stationarity"] < rows[0]["final_stationarity"]),
      "note":"The toy teacher is exact and quadratic; this isolates the proposed proximal-mismatch mechanism, not image-denoising generalization."}
    Path("results.json").write_text(json.dumps(out,indent=2))
    print(json.dumps({"math_check":out["math_check"],"adaptation":out["adaptation"],"interpretation":out["interpretation"]},indent=2))

if __name__ == '__main__': main()