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

SEED = 17
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
try:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
except Exception:
    device = torch.device("cpu")


def harmonic_cycle(n=24):
    # D is the vertex-edge incidence matrix; ker(D) is the 1-cycle space.
    D = np.zeros((n, n), dtype=float)
    for e in range(n):
        D[e, e] = -1.0
        D[(e + 1) % n, e] = 1.0
    # Positive diagonal mass matrices, deliberately nonuniform.
    m = 0.5 + np.linspace(0.2, 1.4, n)
    M = np.diag(m)
    # There are no 2-simplices, hence the closed constraint is vacuous.
    # Coclosedness is D^T M h = 0. Its one-dimensional solution is
    # h proportional to M^{-1} times the constant circulation vector.
    h = np.linalg.solve(M, np.ones(n))
    h /= math.sqrt(h @ M @ h)
    H = h[:, None]
    return D, M, H


def math_check():
    D, M, H = harmonic_cycle()
    # This cycle has no 2-simplices, so B_2 is the empty (0 x n) matrix.
    # Thus B_2 H = 0 exactly; D is B_1 and tests the codifferential.
    B2 = np.zeros((0, len(H)))
    closed = np.linalg.norm(B2 @ H)
    coclosed = np.linalg.norm(D.T @ M @ H)
    ortho = float((H.T @ M @ H).item())
    rng = np.random.default_rng(SEED)
    x = rng.normal(size=(len(H), 50))
    c = H.T @ M @ x
    u = x - H @ c
    projection_error = np.max(np.abs(H.T @ M @ u))
    reconstruction = np.max(np.abs(x - (u + H @ c)))
    # Period is chosen as the linear circulation functional whose value on H is 1.
    P = (H.T @ M)
    period_consistency = np.max(np.abs(P @ H - 1.0))
    return {
        "closed_residual": float(closed), "coclosed_residual": float(coclosed),
        "weighted_orthonormality": ortho, "max_gauge_residual": float(projection_error),
        "max_reconstruction_residual": float(reconstruction),
        "period_basis_residual": float(period_consistency)
    }, D, M, H


class LocalGNN(nn.Module):
    def __init__(self, n, hidden=32, harmonic=False, H=None, M=None):
        super().__init__(); self.n=n; self.harmonic=harmonic
        self.inp=nn.Linear(1, hidden)
        self.layers=nn.ModuleList([nn.Linear(hidden, hidden) for _ in range(4)])
        self.local=nn.Linear(hidden*2, 1)
        if harmonic:
            self.coeff=nn.Sequential(nn.Linear(hidden, hidden), nn.Tanh(), nn.Linear(hidden, 1))
            self.register_buffer("H", torch.tensor(H, dtype=torch.float32))
            self.register_buffer("M", torch.tensor(M, dtype=torch.float32))
    def forward(self, x):
        # x: batch,n node scalars. Ring message passing, no global operation.
        z=torch.tanh(self.inp(x.unsqueeze(-1)))
        for layer in self.layers:
            msg=(torch.roll(z,1,1)+torch.roll(z,-1,1))/2
            z=torch.tanh(layer(z+msg))
        edge_z=(z+torch.roll(z,-1,1))/2
        raw=self.local(torch.cat([z, edge_z], dim=-1)).squeeze(-1)
        if not self.harmonic:
            return raw
        # Global latent channel: pooled embedding predicts circulation coefficient.
        a=self.coeff(z.mean(dim=1)).squeeze(-1)
        # Gauge-fix local output using H^T M, then reconstruct y=u+Ha.
        c=raw @ (torch.diagonal(self.M) * self.H.squeeze(1))
        u=raw-c[:,None]*self.H.squeeze(1)
        return u+a[:,None]*self.H.squeeze(1)


def run_training(H, M, n=24, steps=700):
    # A source scalar is placed at one node. The target has global circulation
    # a on every edge; targets are randomized independently.
    rng=np.random.default_rng(SEED)
    Ntr,Nte=256,256
    xtr=np.zeros((Ntr,n),np.float32); xte=np.zeros((Nte,n),np.float32)
    atr=rng.uniform(-2,2,Ntr).astype(np.float32); ate=rng.uniform(-2,2,Nte).astype(np.float32)
    xtr[:,0]=atr; xte[:,0]=ate
    ytr=(atr[:,None]*H[:,0][None,:]).astype(np.float32); yte=(ate[:,None]*H[:,0][None,:]).astype(np.float32)
    xtr=torch.tensor(xtr,device=device); xte=torch.tensor(xte,device=device)
    ytr=torch.tensor(ytr,device=device); yte=torch.tensor(yte,device=device)
    out={}
    for name,is_h in [("baseline",False),("harmonic",True)]:
        torch.manual_seed(SEED)
        model=LocalGNN(n,hidden=24,harmonic=is_h,H=H,M=M).to(device)
        opt=torch.optim.Adam(model.parameters(),lr=3e-3)
        for step in range(steps):
            # minibatches keep memory modest
            ix=torch.randint(0,Ntr,(64,),device=device)
            pred=model(xtr[ix]); loss=((pred-ytr[ix])**2).mean()
            opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            pred=model(xte)
            field_rmse=float(torch.sqrt(((pred-yte)**2).mean()).cpu())
            P=torch.tensor((H.T@M).astype(np.float32),dtype=torch.float32,device=device)
            periods=pred@P.T; target_period=yte@P.T
            period_rmse=float(torch.sqrt(((periods-target_period)**2).mean()).cpu())
            # Error separately on edges more than four hops from source.
            far=list(range(5,n-4))
            far_rmse=float(torch.sqrt(((pred[:,far]-yte[:,far])**2).mean()).cpu())
        out[name]={"field_rmse":field_rmse,"period_rmse":period_rmse,"far_edge_rmse":far_rmse,
                   "parameters":sum(p.numel() for p in model.parameters())}
    return out

if __name__ == "__main__":
    checks,D,M,H=math_check()
    results=run_training(H,M)
    report={"device":str(device),"math_check":checks,"results":results,
            "setup":"24-edge weighted cycle with no 2-simplices; source scalar at one node; randomized global harmonic circulation target; 4 message-passing layers; 700 Adam steps; harmonic head has 3122 vs baseline 2497 parameters"}
    print(json.dumps(report,indent=2))
    with open("results.json","w") as f: json.dump(report,f,indent=2)