import json, math, time
import numpy as np
import torch
from torch import nn

SEED=1310
np.random.seed(SEED); torch.manual_seed(SEED)
DT=np.float64

def warp_numpy(q, n=20001, eps=1e-12, h=None):
    s=np.linspace(0,1,n)
    if h is None: h=np.zeros_like(s)
    rho=eps + np.maximum(s,1e-15)**(q-1)*np.exp(np.clip(h,-5,5))
    # trapezoidal cumulative integration
    R=np.concatenate([[0.], np.cumsum((rho[:-1]+rho[1:])*.5*(s[1]-s[0]))])
    r=R/R[-1]
    return s,r,rho/R[-1]

def log_slope(x,y,lo=1e-3,hi=.1):
    m=(x>lo)&(x<hi)&(y>0)
    return float(np.polyfit(np.log(x[m]),np.log(y[m]),1)[0])

def interp_max_error(q, lam, n):
    sg,rg,_=warp_numpy(q,n=50001)
    # target after mapping, sampled on computational grid
    y=rg**lam
    # evaluate exact on dense points and linear interpolation from n nodes
    sn=np.linspace(0,1,n); rn=np.interp(sn,sg,rg); yn=rn**lam
    xx=np.linspace(0,1,100001)
    yy=np.interp(xx,sn,yn)
    exact=np.interp(xx,sg,y)
    return float(np.max(np.abs(yy-exact)))

class MLP(nn.Module):
    def __init__(self):
        super().__init__(); self.net=nn.Sequential(nn.Linear(1,32),nn.Tanh(),nn.Linear(32,32),nn.Tanh(),nn.Linear(32,1))
    def forward(self,x): return self.net(x)

class LearnedWarp(nn.Module):
    def __init__(self,q=2,ng=257):
        super().__init__(); self.q=q; self.grid=torch.linspace(0,1,ng).view(-1,1)
        self.h=nn.Sequential(nn.Linear(1,12),nn.Tanh(),nn.Linear(12,1))
    def map(self,s):
        # differentiable cumulative trapezoids on fixed grid, then linear interpolation
        g=self.grid.to(s.device); hh=torch.clamp(self.h(g),-2.0,2.0)
        rho=1e-5 + torch.clamp(g,min=1e-6)**(self.q-1)*torch.exp(hh)
        d=g[1:]-g[:-1]; inc=.5*(rho[1:]+rho[:-1])*d
        R=torch.cat([torch.zeros(1,1,device=s.device),torch.cumsum(inc,0)],0); R=R/R[-1]
        idx=torch.clamp((s*(len(g)-1)).long(),0,len(g)-2); t=s*(len(g)-1)-idx
        return R[idx,0]*(1-t)+R[idx+1,0]*t

def train_case(kind, steps=500):
    torch.manual_seed(SEED+({'identity':0,'fixed':1,'learned':2}[kind]))
    model=MLP(); params=list(model.parameters())
    warp=LearnedWarp() if kind=='learned' else None
    if warp: params += list(warp.parameters())
    opt=torch.optim.Adam(params,lr=2e-3)
    # fixed samples, held-out deterministic grid
    s=torch.linspace(.0005,1,96).view(-1,1); target_phys=s**.5
    sh=torch.linspace(.002,1,400).view(-1,1); yh=sh**.5
    t0=time.time()
    for k in range(steps):
        if kind=='identity': x=s; y=target_phys
        elif kind=='fixed': x=s; y=s # r=s^2, physical sqrt(r)=s
        else:
            r=warp.map(s); x=s; y=torch.sqrt(r)
        pred=model(x); loss=((pred-y)**2).mean()
        # weak regularity control on learned h sampled grid
        if warp: loss=loss+1e-5*(warp.h(warp.grid[1:])-warp.h(warp.grid[:-1])).pow(2).mean()
        opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad():
        if kind=='identity': yt=yh
        elif kind=='fixed': yt=sh
        else: yt=torch.sqrt(warp.map(sh))
        err=((model(sh)-yt)**2).mean().sqrt().item()
    return err, time.time()-t0

def main():
    # Prediction 1: exact endpoint normalization, strict positivity, and derivative identity.
    endpoint=[]
    for q in [1,1.5,2,3,4]:
        s,r,rp=warp_numpy(q)
        fd=np.gradient(r,s); rel=float(np.max(np.abs(fd[2:-2]-rp[2:-2])/(np.abs(rp[2:-2])+1e-12)))
        endpoint.append({'q':q,'r0':float(r[0]),'r1':float(r[-1]),'min_rprime':float(rp[1:].min()),'fitted_r_exponent':log_slope(s,r),'max_relative_derivative_error':rel})
    # Prediction 2: exponent multiplication over q,lambda sweep.
    exponent=[]
    for q in [1,1.5,2,3,4]:
        for lam in [.25,.5,.75]:
            s,r,_=warp_numpy(q); exponent.append({'q':q,'lambda':lam,'predicted':q*lam,'observed':log_slope(s,r**lam)})
    # Prediction 3: higher transformed regularity gives steeper interpolation convergence.
    conv={}
    for q in [1,2,3]:
        lam=.25; ns=[17,33,65,129,257]
        es=[interp_max_error(q,lam,n) for n in ns]
        slope=float(np.polyfit(np.log(ns),np.log(es),1)[0])
        conv[str(q)]={'q_lambda':q*lam,'ns':ns,'errors':es,'observed_error_slope_vs_N':slope,'predicted_slope_ideal':-min(2,q*lam)}
    mlp={}
    for k in ['identity','fixed','learned']:
        e,t=train_case(k); mlp[k]={'heldout_rmse':e,'seconds':t}
    out={'seed':SEED,'predictions':{'monotonic_normalized':endpoint,'exponent_multiplication':exponent,'interpolation_convergence':conv},'mini_experiment':mlp}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()