import json, math, time
import numpy as np
import torch
from hodge_dual import curl_2d, divergence_2d

torch.set_default_dtype(torch.float64)
torch.manual_seed(7); np.random.seed(7)

def spectrum(N, eps):
    h=1.0/N
    # central-difference curl symbol; zero mode is gauge nullspace
    vals=[]
    for k in range(N):
        sx=math.sin(2*math.pi*k/N)/h
        for l in range(N):
            sy=math.sin(2*math.pi*l/N)/h
            vals.append((sx*sx+sy*sy)/eps)
    return np.array(vals)/(N*N)

def divergence_check():
    N=24; h=1/N
    A=torch.randn(1,N,N,3)*3.7
    p0=torch.randn(1,N,N,2)
    r=divergence_2d(curl_2d(A,h),h)
    # random amplitudes demonstrate independence from potential magnitude
    rows=[]
    for scale in [0.,1e-4,1.,1e4]:
        rr=divergence_2d(curl_2d(A*scale,h),h)
        rows.append(float(rr.abs().max()))
    return rows

def gd_mode(N, eps, eta, steps=80):
    h=1/N
    # one Fourier-like real perturbation, with p0=0; energy is purely quadratic
    x=torch.arange(N).reshape(1,N,1).double(); y=torch.arange(N).reshape(1,1,N).double()
    A=torch.zeros(1,N,N,3)
    A[...,2]=torch.sin(2*math.pi*6*x/N)*torch.sin(2*math.pi*6*y/N)
    initial=0.5*(curl_2d(A,h)**2).sum()/eps/(N*N)
    vals=[]
    for _ in range(steps):
        A.requires_grad_(True)
        loss=0.5*(curl_2d(A,h)**2).sum()/eps/(N*N)
        g=torch.autograd.grad(loss,A)[0]
        with torch.no_grad(): A=A-eta*g
        vals.append(float(loss))
    return float(initial), vals[-1], vals

def primal_compare(N=24):
    # Same central differences: primal Poisson objective and dual quadratic have
    # identical positive conditioning scale (dual additionally enforces div p).
    vals=spectrum(N,1.0); positive=vals[vals>1e-12]
    return {"primal_condition_number":float(positive.max()/positive.min()),
            "dual_condition_number":float(positive.max()/positive.min()),
            "dual_gauss_residual":max(divergence_check())}

def main():
    N=24
    divs=divergence_check()
    # Prediction 1: div(curl A)=0 to roundoff at every amplitude.
    div_pred=0.0
    # Prediction 2: lambda_max(eps)*eps is constant.
    epss=[0.5,1.,2.,4.]
    lams=[float(spectrum(N,e).max()) for e in epss]
    products=[e*l for e,l in zip(epss,lams)]
    # Prediction 3: GD changes from contraction to divergence at eta*lambda_max=2.
    lam=lams[1]
    eta_crit=2/lam
    tests=[]
    for mult in [0.90,0.99,1.01,1.10]:
        initial, final, _=gd_mode(N,1.,mult*eta_crit,steps=30)
        tests.append({"eta_over_critical":mult,"final_over_initial":final/initial,
                      "stable_observed":bool(final < initial*10)})
    out={"grid":N,"divergence_max_abs_by_amplitude":divs,
         "prediction_1":{"predicted":"0 exactly (up to floating point)","observed_max":max(divs)},
         "prediction_2":{"predicted":"lambda_max * eps is constant","eps":epss,
                          "lambda_max":lams,"products":products,
                          "relative_spread":(max(products)-min(products))/np.mean(products)},
         "prediction_3":{"predicted":"transition at eta*lambda_max=2","lambda_max":lam,
                          "eta_critical":eta_crit,"tests":tests},
         "baseline_context":primal_compare(N),
         "note":"The primal and dual constant-coefficient quadratic operators have the same nonzero Fourier condition number; the dual advantage tested here is exact constraint satisfaction, not a universal conditioning improvement."}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))

if __name__=='__main__': main()