import json, math, random
import numpy as np
import torch
import torch.nn.functional as F

SEED = 314
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
try:
    if DEVICE == "cuda":
        torch.cuda.empty_cache()
except Exception:
    DEVICE = "cpu"

N = 32
Y, X = torch.meshgrid(torch.arange(N, device=DEVICE), torch.arange(N, device=DEVICE), indexing="ij")

def phase_field(centers):
    # A periodic-grid engineering toy: use two well-separated phase singularities.
    # The opposite signs make the far field approximately constant.
    ph = torch.zeros((N, N), device=DEVICE)
    for cy, cx, charge in centers:
        dy = Y.float() - cy
        dx = X.float() - cx
        ph = ph + charge * torch.atan2(dy, dx)
    return ph

def unit_from_phase(ph):
    return torch.stack((torch.cos(ph), torch.sin(ph)), dim=0)

def wrapped(d):
    return torch.atan2(torch.sin(d), torch.cos(d))

def winding(u):
    ph = torch.atan2(u[1], u[0])
    # plaquette circulation with forward edges, periodic boundaries
    dx = wrapped(torch.roll(ph, -1, 1) - ph)
    dy = wrapped(torch.roll(ph, -1, 0) - ph)
    circulation = dx + torch.roll(dy, -1, 1) - torch.roll(dx, -1, 0) - dy
    return circulation / (2.0 * math.pi)

def frac_energy(u, s=0.5):
    # torch FFT normalization: divide by N^2 so energy is resolution-comparable.
    fy = torch.fft.fftfreq(N, device=u.device) * N
    fx = torch.fft.fftfreq(N, device=u.device) * N
    ky, kx = torch.meshgrid(fy, fx, indexing="ij")
    rad = torch.sqrt(kx*kx + ky*ky)
    weight = torch.where(rad > 0, rad ** (1.0 + s), torch.zeros_like(rad))
    z = torch.fft.fft2(u)
    return (weight[None] * z.abs().square()).sum() / (N*N)**2

def tv_energy(u):
    return (u.roll(-1, 0) - u).abs().sum() + (u.roll(-1, 1) - u).abs().sum()

def task_loss(u, target):
    return (u-target).square().mean()

def defect_stats(w, threshold=0.25):
    # Winding lives on plaquettes. Report all nonzero plaquettes, not vortex centers.
    nz = torch.nonzero(w.abs() >= threshold, as_tuple=False)
    vals = [(int(y), int(x), float(w[y, x])) for y, x in nz]
    return {"count": len(vals), "plaquettes": vals}

# ---- Core numerical checks, before optimization ----
s = 0.5
ks = [1, 2, 4, 8]
mode_energies = []
for k in ks:
    ph = 0.7 * torch.sin(2*math.pi*k*X/N)
    u = unit_from_phase(ph)
    mode_energies.append(float(frac_energy(u, s)))
# A small phase amplitude makes the expected single-frequency power law especially clean.
log_slope = float(np.polyfit(np.log(ks), np.log(np.maximum(mode_energies, 1e-30)), 1)[0])
centers = [(N//2-5, N//2-5, 1), (N//2+5, N//2+5, -1)]
target_phase = phase_field(centers)
target = unit_from_phase(target_phase)
target_w = winding(target).detach()
initial_phase = target_phase + 0.65 * torch.randn((N,N), device=DEVICE)
initial = unit_from_phase(initial_phase)
initial_w = winding(initial).detach()
math_check = {
    "frequency_ks": ks,
    "fractional_energies": mode_energies,
    "loglog_slope_expected_1_plus_s": 1+s,
    "loglog_slope_observed": log_slope,
    "target_winding": defect_stats(target_w),
    "initial_winding": defect_stats(initial_w),
    "target_integer_defect_count": int((target_w.abs() > .5).sum().item())
}

# Direct field fitting isolates the regularizer/topology behavior and has matched parameters.
def optimize(kind, steps=700):
    # inverse stereographic-style unconstrained channels, normalized in the loss
    p = torch.nn.Parameter(torch.logit(((initial[0]+1)/2).clamp(.01,.99)).detach().clone())
    q = torch.nn.Parameter(torch.logit(((initial[1]+1)/2).clamp(.01,.99)).detach().clone())
    # A direct angle parameter is better conditioned and still exactly S1-valued.
    angle = torch.nn.Parameter(torch.atan2(initial[1], initial[0]).detach().clone())
    opt = torch.optim.Adam([angle], lr=0.08)
    for step in range(steps):
        u = unit_from_phase(angle)
        fit = task_loss(u, target)
        if kind == "tv":
            reg = tv_energy(u) / (2*N*N)
            loss = fit + 0.018 * reg
        else:
            reg = frac_energy(u, s)
            # Squared winding mismatch; target is the discrete current.
            jac = winding(u)
            jacloss = (jac-target_w).square().mean()
            loss = fit + 0.004 * reg + 10.0 * jacloss
        opt.zero_grad(); loss.backward(); opt.step()
    with torch.no_grad():
        u = unit_from_phase(angle)
        w = winding(u)
        # high-frequency fraction: energy above quarter Nyquist
        fy = torch.fft.fftfreq(N, device=DEVICE)*N; fx=fy
        ky,kx=torch.meshgrid(fy,fx,indexing='ij'); rad=torch.sqrt(kx*kx+ky*ky)
        z=torch.fft.fft2(u); spec=z.abs().square().sum(0)/(N*N)**2
        hf=float(spec[rad > N/4].sum().item())
        return {
            "task_mse": float(task_loss(u,target).item()),
            "fractional_energy": float(frac_energy(u,s).item()),
            "tv_energy": float((tv_energy(u)/(2*N*N)).item()),
            "winding_mse": float((w-target_w).square().mean().item()),
            "winding_defects": defect_stats(w),
            "integer_defect_count": int((w.abs()>.5).sum().item()),
            "high_frequency_spectrum": hf,
            "final_winding_range": [float(w.min()), float(w.max())]
        }

# Wrapped phase circulation is piecewise constant. At a smooth constant field,
# its gradient is generally exactly zero, so this engineering surrogate cannot
# nucleate a defect from a topology-free initialization.
probe = torch.nn.Parameter(torch.zeros((N, N), device=DEVICE))
probe_u = unit_from_phase(probe)
probe_loss = (winding(probe_u) - target_w).square().mean()
probe_loss.backward()
math_check["jacobian_loss_at_constant_field"] = float(probe_loss.item())
math_check["jacobian_gradient_norm_at_constant_field"] = float(probe.grad.norm().item())

results = {"device": DEVICE, "N": N, "math_check": math_check,
           "baseline_tv": optimize("tv"), "idea_fractional_jacobian": optimize("idea")}
print(json.dumps(results, indent=2))
with open("results.json", "w") as f: json.dump(results, f, indent=2)