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

SEED = 297
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'


def make_operator(d=32):
    r = np.random.default_rng(SEED)
    qm2 = r.normal(size=(d,d))/math.sqrt(d) + .75*np.eye(d)
    qm1 = r.normal(size=(d,d))/math.sqrt(d)
    q0 = r.normal(size=(d,d))/math.sqrt(d)
    return qm2, qm1, q0


def math_check():
    qm2,qm1,q0 = make_operator(); r=np.random.default_rng(SEED+1)
    v=r.normal(size=32); v/=np.linalg.norm(v)
    ds=np.array([1e-1,3e-2,1e-2,3e-3,1e-3,3e-4,1e-4,3e-5,1e-5])
    unsafe=[]; safe=[]; errors=[]
    for d in ds:
        q=qm2/d**2+qm1/d+q0
        unsafe.append(np.linalg.norm(q@v))
        safe.append(np.linalg.norm(q@(d*d*v)))
        errors.append(np.linalg.norm(q@(d*d*v)-qm2@v))
    unsafe_slope=np.polyfit(np.log(ds),np.log(unsafe),1)[0]
    safe_slope=np.polyfit(np.log(ds),np.log(safe),1)[0]
    return {'unsafe_growth_slope':float(unsafe_slope),'safe_variation_slope':float(safe_slope),
            'safe_limit_error_at_1e-5':float(errors[-1]),'unsafe_norm_ratio':float(unsafe[-1]/unsafe[0]),
            'safe_norm_ratio':float(safe[-1]/safe[0])}


class Branch(nn.Module):
    def __init__(self,d=32, mode='unsafe'):
        super().__init__(); self.mode=mode
        self.h=nn.Sequential(nn.Linear(8,32),nn.Tanh(),nn.Linear(32,d))
    def forward(self,x,z):
        v=self.h(x)
        if self.mode=='safe':
            # Fixed beta=1, order m=2; normalize only to prevent feature scale cheating.
            v=v/(v.norm(dim=1,keepdim=True)+1e-6)
            return v*(z-1.0).pow(2)
        return v


def run_branch(mode, q, train_x, train_z, train_y, val_x, val_z, val_y, steps=450):
    model=Branch(mode=mode).to(DEVICE)
    opt=torch.optim.Adam(model.parameters(),lr=2e-3)
    qm2,qm1,q0=[torch.tensor(a,dtype=torch.float32,device=DEVICE) for a in q]
    def forward(x,z):
        psi=model(x,z); t=z-1.0
        return (psi@qm2.T)/(t*t)+(psi@qm1.T)/t+psi@q0.T
    maxact=0.; maxgrad=0.; nan_count=0
    for _ in range(steps):
        opt.zero_grad(); pred=forward(train_x,train_z)
        loss=((pred-train_y)**2).mean()
        if not torch.isfinite(loss): nan_count+=1; break
        loss.backward()
        g=torch.nn.utils.clip_grad_norm_(model.parameters(),1e3) if mode=='clip' else 0.
        maxgrad=max(maxgrad,float(g)); maxact=max(maxact,float(pred.detach().norm(dim=1).max()))
        opt.step()
    with torch.no_grad():
        pred=forward(val_x,val_z); val_loss=((pred-val_y)**2).mean()
        near=(torch.abs(val_z-1)<0.002).squeeze(1)
        near_act=float(pred[near].norm(dim=1).max())
    return {'val_mse':float(val_loss),'max_output_norm':maxact,'near_pole_output_norm':near_act,
            'max_grad_norm':maxgrad,'nan_steps':nan_count,'parameters':sum(p.numel() for p in model.parameters())}


def experiment():
    qm2,qm1,q0=make_operator(); rng=np.random.default_rng(SEED+2)
    ntr,nva=256,256; d=32
    xtr=torch.tensor(rng.normal(size=(ntr,8)),dtype=torch.float32,device=DEVICE)
    xva=torch.tensor(rng.normal(size=(nva,8)),dtype=torch.float32,device=DEVICE)
    # Deliberately include very near-pole frequencies to test stability.
    ztr=torch.tensor(1+rng.uniform(-.05,.05,ntr),dtype=torch.float32,device=DEVICE).view(-1,1)
    zva=torch.tensor(1+rng.uniform(-.05,.05,nva),dtype=torch.float32,device=DEVICE).view(-1,1)
    target=nn.Sequential(nn.Linear(8,d),nn.Tanh(),nn.Linear(d,d)).to(DEVICE)
    with torch.no_grad():
        ytr=target(xtr); yva=target(xva)
    q=(qm2,qm1,q0); out={}
    # clip is a standard control; its architecture is unconstrained and differs only in update clipping.
    for mode in ('unsafe','clip','safe'):
        out[mode]=run_branch(mode,q,xtr,ztr,ytr,xva,zva,yva)
    return out


def main():
    try:
        math_result=math_check(); exp=experiment()
    except RuntimeError as e:
        if 'CUDA' in str(e) or 'cuda' in str(e).lower():
            global DEVICE; DEVICE='cpu'; math_result=math_check(); exp=experiment()
        else: raise
    result={'seed':SEED,'device':DEVICE,'math_check':math_result,'experiment':exp}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()