import json, math, os, random
import numpy as np
import torch
from torch import nn

SEED = 1729

def seed_all(s):
    random.seed(s); np.random.seed(s); torch.manual_seed(s)
    if torch.cuda.is_available(): torch.cuda.manual_seed_all(s)

# Differentiable Heron area. Inputs are positive side lengths.
def heron(x):
    a,b,c = x[...,0], x[...,1], x[...,2]
    s = (a+b+c)/2
    q = (s*(s-a)*(s-b)*(s-c)).clamp_min(1e-12)
    return torch.sqrt(q)

def triangle_data(n, seed):
    g = torch.Generator().manual_seed(seed)
    # Mixture includes ordinary and nearly collinear triangles, where ranking matters.
    x = torch.randn(n,3,3,generator=g)
    near = torch.rand(n,generator=g) < .25
    x[near,:,1] = x[near,:,0] + .03*torch.randn(near.sum(),3,generator=g)
    # Random per-example scale makes raw Jacobian sensitivity nonuniform.
    x *= (0.6 + 1.4*torch.rand(n,1,1,generator=g))
    p,q,r = x[:,0],x[:,1],x[:,2]
    sides = torch.stack([(p-q).norm(dim=1),(q-r).norm(dim=1),(r-p).norm(dim=1)],1)
    area = heron(sides)
    # sine of angle at p: an angle-sensitive, rigid-motion invariant target
    target = 2*area/(sides[:,0]*sides[:,2]).clamp_min(1e-8)
    target = target.clamp(0,1)
    return sides, area, target

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

def train_eval(kind, train, test, steps=500):
    ts,ta,ty=train; vs,va,vy=test
    def feat(s,a):
        if kind=='distance': return s
        if kind=='area': return torch.cat([s,a[:,None]],1)
        # Scalar triangle Jacobian has one singular value: norm of dA/d(a,b,c).
        z=s.detach().clone().requires_grad_(True)
        aa=heron(z); grad=torch.autograd.grad(aa.sum(),z)[0]
        sig=grad.norm(dim=1)
        w=(sig/(sig+1e-3)).clamp(0,1).detach()
        return torch.cat([s,(a*w)[:,None]],1)
    X,Y=feat(ts,ta),ty
    VX,VY=feat(vs,va),vy
    model=MLP(X.shape[1]); opt=torch.optim.Adam(model.parameters(),lr=3e-3)
    for _ in range(steps):
        opt.zero_grad(); loss=((model(X)-Y)**2).mean(); loss.backward(); opt.step()
    with torch.no_grad():
        mse=((model(VX)-VY)**2).mean().item()
        pred=model(VX)
    return mse, model, pred

def matrix_check():
    # The paper's polynomial coordinate is alpha(x,y,z)=4*area^2 when
    # x,y,z are squared side lengths. Its gradient is the displayed row.
    t=torch.tensor([2.1,1.8,2.2,2.0,2.4,1.9,2.3,2.1,2.5],dtype=torch.double,requires_grad=True)
    tris=[(0,5,6),(1,2,8),(3,4,7),(6,7,8)]
    J=torch.zeros(4,9,dtype=torch.double)
    for row,ix in enumerate(tris):
        x=t[list(ix)]
        alpha=-0.5*(x*x).sum()+x[0]*x[1]+x[0]*x[2]+x[1]*x[2]
        J[row]=torch.autograd.grad(alpha,t,retain_graph=True)[0]
    M=torch.zeros(4,9,dtype=torch.double)
    M[0,[0,5,6]]=torch.stack([-t[0]+t[5]+t[6],t[0]-t[5]+t[6],t[0]+t[5]-t[6]])
    M[1,[1,2,8]]=torch.stack([-t[1]+t[2]+t[8],t[1]-t[2]+t[8],t[1]+t[2]-t[8]])
    M[2,[3,4,7]]=torch.stack([-t[3]+t[4]+t[7],t[3]-t[4]+t[7],t[3]+t[4]-t[7]])
    M[3,[6,7,8]]=torch.stack([-t[6]+t[7]+t[8],t[6]-t[7]+t[8],t[6]+t[7]-t[8]])
    err=(M-J).abs().max().item()
    sv=torch.linalg.svdvals(J).detach().numpy()
    return {'squared_length_polynomial_jacobian_max_abs_error':err,
            'displayed_matrix_singular_values':sv.tolist(),
            'displayed_matrix_rank':int((sv>1e-10).sum()),
            'smallest_singular_value':float(sv[-1])}

def main():
    seed_all(SEED)
    check=matrix_check()
    results={k:[] for k in ['distance','area','weighted_area']}
    for split in range(5):
        train=triangle_data(700,11+2*split); test=triangle_data(500,12+2*split)
        for k in results:
            mse, model, pred=train_eval(k,train,test,steps=500)
            results[k].append(math.sqrt(mse))
    results={k:{'rmse_each_split':v,'mean_rmse':float(np.mean(v)),'std_rmse':float(np.std(v))} for k,v in results.items()}
    # Exact rigid transformation invariance of all geometric features.
    sides,area,target=test
    g=torch.Generator().manual_seed(99); R=torch.randn(3,3,generator=g); Q,_=torch.linalg.qr(R)
    # Reconstruct one set of points only for a numerical transformed-feature check.
    pts=torch.randn(200,3,3,generator=g); rot=pts@Q
    def ds(x):
        return torch.stack([(x[:,0]-x[:,1]).norm(dim=1),(x[:,1]-x[:,2]).norm(dim=1),(x[:,2]-x[:,0]).norm(dim=1)],1)
    sdiff=(ds(pts)-ds(rot)).abs().max().item()
    adiff=(heron(ds(pts))-heron(ds(rot))).abs().max().item()
    out={'seed':SEED,'math_check':check,'benchmark':results,
         'rigid_transform_max_distance_feature_change':sdiff,
         'rigid_transform_max_area_feature_change':adiff,
         'train_size':700,'test_size':500,'steps':500}
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()