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

SEED = 2967
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)


def projectors(normals):
    I = torch.eye(3, dtype=normals.dtype, device=normals.device)
    return I[None] - normals[..., :, None] * normals[..., None, :]


def estimate_normals(x, k=12):
    # Small, explicit kNN covariance estimator; normal orientation is irrelevant to P.
    d = torch.cdist(x, x)
    nn_idx = d.topk(k + 1, largest=False).indices[:, 1:]
    neigh = x[nn_idx] - x[:, None, :]
    cov = torch.einsum('nki,nkj->nij', neigh, neigh) / float(k)
    vals, vecs = torch.linalg.eigh(cov)
    n = vecs[:, :, 0]
    return n / (n.norm(dim=-1, keepdim=True) + 1e-12), nn_idx


def rotation(seed):
    g = torch.Generator().manual_seed(seed)
    z = torch.randn(3, 3, generator=g)
    q, r = torch.linalg.qr(z)
    q = q @ torch.diag(torch.sign(torch.diag(r)))
    if torch.linalg.det(q) < 0: q[:, 0] *= -1
    return q


def make_cloud(seed, n=240, noise=0.8):
    g = torch.Generator().manual_seed(seed)
    x = torch.randn(n, 3, generator=g)
    x = x / x.norm(dim=1, keepdim=True)
    R = rotation(10000 + seed)
    x = x @ R.T
    # A smooth scalar field and its exact tangent gradient on the unit sphere.
    a = torch.tensor([0.8, -0.3, 0.5])
    f = x @ a
    normal = x
    clean_v = a[None] - f[:, None] * normal
    # Deliberate ambient normal corruption, which intrinsic projection should discard.
    eps = noise * torch.randn(n, 1, generator=g)
    v = clean_v + eps * normal
    target = clean_v.norm(dim=1)
    return x, f[:, None], v, target


class LocalLayer(nn.Module):
    def __init__(self, intrinsic=False, hidden=48):
        super().__init__(); self.intrinsic = intrinsic
        # [s_i,s_j,v_i,v_j,r,d] = 12 numbers
        self.mlp = nn.Sequential(nn.Linear(12, hidden), nn.SiLU(), nn.Linear(hidden, hidden), nn.SiLU(), nn.Linear(hidden, 1))

    def forward(self, x, s, v):
        n, idx = estimate_normals(x)
        P = projectors(n)
        if self.intrinsic: v = torch.einsum('nij,nj->ni', P, v)
        xj = x[idx]; sj = s[idx]; vj = v[idx]
        delta = xj - x[:, None, :]
        if self.intrinsic: r = torch.einsum('nij,nkj->nki', P, delta)
        else: r = delta
        d = r.norm(dim=-1, keepdim=True)
        vi = v[:, None, :].expand_as(vj)
        inp = torch.cat([s[:, None, :].expand_as(sj), sj, vi, vj, r, d], dim=-1)
        msg = self.mlp(inp).mean(dim=1)
        return msg[:, 0], n, P


def math_check():
    # Verify P is symmetric/idempotent, removes normal output, and PJP removes
    # both normal output and normal input derivative directions.
    torch.manual_seed(11)
    n = torch.randn(100, 3); n = n / n.norm(dim=1, keepdim=True); P = projectors(n)
    sym = (P - P.transpose(1, 2)).abs().max().item()
    idem = (P @ P - P).abs().max().item()
    normal_removed = (torch.einsum('ni,nij->nj', n, P)).abs().max().item()
    J = torch.randn(100, 3, 3)
    K = P @ J @ P
    left = torch.einsum('ni,nij->nj', n, K).abs().max().item()
    right = torch.einsum('nij,nj->ni', K, n).abs().max().item()
    return dict(projector_symmetry=sym, projector_idempotence=idem,
                normal_removed=normal_removed, jacobian_left_normal=left,
                jacobian_right_normal=right)


def train_eval(intrinsic, epochs=65):
    model = LocalLayer(intrinsic=intrinsic)
    opt = torch.optim.Adam(model.parameters(), lr=3e-3, weight_decay=1e-5)
    train = [make_cloud(i, noise=0.8) for i in range(8)]
    test = [make_cloud(100+i, noise=1.5) for i in range(4)]
    for ep in range(epochs):
        random.shuffle(train); model.train()
        for x,s,v,y in train:
            pred,_,_ = model(x,s,v)
            loss = ((pred-y)**2).mean()
            opt.zero_grad(); loss.backward(); opt.step()
    model.eval(); errs=[]
    with torch.no_grad():
        for x,s,v,y in test:
            p,n,P = model(x,s,v); errs.append(torch.mean((p-y)**2).sqrt().item() / (y.pow(2).mean().sqrt().item()+1e-8))
    return float(np.mean(errs)), model


def main():
    check = math_check()
    b, bm = train_eval(False); i, im = train_eval(True)
    # Mechanism observable: projection gives machine-zero normal velocity after each update.
    x,s,v,y = make_cloud(999, noise=2.0)
    with torch.no_grad():
        _, n, P = im(x,s,v); vp = torch.einsum('nij,nj->ni', P, v)
        leakage = (torch.abs((n*vp).sum(1))).max().item()
        raw_leakage = (torch.abs((n*v).sum(1))).mean().item()
    out = {'math_check': check, 'relative_rmse_baseline': b,
           'relative_rmse_intrinsic': i, 'relative_rmse_improvement_pct': 100*(b-i)/b,
           'raw_mean_normal_velocity': raw_leakage, 'projected_max_normal_velocity': leakage,
           'seed': SEED}
    print(json.dumps(out, indent=2))

if __name__ == '__main__': main()