import random
import numpy as np
import torch
from torch import nn
from scipy.spatial import Delaunay

SEED = 113
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)

def make_mesh(nx=2, ny=2, nz=1):
    pts = np.array([(i/nx, j/ny, k/max(1,nz))
                    for k in range(nz+1) for j in range(ny+1) for i in range(nx+1)], float)
    tet = Delaunay(pts).simplices.copy()
    for q, t in enumerate(tet):
        a,b,c,d = pts[t]
        if np.linalg.det(np.stack([b-a,c-a,d-a])) < 0:
            tet[q,[0,1]] = tet[q,[1,0]]
    return pts, tet

def build_incidence(tet, nvert):
    edges = sorted({tuple(sorted((int(t[i]),int(t[j]))))
                    for t in tet for i in range(4) for j in range(i+1,4)})
    faces = sorted({tuple(sorted((int(t[i]),int(t[j]),int(t[k]))))
                    for t in tet for i in range(4) for j in range(i+1,4)
                    for k in range(j+1,4)})
    ei = {e:i for i,e in enumerate(edges)}
    fi = {f:i for i,f in enumerate(faces)}
    D0 = np.zeros((len(edges), nvert))
    for r,(a,b) in enumerate(edges): D0[r,a],D0[r,b] = -1,1
    D1 = np.zeros((len(faces),len(edges)))
    for r,f in enumerate(faces):
        for omit in range(3):
            sub = tuple(f[j] for j in range(3) if j != omit)
            e = tuple(sorted(sub)); sign = (-1)**omit
            if sub != e: sign *= -1
            D1[r,ei[e]] += sign
    D2 = np.zeros((len(tet),len(faces)))
    for r,t in enumerate(tet):
        # positively oriented tetrahedron boundary: omit vertex with alternating sign.
        for omit in range(4):
            sub = tuple(int(t[j]) for j in range(4) if j != omit)
            f = tuple(sorted(sub)); sign = (-1)**omit
            inv = sum(sub[i] > sub[j] for i in range(3) for j in range(i+1,3))
            if inv % 2: sign *= -1
            D2[r,fi[f]] += sign
    return D0,D1,D2,edges,faces

class ExactComplex(nn.Module):
    def __init__(self, D0, D1, D2):
        super().__init__()
        self.register_buffer('D0', torch.tensor(D0,dtype=torch.float32))
        self.register_buffer('D1', torch.tensor(D1,dtype=torch.float32))
        self.register_buffer('D2', torch.tensor(D2,dtype=torch.float32))
        self.net = nn.Sequential(nn.Linear(1,16),nn.Tanh(),nn.Linear(16,1))
    def forward(self, x):
        # Exact compatible signal is annihilated by D1 after D0; predict a scalar
        # response from a defect magnitude, with no learned cross-order operator.
        defect = self.D1 @ x
        return self.net((defect.pow(2).mean().sqrt()).reshape(1,1)).reshape(1)

class Unconstrained(nn.Module):
    def __init__(self, ne, nf):
        super().__init__()
        self.map = nn.Parameter(torch.randn(nf,ne)*0.15)
        self.net = nn.Sequential(nn.Linear(1,16),nn.Tanh(),nn.Linear(16,1))
    def forward(self,x):
        g = self.map @ x
        return self.net((g.pow(2).mean().sqrt()).reshape(1,1)).reshape(1)

def run():
    pts,tet = make_mesh(); D0,D1,D2,edges,faces = build_incidence(tet,len(pts))
    r10=np.linalg.norm(D1@D0)/(np.linalg.norm(D0)+1e-12)
    r21=np.linalg.norm(D2@D1)/(np.linalg.norm(D1)+1e-12)
    rng=np.random.default_rng(SEED)
    # Direct compatible-field test: D0 u must be killed by D1. A same-size
    # unconstrained cross-order map has no such algebraic guarantee.
    compatible_leaks=[]; unconstrained_leaks=[]
    A=rng.normal(size=(D1.shape[0],D0.shape[0]))
    for _ in range(100):
        u=rng.normal(size=D0.shape[1]); e=D0@u
        compatible_leaks.append(np.linalg.norm(D1@e)/(np.linalg.norm(e)+1e-12))
        unconstrained_leaks.append(np.linalg.norm(A@e)/(np.linalg.norm(e)+1e-12))
    exact_leak=float(np.mean(compatible_leaks)); uncon_leak=float(np.mean(unconstrained_leaks))
    # Compatible vertex displacement fields versus random edge perturbations.
    xs=[]; ys=[]
    for _ in range(160):
        u=rng.normal(size=len(pts)); compatible=D0@u
        incompatible=compatible + (0.35 if _%2 else 0.0)*rng.normal(size=len(edges))
        x=torch.tensor(incompatible,dtype=torch.float32)
        xs.append(x); ys.append(float(_%2))
    def train(model):
        opt=torch.optim.Adam(model.parameters(),lr=0.02)
        for _ in range(250):
            loss=0.
            for x,y in zip(xs,ys):
                pred=model(x); loss=loss+(pred-y)**2
            opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            pred=np.array([float(model(x)) for x in xs])
        return float(np.mean((pred-np.array(ys))**2)), float(np.mean((pred[:80]<.5)==(np.array(ys[:80])<.5)))
    exact=train(ExactComplex(D0,D1,D2)); uncon=train(Unconstrained(len(edges),len(faces)))
    print({'vertices':len(pts),'edges':len(edges),'faces':len(faces),'cells':len(tet),
           'relative_D1D0':r10,'relative_D2D1':r21,
           'compatible_D1D0_leak':exact_leak,
           'unconstrained_leak':uncon_leak,
           'exact_mse_accuracy':exact,'unconstrained_mse_accuracy':uncon})

if __name__ == '__main__': run()