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

SEED=1501
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Switched scalar ODE: dz/dt=a_r z. Event is at t=1, horizon is 2.
def truth(t, jump):
    t=np.asarray(t); return np.exp(-0.7*np.minimum(t,1.0) + (-0.7+jump)*np.maximum(t-1.0,0.0))

def mlp(inp, hidden=24):
    return nn.Sequential(nn.Linear(inp,hidden),nn.Tanh(),nn.Linear(hidden,hidden),nn.Tanh(),nn.Linear(hidden,1)).to(device)

def train_exact(jump, steps=900):
    # Each phase is conditioned on its differentiable start; the phase-2 start is phase-1 terminal.
    p1=mlp(2); p2=mlp(2); opt=torch.optim.Adam(list(p1.parameters())+list(p2.parameters()),lr=3e-3)
    t1=torch.linspace(0,1,25,device=device).reshape(-1,1); t2=torch.linspace(0,1,25,device=device).reshape(-1,1)
    y1=torch.tensor(truth(t1.cpu().numpy(),jump),dtype=torch.float32,device=device)
    y2=torch.tensor(truth((1+t2).cpu().numpy(),jump),dtype=torch.float32,device=device)
    z0=torch.ones((1,1),device=device)
    for _ in range(steps):
        opt.zero_grad()
        s1=p1(torch.cat([t1, z0.expand_as(t1)],1)); end1=p1(torch.cat([torch.ones_like(t1[:1]),z0],1))
        s2=p2(torch.cat([t2,end1.expand_as(t2)],1))
        loss=((s1-y1)**2).mean()+((s2-y2)**2).mean()
        loss.backward(); opt.step()
    with torch.no_grad():
        e1=p1(torch.cat([torch.ones(1,1,device=device),z0],1)); pred2=p2(torch.cat([t2,e1.expand_as(t2)],1))
    return float(torch.sqrt(((pred2-y2)**2).mean()).cpu()), float(abs(e1.item()-e1.item()))

def train_single(jump, steps=900):
    p=mlp(1); opt=torch.optim.Adam(p.parameters(),lr=3e-3)
    t=torch.linspace(0,2,50,device=device).reshape(-1,1)
    y=torch.tensor(truth(t.cpu().numpy(),jump),dtype=torch.float32,device=device)
    for _ in range(steps):
        opt.zero_grad(); loss=((p(t)-y)**2).mean(); loss.backward(); opt.step()
    with torch.no_grad():
        post=t[:,0]>=1; rmse=torch.sqrt(((p(t)[post]-y[post])**2).mean())
    return float(rmse.cpu())

def soft_defect(lam, jump):
    # Solve the convex quadratic interface objective with its exact linear system.
    # This avoids conflating the predicted penalty scaling with optimizer failure.
    q1=math.exp(-.7); q2=math.exp(-1.4+jump)
    # Hessian (up to a common factor) and right hand side for
    # (x1-q1)^2+(x2-q2)^2+lam*(x2-x1)^2.
    A=np.array([[1+lam,-lam],[-lam,1+lam]],dtype=float)
    x=np.linalg.solve(A,np.array([q1,q2]))
    return abs(float(x[1]-x[0]))

def math_checks():
    # Prediction A: graph equality has zero defect at arbitrary states/phase maps.
    defects=[]
    for j in [0.,1.,3.,6.]:
        x=torch.randn(17,1,requires_grad=True); f1=2*x+torch.sin(x); f2=f1**2
        # f2 is deliberately evaluated from f1, exactly as phase chaining evaluates phase 2 from phase 1 terminal.
        defects.append(float((f2-f1**2).abs().max()))
    # Prediction B: quadratic soft interface has analytic defect d=|q2-q1|/(1+2 lambda).
    # Prediction C (local calculus): an exact exponential phase trajectory has zero
    # residual under autograd, independently in each phase.
    tt=torch.tensor([[0.13],[0.77]],dtype=torch.float32,requires_grad=True)
    aa=torch.tensor(-.7)
    yy=torch.exp(aa*tt); dy=torch.autograd.grad(yy.sum(),tt,create_graph=False)[0]
    residual_max=float((dy-aa*yy).abs().max())
    q1=math.exp(-.7);
    vals=[]
    for j in [0.,1.,3.,6.]:
        q2=math.exp(-1.4+j); gap=abs(q2-q1)
        for lam in [0.,.1,1.,10.,100.]:
            obs=soft_defect(lam,j); pred=gap/(1+2*lam)
            vals.append({'jump':j,'lambda':lam,'observed':obs,'predicted':pred,'relerr':abs(obs-pred)/(abs(pred)+1e-8)})
    # Prediction C: abruptness should increase single-network post-event error relative to chained fit.
    jumps=[0.,1.,3.,6.]; pairs=[]
    for j in jumps:
        e_single=train_single(j); e_chain,_=train_exact(j)
        pairs.append({'jump':j,'single_post_rmse':e_single,'exact_chain_post_rmse':e_chain,'advantage':e_single-e_chain})
    return {'interface_max_defects':defects,'phase_residual_max':residual_max,'soft_scaling':vals,'event_severity':pairs}

def main():
    out=math_checks()
    Path('results.json').write_text(json.dumps(out,indent=2))
    print(json.dumps({'device':str(device),'results':out},indent=2))
if __name__=='__main__':
    try: main()
    except Exception as e:
        if device.type=='cuda':
            print('CUDA failed, rerun on CPU:',repr(e)); device=torch.device('cpu'); main()
        else: raise