import math, random, time
from collections import deque
import numpy as np
import torch
import torch.nn as nn

SEED = 7
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)

class HankelResidualObserver:
    def __init__(self, p=1, t_ini=8, horizon=8, max_cols=40, ridge=1e-3, tau=1.0):
        self.p, self.t_ini, self.horizon = p, t_ini, horizon
        self.max_cols, self.ridge, self.tau = p*0 + max_cols, ridge, tau
        self.residuals = deque(maxlen=t_ini + horizon + max_cols + 4)
    def update(self, residual):
        self.residuals.append(np.asarray(residual, dtype=np.float64).reshape(self.p))
    def predict(self):
        if len(self.residuals) < self.t_ini + self.horizon:
            return np.zeros((self.horizon, self.p)), 0.0, 0
        a = np.asarray(self.residuals)
        n = len(a) - self.t_ini - self.horizon + 1
        starts = np.arange(max(0, n - self.max_cols), n)
        dp = np.stack([a[j:j+self.t_ini].reshape(-1) for j in starts], axis=1)
        df = np.stack([a[j+self.t_ini:j+self.t_ini+self.horizon].reshape(-1) for j in starts], axis=1)
        d_ini = a[-self.t_ini:].reshape(-1)
        g = np.linalg.solve(dp.T @ dp + self.ridge*np.eye(dp.shape[1]), dp.T @ d_ini)
        recon = dp @ g
        q = np.linalg.norm(recon-d_ini)/(np.linalg.norm(d_ini)+1e-8)
        gamma = float(np.clip(1.0-q/self.tau, 0.0, 1.0))
        return (df @ g).reshape(self.horizon, self.p), gamma, dp.shape[1]

class SmallGRU(nn.Module):
    def __init__(self, hidden=24, horizon=8):
        super().__init__(); self.gru=nn.GRU(1,hidden,batch_first=True); self.head=nn.Linear(hidden,horizon)
    def forward(self,x):
        z,_=self.gru(x); return self.head(z[:,-1]).unsqueeze(-1)

def make_train(n=5000,w=16,h=8):
    y=np.zeros(n+w+h+1); noise=np.random.default_rng(SEED).normal(0,.08,len(y))
    for t in range(1,len(y)): y[t]=.82*y[t-1]+noise[t]
    return (np.asarray([y[t:t+w] for t in range(n)],np.float32)[...,None],
            np.asarray([y[t+w:t+w+h] for t in range(n)],np.float32)[...,None])

def fit_model(device):
    X,Y=make_train(); model=SmallGRU().to(device); opt=torch.optim.Adam(model.parameters(),lr=3e-3); g=torch.Generator().manual_seed(SEED)
    for _ in range(18):
        perm=torch.randperm(len(X),generator=g)
        for i in range(0,len(X),128):
            ix=perm[i:i+128]; xb=torch.from_numpy(X[ix]).to(device); yb=torch.from_numpy(Y[ix]).to(device)
            loss=(model(xb)-yb).square().mean(); opt.zero_grad(); loss.backward(); opt.step()
    return model

def run_stream(model,device,disturbed=True,n=700,w=16,h=8):
    rng=np.random.default_rng(101 if disturbed else 102); y=np.zeros(n+w+h+2)
    for t in range(1,len(y)):
        d=.55*math.sin(.42*t+.4) if disturbed else 0.; y[t]=.82*y[t-1]+d+rng.normal(0,.035)
    obs=HankelResidualObserver(t_ini=8,horizon=h,max_cols=40,ridge=1e-2,tau=1.0)
    base_err=[]; corr_err=[]; gammas=[]; lat=[]; model.eval()
    with torch.no_grad():
        for t in range(w,w+n-h):
            x=torch.from_numpy(y[t-w:t].astype(np.float32)[None,:,None]).to(device); pred=model(x).cpu().numpy()[0,:,0]
            if t>w: obs.update([y[t]-prior_pred[0]])
            tic=time.perf_counter(); dh,gamma,_=obs.predict(); lat.append(time.perf_counter()-tic)
            target=y[t:t+h]; base_err.extend(pred-target); corr_err.extend(pred+gamma*dh[:,0]-target); gammas.append(gamma); prior_pred=pred
    be=np.asarray(base_err); ce=np.asarray(corr_err)
    return {'baseline_rmse':float(np.sqrt(np.mean(be**2))),'observer_rmse':float(np.sqrt(np.mean(ce**2))),
      'baseline_long_rmse':float(np.sqrt(np.mean(be.reshape(-1,h)[:,-3:]**2))), 'observer_long_rmse':float(np.sqrt(np.mean(ce.reshape(-1,h)[:,-3:]**2))),
      'mean_gamma':float(np.mean(gammas)),'p95_latency_ms':float(np.percentile(lat,95)*1000),'n_predictions':len(gammas)}

def math_check():
    # Exactly period-4 residuals: every length-8 past and future window is a dictionary shift.
    pattern=np.array([1.,.25,-1.,-.25]); r=np.resize(pattern,96)[:,None]
    o=HankelResidualObserver(t_ini=8,horizon=8,max_cols=40,ridge=1e-10,tau=1.)
    for x in r: o.update(x)
    pred,gamma,cols=o.predict(); truth=np.resize(pattern,8)[:,None]
    # Random unrelated history should fail reconstruction and suppress correction.
    z=np.random.default_rng(99).normal(size=(96,1)); qobs=HankelResidualObserver(t_ini=8,horizon=8,max_cols=40,ridge=1e-2,tau=1.)
    for x in z:qobs.update(x)
    _,bad_gamma,_=qobs.predict()
    return {'periodic_prediction_rmse':float(np.sqrt(np.mean((pred-truth)**2))), 'periodic_gamma':gamma,
            'random_mismatch_gamma':bad_gamma,'dictionary_columns':cols}

def main():
    device='cuda' if torch.cuda.is_available() else 'cpu'
    try: model=fit_model(device)
    except Exception: device='cpu'; model=fit_model(device)
    if device=='cuda': torch.cuda.empty_cache()
    print(__import__('json').dumps({'device':device,'math_check':math_check(),'disturbed':run_stream(model,device,True),'clean':run_stream(model,device,False)},indent=2))
if __name__=='__main__': main()