import json, math, time, random
import numpy as np
import torch

SEED = 641
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')

# Small normalized 3-D Fourier grid.  P1 is a plane and P2 is a homogeneous cubic
# whose restriction to that plane gives a curved/straight algebraic locus.
N = 16
f = torch.fft.fftfreq(N, d=1.0/N)
k1 = f / (N/2)
Kz, Ky, Kx = torch.meshgrid(k1, k1, k1, indexing='ij')
K = torch.stack([Kx, Ky, Kz], dim=-1).to(device)

def poly_residual(K, a, b, cubic):
    a = a / (torch.linalg.vector_norm(a) + 1e-8)
    cubic = cubic / (torch.linalg.vector_norm(cubic) + 1e-8)
    x,y,z = K[...,0], K[...,1], K[...,2]
    # ten monomials of a homogeneous cubic
    mons = torch.stack([x**3, y**3, z**3, x*x*y, x*x*z, y*y*x, y*y*z,
                        z*z*x, z*z*y, x*y*z], -1)
    p1 = (K*a).sum(-1)
    p2 = (K*b).sum(-1) + (mons*cubic).sum(-1)
    return p1*p1 + p2*p2

def gate_from_coeff(K, a, b, cubic, tau):
    r = poly_residual(K, a, b, cubic)
    # Explicit symmetrization is the advertised inversion protection.
    rev = poly_residual(-K, a, b, cubic)
    r = .5*(r + rev)
    return torch.exp(-r / tau), r

# Ground-truth curved pinch curve: plane z=0 and cubic x^3-3xy^2=0.
a0 = torch.tensor([0.,0.,1.], device=device)
b0 = torch.zeros(3, device=device)
c0 = torch.zeros(10, device=device); c0[0] = 1.; c0[5] = -3.
tau0 = 0.045
with torch.no_grad():
    target_gate, target_r = gate_from_coeff(K,a0,b0,c0,tau0)

# Stage-1 math checks: inversion invariance, and concentration among low residual points.
with torch.no_grad():
    inv_err = (target_r - poly_residual(-K, a0, b0, c0)).abs().max().item()
    flat = target_r.flatten()
    q10 = torch.quantile(flat, .10)
    q90 = torch.quantile(flat, .90)
    gate_flat = target_gate.flatten()
    low_gate_mean = gate_flat[flat <= q10].mean().item()
    high_gate_mean = gate_flat[flat >= q90].mean().item()
    concentration_ratio = low_gate_mean / (high_gate_mean + 1e-12)
    active_fraction = (target_gate > .5).float().mean().item()

class DenseMultiplier(torch.nn.Module):
    def __init__(self, shape):
        super().__init__()
        self.re = torch.nn.Parameter(torch.zeros(shape))
        self.im = torch.nn.Parameter(torch.zeros(shape))
        torch.nn.init.normal_(self.re, std=.02)
    def forward(self, X):
        return X * torch.complex(self.re, self.im)

class PinchMultiplier(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.a = torch.nn.Parameter(torch.randn(3))
        self.b = torch.nn.Parameter(torch.randn(3)*.05)
        self.cubic = torch.nn.Parameter(torch.randn(10)*.05)
        self.log_tau = torch.nn.Parameter(torch.tensor(math.log(.06)))
    def forward(self, X):
        tau = torch.exp(self.log_tau).clamp(.005, .5)
        g, _ = gate_from_coeff(K, self.a, self.b, self.cubic, tau)
        return X * g

# Random fields make the Fourier regression identify every multiplier entry.
Btrain, Bval = 96, 48
xtrain = torch.randn(Btrain,N,N,N,device=device)
xval = torch.randn(Bval,N,N,N,device=device)
with torch.no_grad():
    ytrain = torch.fft.ifftn(torch.fft.fftn(xtrain)*target_gate).real
    yval = torch.fft.ifftn(torch.fft.fftn(xval)*target_gate).real

def fit(model, steps=350):
    opt = torch.optim.Adam(model.parameters(), lr=.035)
    model.train(); t0=time.perf_counter()
    for step in range(steps):
        # fixed full batch keeps the comparison deterministic and tiny
        pred = torch.fft.ifftn(model(torch.fft.fftn(xtrain))).real
        loss = torch.mean((pred-ytrain)**2)
        opt.zero_grad(); loss.backward(); opt.step()
    elapsed=time.perf_counter()-t0
    model.eval()
    with torch.no_grad():
        pred = torch.fft.ifftn(model(torch.fft.fftn(xval))).real
        mse = torch.mean((pred-yval)**2).item()
        rel = (torch.linalg.vector_norm(pred-yval)/torch.linalg.vector_norm(yval)).item()
    return mse, rel, elapsed

# CUDA OOM or driver problems are explicitly handled by rerunning on CPU.
def run():
    dense = DenseMultiplier((N,N,N)).to(device)
    pinch = PinchMultiplier().to(device)
    dmse,drel,dt = fit(dense)
    pmse, prel, pt = fit(pinch)
    return dmse,drel,dt,pmse,prel,pt
try:
    dmse,drel,dt,pmse,prel,pt = run()
except Exception as e:
    device=torch.device('cpu')
    Kz, Ky, Kx = torch.meshgrid(k1,k1,k1,indexing='ij'); K=torch.stack([Kx,Ky,Kz],-1)
    target_gate,target_r=gate_from_coeff(K,a0.cpu(),b0.cpu(),c0.cpu(),tau0)
    xtrain=torch.randn(Btrain,N,N,N); xval=torch.randn(Bval,N,N,N)
    ytrain=torch.fft.ifftn(torch.fft.fftn(xtrain)*target_gate).real
    yval=torch.fft.ifftn(torch.fft.fftn(xval)*target_gate).real
    dmse,drel,dt,pmse,prel,pt=run()

result = {
 'device': str(device), 'grid': [N,N,N],
 'math_check': {'max_inversion_residual_error': inv_err,
                'mean_gate_lowest_residual_decile': low_gate_mean,
                'mean_gate_highest_residual_decile': high_gate_mean,
                'low_to_high_gate_ratio': concentration_ratio,
                'fraction_gate_above_half': active_fraction},
 'benchmark': {'dense_mse': dmse, 'dense_relative_l2': drel, 'dense_seconds': dt,
               'pinch_mse': pmse, 'pinch_relative_l2': prel, 'pinch_seconds': pt},
 'parameters': {'dense_complex_entries': 2*N**3, 'pinch_trainable_scalars': 3+3+10+1},
 'note': 'Target is exactly an algebraic pinch gate applied to white-noise volumes; this tests representation and optimization on the claimed synthetic phenomenon, not a physical PDE.'
}
print(json.dumps(result, indent=2))