import json, math, random
from pathlib import Path
import numpy as np
from sklearn.metrics import roc_auc_score

SEED=2467
np.random.seed(SEED); random.seed(SEED)


def hankel_pair(y,p):
    y=np.asarray(y,float).reshape(-1)
    n=len(y)-2*p+1
    ym=np.stack([y[t:t+p] for t in range(n)],axis=1)
    yp=np.stack([y[t+p:t+2*p] for t in range(n)],axis=1)
    return ym,yp


def ridge_map(X,Y,lam=1e-8):
    return Y@X.T@np.linalg.inv(X@X.T+lam*np.eye(X.shape[0]))


def math_checks():
    # AR(2): delay coordinates are an exact finite Koopman invariant subspace.
    a,b=0.82,-0.21
    y=np.zeros(500); y[:2]=[0.7,-0.1]
    for t in range(2,len(y)): y[t]=a*y[t-1]+b*y[t-2]
    X,Y=hankel_pair(y,2); K=ridge_map(X,Y,1e-10)
    exact_rel=np.linalg.norm(Y-K@X)/np.linalg.norm(Y)
    # Additive observation noise: irreducible one-step residual should scale linearly in sigma.
    sigmas=np.array([0.002,0.005,0.01,0.02,0.05,0.1])
    rms=[]
    for s in sigmas:
        vals=[]
        for rep in range(12):
            yn=y+s*np.random.default_rng(1000+rep).normal(size=len(y))
            xx,yy=hankel_pair(yn,2); kk=ridge_map(xx,yy,1e-3)
            vals.append(np.sqrt(np.mean((yy-kk@xx)**2)))
        rms.append(float(np.mean(vals)))
    slope,inter=np.polyfit(sigmas,rms,1)
    pred_noise={"prediction":"Koopman residual is linear in observation noise sigma","slope":float(slope),"intercept":float(inter),"r2":float(1-np.sum((rms-(slope*sigmas+inter))**2)/np.sum((rms-np.mean(rms))**2)),"observed_pairs":[[float(s),float(r)] for s,r in zip(sigmas,rms)]}
    # Scalar latent rollout: log norm slope is log|a|, boundary is |a|=1.
    factors=np.array([0.7,0.9,0.98,1.0,1.02,1.1,1.3])
    slopes=[]
    H=60
    for q in factors:
        vals=np.abs(q**np.arange(H)+1e-30)
        slopes.append(float(np.polyfit(np.arange(H),np.log(vals),1)[0]))
    boundary_factor=float(factors[np.argmin(np.abs(np.array(slopes)))])
    stability={"prediction":"rollout growth-rate log(norm_t)/t equals log|a| and changes sign at |a|=1","factors":[float(x) for x in factors],"observed_log_slopes":slopes,"predicted_log_slopes":[float(np.log(x)) for x in factors],"estimated_boundary":boundary_factor,"predicted_boundary":1.0,"max_boundary_error":abs(boundary_factor-1.0)}
    # Hankel delay depth: exact AR(2) residual is zero at p>=2, but p=1 cannot represent it.
    ps=[1,2,3,4,6]
    delay_res=[]
    for p in ps:
        xx,yy=hankel_pair(y,p); kk=ridge_map(xx,yy,1e-9)
        delay_res.append(float(np.linalg.norm(yy-kk@xx)/np.linalg.norm(yy)))
    delay={"prediction":"delay depth p>=2 closes the AR(2) Koopman residual; p=1 does not","p":ps,"observed_relative_residual":delay_res,"p2_to_p4_ratio":delay_res[1]/max(delay_res[3],1e-15)}
    return {"noise_scaling":pred_noise,"stability_boundary":stability,"delay_transition":delay}


def make_data(ntraj=100,length=34,fault=False):
    rng=np.random.default_rng(88 if not fault else 99)
    out=[]
    for _ in range(ntraj):
        y=np.zeros(length); y[:2]=rng.normal(0,0.7,2)
        aa,bb=(0.82,-0.21) if not fault else (1.03,-0.21)
        for t in range(2,length): y[t]=aa*y[t-1]+bb*y[t-2]+rng.normal(0,0.025)
        for t in range(length-7):
            out.append((y[t:t+4].astype(np.float32),y[t+4:t+8].astype(np.float32)))
    return np.array([x for x,_ in out]),np.array([z for _,z in out])


def train_compare():
    import torch
    import torch.nn as nn
    device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        torch.manual_seed(SEED)
        if device=='cuda': torch.cuda.manual_seed_all(SEED)
        xp,xf=make_data(90)
        vp,vf=make_data(25)
        ap,af=make_data(35, fault=True)
        class Model(nn.Module):
            def __init__(self,dual):
                super().__init__(); self.dual=dual
                self.e=nn.Sequential(nn.Linear(4,24),nn.Tanh(),nn.Linear(24,3))
                self.dp=nn.Sequential(nn.Linear(3,24),nn.Tanh(),nn.Linear(24,4))
                self.df=nn.Sequential(nn.Linear(3,24),nn.Tanh(),nn.Linear(24,4))
            def forward(self,x):
                z=self.e(x); return self.dp(z),self.df(z)
        def fit(dual):
            m=Model(dual).to(device); opt=torch.optim.Adam(m.parameters(),lr=3e-3)
            X=torch.tensor(xp,device=device); F=torch.tensor(xf,device=device)
            g=torch.Generator(device=device); g.manual_seed(SEED)
            for step in range(700):
                ix=torch.randint(0,len(X),(128,),generator=g,device=device)
                hp,hf=m(X[ix]); loss=((hp-X[ix])**2).mean()+(1.0 if dual else 0.0)*((hf-F[ix])**2).mean()
                opt.zero_grad(); loss.backward(); opt.step()
            return m
        def score(m,p,f):
            with torch.no_grad():
                hp,hf=m(torch.tensor(p,device=device))
                # Both models are assessed on the same future prediction task; past AE has no trained future head.
                if m.dual: return ((hf-torch.tensor(f,device=device))**2).mean(1).cpu().numpy()
                return ((hp-torch.tensor(p,device=device))**2).mean(1).cpu().numpy()
        base=fit(False); dual=fit(True)
        # Calibrate each score on nominal validation, then compare fault-vs-nominal AUC.
        bnom=score(base,vp,vf); bfault=score(base,ap,af)
        dnom=score(dual,vp,vf); dfault=score(dual,ap,af)
        result={"device":device,"baseline_past_AUC":float(roc_auc_score(np.r_[np.zeros(len(bnom)),np.ones(len(bfault))],np.r_[bnom,bfault])),"dual_future_AUC":float(roc_auc_score(np.r_[np.zeros(len(dnom)),np.ones(len(dfault))],np.r_[dnom,dfault])),"baseline_nominal_mean":float(bnom.mean()),"baseline_fault_mean":float(bfault.mean()),"dual_nominal_future_mean":float(dnom.mean()),"dual_fault_future_mean":float(dfault.mean())}
        return result
    except Exception as e:
        return {"error":repr(e),"fallback":"training failed; math checks remain valid"}

if __name__=='__main__':
    result={"seed":SEED,"math":math_checks(),"mini_experiment":train_compare()}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))