import math, random, json
import numpy as np
import torch

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

def poles_from_corr(c, rank=2, ridge=1e-6):
    c = np.asarray(c, dtype=np.float64)
    m = min(rank, (len(c)-1)//2)
    H0 = np.array([[c[i+j] for j in range(m)] for i in range(m)])
    H1 = np.array([[c[i+j+1] for j in range(m)] for i in range(m)])
    R = np.linalg.solve(H0.T @ H0 + ridge*np.eye(m), H0.T @ H1)
    return np.linalg.eigvals(R), H0

def math_check():
    true = np.array([.93, .62]); weights = np.array([.7, .3]); t=np.arange(40)
    c=sum(w*r**t for w,r in zip(weights,true))
    est,H=poles_from_corr(c,2); est=est[np.argsort(-np.abs(est))]
    mods=np.abs(est); dominant=float(mods[0])
    half_pred=math.log(.5)/math.log(dominant)
    half_emp=int(np.where(np.abs(c/c[0])<=.5)[0][0])
    return {"true_poles":true.tolist(), "estimated_moduli":mods.tolist(),
      "dominant_tau":float(-1/math.log(dominant)), "predicted_half_life":float(half_pred),
      "empirical_half_life_steps":half_emp, "hankel_condition":float(np.linalg.cond(H)),
      "pole_recovery_max_error":float(np.max(np.abs(np.sort(mods)[::-1]-true)))}

class LinearMemory(torch.nn.Module):
    def __init__(self,n=8):
        super().__init__()
        self.W=torch.nn.Parameter(.98*torch.eye(n)+.04*torch.randn(n,n))
        self.inp=torch.nn.Parameter(.2*torch.randn(n)); self.out=torch.nn.Parameter(.2*torch.randn(n))
    def forward(self,x):
        h=torch.zeros(x.shape[0],self.W.shape[0],device=x.device); ys=[]
        for k in range(x.shape[1]):
            h=torch.tanh(h+x[:,k,0:1]*self.inp); ys.append((h*self.out).sum(1))
        return torch.stack(ys,1)

def one_task(use_res, device, steps=350):
    torch.manual_seed(SEED + int(use_res)); model=LinearMemory().to(device)
    opt=torch.optim.Adam(model.parameters(),lr=.012); rng=np.random.default_rng(SEED+int(use_res))
    pole_log=[]
    for step in range(steps):
        x=rng.normal(size=(64,24,1)).astype('float32'); target=np.zeros((64,24),dtype='float32'); target[:,-1]=x[:,11,0]
        xb=torch.tensor(x,device=device); yb=torch.tensor(target,device=device)
        pred=model(xb); loss=((pred-yb)**2).mean(); total=loss
        if use_res and step%5==0:
            with torch.no_grad():
                h=torch.zeros(64,8,device=device); states=[]
                for _ in range(18):
                    h=torch.tanh(h+.15*torch.randn(64,1,device=device)*model.inp); states.append(h[:,0])
                cc=torch.stack([(states[0]*states[t]).mean() for t in range(18)])
                poles,_=poles_from_corr(cc.cpu().numpy(),rank=3)
                dominant=float(np.max(np.abs(poles))); pole_log.append(dominant)
            # Differentiable proxy: penalize excess operator norm only when fitted pole is too slow.
            excess=max(0.,dominant-.92)
            total=loss + .08*excess**2*torch.linalg.matrix_norm(model.W,ord=2)**2
        opt.zero_grad(); total.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),5); opt.step()
    with torch.no_grad():
        x=rng.normal(size=(256,24,1)).astype('float32'); target=torch.tensor(x[:,11,0],device=device)
        pred=model(torch.tensor(x,device=device)); mse=float(((pred[:,-1]-target)**2).mean().cpu())
        rho=float(np.max(np.abs(np.linalg.eigvals(model.W.detach().cpu().numpy()))))
    return {"mse":mse,"spectral_radius":rho,"mean_fitted_pole":float(np.mean(pole_log)) if pole_log else None,
            "max_fitted_pole":float(np.max(pole_log)) if pole_log else None,"device":device}

def run_task(use_res):
    preferred='cuda' if torch.cuda.is_available() else 'cpu'
    if preferred=='cuda':
        try: return one_task(use_res,'cuda')
        except Exception as e:
            torch.cuda.empty_cache(); fallback=one_task(use_res,'cpu'); fallback['cuda_fallback']=str(e); return fallback
    return one_task(use_res,'cpu')

def main():
    result={"math_check":math_check(),"task_baseline":run_task(False),"task_resonance":run_task(True)}
    with open('results.json','w') as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))
if __name__=='__main__': main()