import json
import numpy as np
from residual_screened_koopman import eig_residuals, screen_mask, spectral_forecast


def trajectory(lams, x0, T):
    z=np.zeros((T,len(lams))); z[0]=x0
    for t in range(T-1): z[t+1]=lams*z[t]
    return z


def run():
    # Prediction 1: exact eigenpairs have zero residual without corruption.
    lams=np.array([0.8,0.98]); clean=trajectory(lams,[1.,1.],500)
    r0=eig_residuals(np.diag(lams),clean)[2]

    # Prediction 2: iid observation noise makes residual energy scale as sigma^2.
    sigmas=np.array([0.,.01,.02,.04,.08]); obs=[]
    for j,s in enumerate(sigmas):
        a=[]
        for q in range(40):
            g=np.random.default_rng(1000+j*100+q)
            y=clean+s*g.normal(size=clean.shape)
            a.append(np.mean(eig_residuals(np.diag(lams),y)[2]**2))
        obs.append(float(np.mean(a)))
    slope=float(np.polyfit(sigmas[1:4]**2,np.array(obs[1:4]),1)[0])
    # Theoretical small-noise slope, accounting for finite clean denominator.
    preds=[]
    n=len(clean)-1
    for s in sigmas:
        p=[]
        for i,lam in enumerate(lams):
            den=np.sum(clean[:-1,i]**2)
            p.append(s*s*(1+lam*lam)*n/(den+s*s*n))
        preds.append(float(np.mean(p)))

    # Prediction 3: sparse corruption produces residual energy approximately
    # proportional to corruption rate at fixed impulse amplitude.
    rates=np.array([0.,.002,.005,.01,.02,.04]); impulse=.25; rate_obs=[]
    for j,p in enumerate(rates):
        a=[]
        for q in range(40):
            g=np.random.default_rng(5000+j*100+q)
            hit=(g.random(clean.shape)<p)*g.choice([-1.,1.],size=clean.shape)
            y=clean+impulse*hit
            a.append(np.mean(eig_residuals(np.diag(lams),y)[2]**2))
        rate_obs.append(float(np.mean(a)))
    rate_slope=float(np.polyfit(rates[1:4],rate_obs[1:4],1)[0])
    rate_preds=[]
    n=len(clean)-1
    for p in rates:
        per=[]
        for i,lam in enumerate(lams):
            den=np.sum(clean[:-1,i]**2)
            per.append(p*impulse**2*(1+lam*lam)*n/(den+p*impulse**2*n))
        rate_preds.append(float(np.mean(per)))

    # Held-out screening: mode 0 is reliable, mode 1 is corrupted in validation.
    train=clean[:220].copy(); val=clean[220:].copy()
    g=np.random.default_rng(77)
    train[:,1]+=0.10*g.normal(size=len(train)); val[:,1]+=0.10*g.normal(size=len(val))
    K=np.diag(lams)
    rv=eig_residuals(K,val)[2]; mask,tau=screen_mask(rv,.5)
    z0=val[0]; truth=clean[221:271]
    full=spectral_forecast(K,z0,50)
    screened=spectral_forecast(K,z0,50,mask)
    mse_full=float(np.mean((full-truth)**2)); mse_screen=float(np.mean((screened-truth)**2))
    result={'zero_residual':r0.tolist(),'variance':{'sigma':sigmas.tolist(),'observed':obs,'predicted':preds,'small_noise_slope':slope},'impulse_rate':{'rate':rates.tolist(),'observed':rate_obs,'predicted':rate_preds,'small_rate_slope':rate_slope},'screening':{'validation_residuals':rv.tolist(),'tau':float(tau),'mask':mask.tolist(),'retained':int(mask.sum()),'full_rollout_mse':mse_full,'screened_rollout_mse':mse_screen}}
    print(json.dumps(result,indent=2))
    with open('results.json','w') as f: json.dump(result,f,indent=2)

if __name__=='__main__': run()