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

SEED = 1286
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_default_dtype(torch.float64)
device = "cuda" if torch.cuda.is_available() else "cpu"


def basis(d, p):
    from itertools import combinations
    return list(combinations(range(d), p))


def wedge_matrix(w, d, p):
    B, T = basis(d, p), basis(d, p + 1)
    K = np.zeros((len(T), len(B)))
    for j, a in enumerate(B):
        for r in range(d):
            if r in a:
                continue
            seq = (r,) + a
            inversions = sum(seq[u] > seq[v] for u in range(len(seq)) for v in range(u + 1, len(seq)))
            K[T.index(tuple(sorted(seq))), j] = ((-1) ** inversions) * w[r]
    return K


def relu_power(s, m):
    return torch.relu(s) ** m / math.factorial(m)


class ExactBackbone(nn.Module):
    def __init__(self, w, b, k=2):
        super().__init__()
        self.register_buffer("w", torch.as_tensor(w))
        self.register_buffer("b", torch.as_tensor(b))
        self.k = k

    def features(self, x, m):
        return relu_power(x @ self.w.T + self.b, m)

    def scalar(self, x, coef):
        return self.features(x, self.k) @ coef

    def gradient_form(self, x, coef):
        # d(sum_i coef_i sigma_k(s_i)) = sum_i coef_i sigma_{k-1}(s_i) ds_i
        return self.features(x, self.k - 1) @ (coef[:, None] * self.w)


def math_checks():
    rng = np.random.default_rng(SEED)
    d, n = 3, 7
    w = rng.normal(size=(n, d)); w /= np.linalg.norm(w, axis=1, keepdims=True)
    max_nil = 0.0
    for wi in w:
        for p in range(d - 1):
            max_nil = max(max_nil, np.max(np.abs(wedge_matrix(wi, d, p + 1) @ wedge_matrix(wi, d, p))))
    # Finite-difference verification of d sigma_m(s)=sigma_{m-1}(s) ds away from kink.
    x = torch.tensor([[0.37, -0.21, 0.19]], requires_grad=True)
    wt = torch.tensor(w[:1]); bt = torch.tensor([0.41]); s = x @ wt.T + bt
    m = 3
    y = relu_power(s, m).sum(); grad = torch.autograd.grad(y, x)[0].detach().numpy()[0]
    rhs = (relu_power(s.detach(), m - 1).numpy()[0, 0] * w[0])
    identity_err = float(np.max(np.abs(grad - rhs)))
    # Predicted operator scaling: ||K_0|| = ||w||, and composition remains zero.
    scales = np.array([0.25, 0.5, 1., 2., 4.])
    norms = []
    nil_scaled = []
    for a in scales:
        wi = a * w[0]
        k0, k1 = wedge_matrix(wi, d, 0), wedge_matrix(wi, d, 1)
        norms.append(np.linalg.norm(k0, 2)); nil_scaled.append(np.max(np.abs(k1 @ k0)))
    slope = float(np.polyfit(np.log(scales), np.log(norms), 1)[0])
    return {"nilpotence_max_abs": float(max_nil), "derivative_identity_max_abs": identity_err,
            "scale_sweep": [{"scale": float(a), "K0_norm": float(v), "K1K0_max": float(z)} for a,v,z in zip(scales,norms,nil_scaled)],
            "loglog_scaling_exponent_observed": slope, "predicted_scaling_exponent": 1.0}


def make_data(n=512):
    x = torch.rand(n, 2, device=device)
    # exact potential; target vector field is curl-free
    phi = torch.sin(math.pi*x[:,0]) * torch.sin(math.pi*x[:,1])
    target = torch.stack((math.pi*torch.cos(math.pi*x[:,0])*torch.sin(math.pi*x[:,1]),
                          math.pi*torch.sin(math.pi*x[:,0])*torch.cos(math.pi*x[:,1])), 1)
    return x, target


def train_compare(steps=700):
    rng = np.random.default_rng(SEED)
    nfeat = 28
    w = rng.normal(size=(nfeat,2)); w /= np.linalg.norm(w,axis=1,keepdims=True)
    b = rng.uniform(-0.7,0.7,size=nfeat)
    x, target = make_data()
    exact = ExactBackbone(w,b,k=2).to(device)
    coef = nn.Parameter(torch.zeros(nfeat,device=device)); opt = torch.optim.Adam([coef],lr=0.04)
    mlp = nn.Sequential(nn.Linear(2,24),nn.Tanh(),nn.Linear(24,24),nn.Tanh(),nn.Linear(24,2)).to(device)
    opt2 = torch.optim.Adam(mlp.parameters(),lr=0.01)
    for _ in range(steps):
        opt.zero_grad(); pred=exact.gradient_form(x,coef); loss=((pred-target)**2).mean(); loss.backward(); opt.step()
        opt2.zero_grad(); pred2=mlp(x); loss2=((pred2-target)**2).mean(); loss2.backward(); opt2.step()
    with torch.no_grad():
        p=exact.gradient_form(x,coef); q=mlp(x)
        exact_mse=float(((p-target)**2).mean()); mlp_mse=float(((q-target)**2).mean())
    # Autodiff divergence/curl of vector outputs on a smaller probe batch.
    z=x[:64].detach().requires_grad_(True)
    def violations(model_kind):
        if model_kind=="exact": out=exact.gradient_form(z,coef)
        else: out=mlp(z)
        grads=[]
        for j in range(2): grads.append(torch.autograd.grad(out[:,j].sum(),z,retain_graph=True)[0])
        g0,g1=grads
        curl=g1[:,0]-g0[:,1]
        return float(torch.sqrt(torch.mean(curl**2)))
    return {"exact_params": int(nfeat), "mlp_params": sum(p.numel() for p in mlp.parameters()),
            "exact_mse": exact_mse, "mlp_mse": mlp_mse,
            "exact_curl_rms": violations("exact"), "mlp_curl_rms": violations("mlp")}


if __name__ == "__main__":
    try:
        result={"device":device,"math":math_checks(),"comparison":train_compare()}
    except Exception as e:
        if device == "cuda":
            device="cpu"; torch.cuda.empty_cache(); result={"device":device,"math":math_checks(),"comparison":train_compare()}
        else:
            raise
    with open("results.json","w") as f: json.dump(result,f,indent=2)
    print(json.dumps(result,indent=2))