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

SEED = 1144
np.random.seed(SEED); torch.manual_seed(SEED)
K = 1.0
EPS = 1e-6

def reconstruct(x, K=1.0, eps=1e-6):
    # x is the proper-velocity spatial coordinate
    q = (x*x).sum(dim=-1)
    t = torch.sqrt(torch.clamp(q - 1.0/K, min=eps))
    return torch.cat([t.unsqueeze(-1), x], dim=-1)

def lorentz_residual(z, K=1.0):
    # Lorentz convention: -t^2 + ||s||^2 = 1/K
    return (z[..., 1:]**2).sum(-1) - z[..., 0]**2 - 1.0/K

def math_checks():
    out = {}
    # Prediction 1: for norms above threshold, reconstruction residual is roundoff.
    d = 7
    radii = np.array([1.001, 1.01, 1.1, 2., 5.])
    x = torch.zeros((len(radii), d), dtype=torch.float64)
    x[:, 0] = torch.tensor(radii, dtype=torch.float64)
    z = reconstruct(x.double(), K, 1e-30)
    residual = lorentz_residual(z.double(), K).abs().numpy()
    out['reconstruction_max_abs_residual'] = float(residual.max())
    out['reconstruction_residuals'] = residual.tolist()
    # Prediction 2: validity transition is r*=1/sqrt(K), estimated by bisection.
    threshold = 1/math.sqrt(K)
    lo, hi = 0.0, 2.0*threshold
    for _ in range(70):
        mid=(lo+hi)/2
        if mid*mid >= 1/K: hi=mid
        else: lo=mid
    out['predicted_boundary'] = threshold
    out['observed_boundary_bisection'] = hi
    # Prediction 3: dt/dr=r/sqrt(r^2-1/K), diverges as delta=r-r* shrinks.
    deltas = np.array([1e-1, 1e-2, 1e-3, 1e-4, 1e-5])
    r = threshold + deltas
    analytic = r/np.sqrt(r*r-1/K)
    # central finite differences, with steps small relative to each delta
    numeric=[]
    for rr, dd in zip(r, deltas):
        h=dd*1e-4
        f=lambda v: math.sqrt(max(v*v-1/K, 1e-30))
        numeric.append((f(rr+h)-f(rr-h))/(2*h))
    out['derivative_deltas'] = deltas.tolist()
    out['derivative_predicted'] = analytic.tolist()
    out['derivative_observed'] = numeric
    out['derivative_growth_ratio_predicted'] = float(analytic[-1]/analytic[0])
    out['derivative_growth_ratio_observed'] = float(numeric[-1]/numeric[0])
    # Sweep violation for the literal epsilon-clamped formula: below boundary is not exact.
    rs=np.linspace(.2,1.8,17)
    xx=torch.zeros((len(rs),3),dtype=torch.float64); xx[:,0]=torch.tensor(rs)
    rr=lorentz_residual(reconstruct(xx.double(),K,EPS),K).numpy()
    out['violation_sweep']=[{'r':float(a),'abs_residual':float(abs(b))} for a,b in zip(rs,rr)]
    return out

class PVNet(nn.Module):
    def __init__(self, depth=4, width=32, dim=8):
        super().__init__(); self.inp=nn.Linear(2,dim); self.layers=nn.ModuleList([nn.Linear(dim,dim) for _ in range(depth-1)]); self.out=nn.Linear(dim,2)
    def forward(self,x):
        h=torch.tanh(self.inp(x))
        for l in self.layers: h=torch.tanh(l(h))
        # PV spatial tensor is h; manifold boundary is only reconstructed at interface.
        z=reconstruct(h)
        return self.out(z[:,1:]), z

class ProjectedLorentzNet(nn.Module):
    def __init__(self, depth=4, width=32, dim=8):
        super().__init__(); self.inp=nn.Linear(2,dim); self.layers=nn.ModuleList([nn.Linear(dim,dim) for _ in range(depth-1)]); self.out=nn.Linear(dim,2)
    def project(self,h):
        # Standard radial projection makes every intermediate spatial vector valid.
        n=torch.linalg.vector_norm(h,dim=-1,keepdim=True)
        return h * torch.clamp(1.0/(n+1e-12), min=1.001)
    def forward(self,x):
        h=self.project(torch.tanh(self.inp(x)))
        for l in self.layers: h=self.project(torch.tanh(l(h)))
        z=reconstruct(h)
        return self.out(z[:,1:]), z

def mini_experiment():
    torch.manual_seed(SEED)
    n=1024
    x=torch.randn(n,2); y=((x[:,0]*x[:,1]>0).long())
    tr,va=torch.arange(0,768),torch.arange(768,n)
    result={}
    for depth in (4,12):
        for kind, cls in [('baseline',ProjectedLorentzNet),('pv',PVNet)]:
            torch.manual_seed(SEED+depth+(0 if kind=='baseline' else 100))
            model=cls(depth=depth); opt=torch.optim.Adam(model.parameters(),lr=2e-3); losses=[]; grad=[]; invalid=0
            t0=time.perf_counter()
            for step in range(250):
                opt.zero_grad(); logits,z=model(x[tr]); loss=nn.functional.cross_entropy(logits,y[tr]); loss.backward()
                grad.append(float(torch.nn.utils.clip_grad_norm_(model.parameters(),1e9))); opt.step(); losses.append(float(loss))
                invalid += int((lorentz_residual(z).abs()>1e-5).sum())
            elapsed=time.perf_counter()-t0
            with torch.no_grad():
                logits,z=model(x[va]); acc=float((logits.argmax(1)==y[va]).float().mean()); v=float(lorentz_residual(z).abs().max())
            result[f'{kind}_depth{depth}']={'final_train_loss':losses[-1],'val_accuracy':acc,'max_final_residual':v,'invalid_count':invalid,'gradient_std':float(np.std(grad)),'seconds_250_steps':elapsed}
    return result

if __name__=='__main__':
    torch.set_num_threads(min(4,torch.get_num_threads()))
    report={'math_checks':math_checks(),'mini_experiment':mini_experiment()}
    with open('results.json','w') as f: json.dump(report,f,indent=2)
    print(json.dumps(report,indent=2))