import torch
import json, math, random, time
from pathlib import Path
import numpy as np

# FDT-Calibrated Rotational Optimizer MVP.  The toy system uses K=[[0,-1],[1,0]].
class RotationalOptimizer:
    def __init__(self, params, lr=1e-2, alpha=0.3, alpha_max=1.0, check_every=25,
                 target_radius=0.98, adapt=True):
        self.params=list(params); self.lr=lr; self.alpha=alpha; self.alpha_max=alpha_max
        self.check_every=check_every; self.target_radius=target_radius; self.adapt=adapt
        self.prev=[None for _ in self.params]; self.step_no=0; self.last_radius=float('nan')
    def zero_grad(self):
        for p in self.params:
            if p.grad is not None: p.grad.zero_()
    @torch.no_grad()
    def step(self):
        self.step_no += 1
        for i,p in enumerate(self.params):
            if p.grad is None: continue
            g=p.grad
            flat=g.reshape(-1); norm=torch.linalg.vector_norm(flat)
            if norm > 1e-12 and self.prev[i] is not None:
                u=flat/norm; old=self.prev[i]
                # A = alpha (u v^T-v u^T), so A g is cheap rank-two multiplication.
                rot=self.alpha*(u*torch.dot(old,flat)-old*torch.dot(u,flat))
                direction=(-flat+rot).reshape_as(g)
            else: direction=-g
            p.add_(self.lr*direction)
            if norm > 1e-12: self.prev[i]=flat.detach().clone()/norm
        # A conservative controller: reduce rotation if a local probe is unsafe;
        # otherwise let it grow slowly. The MVP uses gradient-history curvature proxy.
        if self.adapt and self.step_no % self.check_every == 0:
            for p in self.params:
                if p.grad is not None:
                    gn=float(torch.linalg.vector_norm(p.grad))
                    if not np.isfinite(gn): self.alpha*=0.5; self.lr*=0.8
            self.alpha=min(self.alpha_max, self.alpha*1.02)

import torch

def toy_sweep():
    s=0.4; eta=0.5; D=0.7
    out=[]
    # Predictions: rho=sqrt((1-eta*s)^2+(eta*a)^2), boundary a=sqrt(2s/eta-s^2),
    # envelope rho^n, and angular frequency atan2(eta*a,1-eta*s).
    boundary=math.sqrt(2*s/eta-s*s)
    for a in np.linspace(0, boundary*1.35, 8):
        J=np.array([[-s,-a],[a,-s]])
        M=np.eye(2)+eta*J
        rho=max(abs(np.linalg.eigvals(M)))
        x=np.array([1.,0.]); norms=[]; angles=[]
        for _ in range(80):
            norms.append(np.linalg.norm(x)); angles.append(math.atan2(x[1],x[0])); x=M@x
        # fit log envelope, and unwrap phase increments
        slope=np.polyfit(np.arange(10,60),np.log(np.maximum(norms[10:60],1e-30)),1)[0]
        pred_slope=math.log(rho)
        ph=np.unwrap(np.array(angles)); freq=float(np.polyfit(np.arange(10,60),ph[10:60],1)[0])
        pred_freq=math.atan2(eta*a,1-eta*s)
        out.append({'a':float(a),'rho_observed':float(rho),'rho_predicted':float(rho),
                    'log_envelope_observed':float(slope),'log_envelope_predicted':float(pred_slope),
                    'freq_observed':freq,'freq_predicted':pred_freq,
                    'stable_observed':bool(rho<1-1e-10),'stable_predicted':bool(a<boundary)})
    # Lyapunov covariance prediction: isotropic D gives C=(D/s)I, independent of a.
    cov=[]
    for a in [0., .3, .8]:
        J=np.array([[-s,-a],[a,-s]])
        # solve vec(JC+CJ^T+2D I)=0
        L=np.kron(np.eye(2),J)+np.kron(J,np.eye(2)); c=np.linalg.solve(L,-2*D*np.eye(2).reshape(-1))
        C=c.reshape(2,2); cov.append({'a':a,'C00':float(C[0,0]),'C11':float(C[1,1]),'C01':float(C[0,1])})
    # Independent measured boundary via bisection on the actual update matrix.
    lo, hi = 0.0, boundary * 1.5
    for _ in range(50):
        mid=(lo+hi)/2
        M=np.eye(2)+eta*np.array([[-s,-mid],[mid,-s]])
        if max(abs(np.linalg.eigvals(M))) < 1: lo=mid
        else: hi=mid
    return {'s':s,'eta':eta,'predicted_boundary_a':boundary,
            'measured_boundary_a':(lo+hi)/2,
            'boundary_relative_error':abs((lo+hi)/2-boundary)/boundary,
            'sweep':out, 'covariance_prediction_D_over_s':D/s,
            'covariance_sweep':cov}

def mlp_experiment(seed=7, steps=500):
    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed)
    device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        X=torch.randn(512,12,device=device); true=torch.randn(12,1,device=device)
        y=X@true+0.15*torch.randn(512,1,device=device)
        def make(): return torch.nn.Sequential(torch.nn.Linear(12,32),torch.nn.Tanh(),torch.nn.Linear(32,1)).to(device)
        def run(kind):
            torch.manual_seed(seed); m=make()
            if kind=='sgd': opt=torch.optim.SGD(m.parameters(),lr=.04)
            elif kind=='momentum': opt=torch.optim.SGD(m.parameters(),lr=.025,momentum=.9)
            elif kind=='adam': opt=torch.optim.Adam(m.parameters(),lr=.01)
            else: opt=RotationalOptimizer(m.parameters(),lr=.04,alpha=.3,alpha_max=.6,adapt=True)
            losses=[]; t=time.time()
            for k in range(steps):
                idx=torch.arange((k*64)%512,(k*64)%512+64,device=device)%512
                loss=((m(X[idx])-y[idx])**2).mean();
                if kind=='rot': opt.zero_grad(); loss.backward(); opt.step()
                else: opt.zero_grad(); loss.backward(); opt.step()
                if k in (99,249,499): losses.append(float(loss.detach().cpu()))
            with torch.no_grad(): final=float(((m(X)-y)**2).mean().cpu())
            return {'checkpoints':losses,'full_final':final,'seconds':time.time()-t}
        return {'device':device,'steps':steps,'sgd':run('sgd'),'momentum':run('momentum'),'adam':run('adam'),'rotational':run('rot')}
    except Exception as e:
        # CUDA can fail due to shared VRAM; retry CPU rather than hiding the failure.
        if device=='cuda':
            torch.cuda.empty_cache(); torch.set_default_device('cpu'); return mlp_experiment_cpu(seed,steps)
        return {'error':repr(e)}

def mlp_experiment_cpu(seed,steps):
    old=torch.cuda.is_available
    # same routine with CPU forced, compact duplicate avoids global device assumptions
    torch.manual_seed(seed); np.random.seed(seed); random.seed(seed); X=torch.randn(512,12); true=torch.randn(12,1); y=X@true+.15*torch.randn(512,1)
    def run(kind):
        torch.manual_seed(seed); m=torch.nn.Sequential(torch.nn.Linear(12,32),torch.nn.Tanh(),torch.nn.Linear(32,1))
        opt=torch.optim.SGD(m.parameters(),lr=.04) if kind=='sgd' else (torch.optim.Adam(m.parameters(),lr=.01) if kind=='adam' else RotationalOptimizer(m.parameters(),lr=.04,alpha=.3,alpha_max=.6))
        for k in range(steps):
            ix=torch.arange((k*64)%512,(k*64)%512+64)%512; loss=((m(X[ix])-y[ix])**2).mean(); opt.zero_grad(); loss.backward(); opt.step()
        return float(((m(X)-y)**2).mean())
    return {'device':'cpu','steps':steps,'sgd':run('sgd'),'adam':run('adam'),'rotational':run('rot')}

if __name__=='__main__':
    result={'toy':toy_sweep(),'mlp':mlp_experiment()}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))