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

SEED = 892
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_default_dtype(torch.float64)
try:
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
except Exception:
    device = torch.device('cpu')


def lap_scalar(z, xy):
    g = torch.autograd.grad(z.sum(), xy, create_graph=True, retain_graph=True)[0]
    vals = []
    for j in range(2):
        gj = torch.autograd.grad(g[:, j].sum(), xy, create_graph=True, retain_graph=True)[0][:, j]
        vals.append(gj)
    return vals[0] + vals[1]


def q_residuals(q, xy):
    qr, qi = q[:, 0], q[:, 1]
    gr = torch.autograd.grad(qr.sum(), xy, create_graph=True, retain_graph=True)[0]
    gi = torch.autograd.grad(qi.sum(), xy, create_graph=True, retain_graph=True)[0]
    lqr, lqi = lap_scalar(qr, xy), lap_scalar(qi, xy)
    er = gr[:, 0]**2 - gi[:, 0]**2 + gr[:, 1]**2 - gi[:, 1]**2
    ei = 2 * (gr[:, 0]*gi[:, 0] + gr[:, 1]*gi[:, 1])
    return lqr, lqi, er, ei


def q_pert(xy, eps):
    x, y = xy[:, 0], xy[:, 1]
    return torch.stack((x + eps*x*x, y), dim=1)


def toy_verification():
    a = torch.linspace(-1, 1, 31, device=device)
    X, Y = torch.meshgrid(a, a, indexing='ij')
    base = torch.stack((X.flatten(), Y.flatten()), 1)
    p = base.clone().requires_grad_(True)
    z = q_pert(p, 0.0)
    vals = q_residuals(z, p)
    exact_max = float(torch.max(torch.stack([v.abs().max() for v in vals])))
    epslist = np.array([1e-3, 2e-3, 5e-3, 1e-2, 2e-2, 5e-2, 1e-1])
    rows = []
    for e in epslist:
        p = base.clone().requires_grad_(True)
        vals = q_residuals(q_pert(p, float(e)), p)
        h = torch.sqrt(vals[0]**2 + vals[1]**2).mean().item()
        en = torch.sqrt(vals[2]**2 + vals[3]**2).mean().item()
        rows.append((float(e), h, en))
    # Predictions: harmonic residual is identically zero for this perturbation;
    # null-gradient residual is 4 eps*x + 4 eps^2*x^2, hence RMS scales linearly.
    earr = np.array([r[0] for r in rows]); nerr = np.array([r[2] for r in rows])
    slope = float(np.polyfit(np.log(earr), np.log(nerr), 1)[0])
    pred_coeff = 4.0 * math.sqrt(float((base[:,0]**2).mean()))
    observed_coeff = float(np.polyfit(earr[:4], nerr[:4], 1)[0])
    return {'exact_identity_max_abs': exact_max, 'sweep': [dict(eps=e,harmonic=h,null_gradient=n) for e,h,n in rows],
            'predicted_null_rms_linear_coefficient': pred_coeff, 'observed_small_eps_coefficient': observed_coeff,
            'observed_loglog_scaling_exponent': slope, 'predictions': {'harmonic stays zero': True, 'null residual O(eps)': True}}


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


class Ansatz(nn.Module):
    def __init__(self, width=20):
        super().__init__()
        self.q = nn.Sequential(nn.Linear(2,width), nn.Tanh(), nn.Linear(width,2))
        self.v = MLP(2,width)
        with torch.no_grad():
            self.q[-1].weight.mul_(0.02)
            self.q[-1].bias.copy_(torch.tensor([0.,0.]))
    def forward(self,xy):
        raw = self.q(xy)
        # affine canonical coordinate plus trainable residual
        q = torch.stack((xy[:,0]+raw[:,0], xy[:,1]+raw[:,1]),1)
        return self.v(q), q


def lap_u(u, xy): return lap_scalar(u, xy)

def train_model(kind, steps=500, n=96):
    model = MLP(2,24) if kind=='baseline' else Ansatz(20)
    model.to(device); opt=torch.optim.Adam(model.parameters(),lr=2e-3)
    for step in range(steps):
        xy=(torch.rand(n,2,device=device)*2-1).requires_grad_(True)
        if kind=='baseline': u=model(xy); q=None
        else: u,q=model(xy)
        # target harmonic solution u=x^2-y^2 with boundary supervision
        target=xy[:,0]**2-xy[:,1]**2
        pde=lap_u(u,xy)
        loss=(pde**2).mean()+10*((u-target)**2).mean()
        if q is not None:
            qr,qi,er,ei=q_residuals(q,xy)
            ramp=min(1.,step/50.)
            loss=loss+ramp*(qr.square().mean()+qi.square().mean()+er.square().mean()+ei.square().mean())
        opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad():
        test=(torch.rand(512,2,device=device)*2-1)
        target=test[:,0]**2-test[:,1]**2
        if kind=='baseline': pred=model(test); q=None
        else: pred,q=model(test)
        mse=((pred-target)**2).mean().item()
    test=test.requires_grad_(True)
    if kind=='baseline': pred=model(test)
    else: pred,q=model(test)
    residual=lap_u(pred,test).abs().mean().item()
    qres=None
    if q is not None:
        rr=q_residuals(q,test); qres=float(torch.sqrt(sum(x.square() for x in rr)).mean())
    return {'test_mse':mse,'mean_abs_laplacian':residual,'coordinate_constraint_rms':qres}


def main():
    verify=toy_verification()
    results={}
    for k in ['baseline','ansatz']:
        try: results[k]=train_model(k)
        except RuntimeError:
            global device
            device=torch.device('cpu'); results[k]=train_model(k)
    out={'device':str(device),'verification':verify,'comparison':results}
    print(json.dumps(out,indent=2))
    with open('results.json','w') as f: json.dump(out,f,indent=2)

if __name__=='__main__': main()