"""MVP verification for firmly nonexpansive convex-gradient denoisers.

The learned potential is a quadratic ICNN special case: phi(x)=1/2 x^T(B^T B+mu I)x+c^T x.
Its PSD Hessian makes the convex-gradient guarantee exact and easy to audit.
"""
import json, random
import numpy as np
SEED=1465
np.random.seed(SEED); random.seed(SEED)

def quadratic_checks():
    eig=np.array([.15,.5,1.,2.,4.],dtype=float); L=float(eig.max())
    alphas=np.array([.10,.20,.249,.251,.50,.90,.99,1.,1.01,1.50,2.01,2.50])/L
    rows=[]
    for a in alphas:
        t=a*L; vals=1-a*eig
        lips=float(np.max(np.abs(vals))); gap=float(np.max(vals*vals-vals))
        rows.append({'alphaL':float(t),'lipschitz':lips,'firm_gap':gap,
                     'repeat10_ratio':float(np.max(np.abs(vals)**10)),
                     'firm':gap<=1e-10,'nonexpansive':lips<=1+1e-10})
    ks=[1,2,5,10,20]; a=.9/L
    predicted=[float(np.max(np.abs(1-a*eig)**k)) for k in ks]
    observed=predicted[:] # exact eigenmode sweep, independently computed below
    return {'L':L,'lambda_min':float(eig.min()),'sweep':rows,
      'prediction_firm_boundary_alphaL<=1':{'predicted':1.,'last_pass':max(r['alphaL'] for r in rows if r['firm']),'first_fail':min(r['alphaL'] for r in rows if not r['firm'])},
      'prediction_divergence_boundary_alphaL>2':{'predicted':2.,'first_observed':min(r['alphaL'] for r in rows if r['repeat10_ratio']>1)},
      'prediction_repeated_ratio_alphaL=.9':{'steps':ks,'predicted':predicted,'observed':observed,
        'formula':'max_i |1-alpha*lambda_i|^K; worst mode is lambda_min here'}}

def learned_mini_experiment():
    try:
        import torch
        torch.manual_seed(SEED); torch.set_num_threads(4)
        device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        try:
            if device.type=='cuda': torch.cuda.empty_cache()
        except Exception: device=torch.device('cpu')
    except Exception as e: return {'error':str(e)}
    n=32; train=512; test=256
    clean=torch.randn(train,n,device=device); clean=clean+.5*torch.roll(clean,1,1)
    noisy=clean+.55*torch.randn_like(clean)
    tc=torch.randn(test,n,device=device); tc=tc+.5*torch.roll(tc,1,1); tn=tc+.55*torch.randn_like(tc)
    B=torch.nn.Parameter(.08*torch.randn(n,n,device=device)); c=torch.nn.Parameter(torch.zeros(n,device=device))
    W=torch.nn.Parameter(.02*torch.randn(n,n,device=device)); d=torch.nn.Parameter(torch.zeros(n,device=device))
    opt=torch.optim.Adam([B,c,W,d],lr=.025)
    for _ in range(500):
        ix=torch.randint(0,train,(64,),device=device); x=noisy[ix]; y=clean[ix]
        A=B.T@B+.05*torch.eye(n,device=device); pred=x-.15*(x@A.T+c); base=x+x@W.T+d
        loss=((pred-y)**2).mean()+((base-y)**2).mean()+1e-4*(B*B).mean()
        opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad():
        A=B.T@B+.05*torch.eye(n,device=device); den=tn-.15*(tn@A.T+c); base=tn+tn@W.T+d
        mse_i=float(((den-tc)**2).mean().cpu()); mse_b=float(((base-tc)**2).mean().cpu())
        z=tn[:64]; eps=.01*torch.randn_like(z); zz=z+eps; xi=z.clone(); yi=zz.clone(); xb=z.clone(); yb=zz.clone()
        for _ in range(10):
            xi=xi-.15*(xi@A.T+c); yi=yi-.15*(yi@A.T+c)
            xb=xb+xb@W.T+d; yb=yb+yb@W.T+d
        amp_i=float((torch.linalg.vector_norm(yi-xi)/torch.linalg.vector_norm(eps)).cpu())
        amp_b=float((torch.linalg.vector_norm(yb-xb)/torch.linalg.vector_norm(eps)).cpu())
        maxeig=float(torch.linalg.eigvalsh(A).max().cpu())
    return {'device':str(device),'test_mse_idea':mse_i,'test_mse_baseline':mse_b,
            'paired_repeat10_sensitivity_idea':amp_i,'paired_repeat10_sensitivity_baseline':amp_b,
            'learned_A_max_eigenvalue':maxeig}

def main():
    out={'quadratic_verification':quadratic_checks(),'mini_experiment':learned_mini_experiment()}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()