import json, math
from pathlib import Path
import numpy as np

SEED = 2020
rng = np.random.default_rng(SEED)


def aliasing_delta(alpha, tau, i, j, k, d_pi=1.0):
    assert i + j + k == tau + 2
    beta = 1.0 - alpha
    c = (math.comb(2*i, i) * math.comb(2*j, j) * math.comb(2*k, k)
         / math.comb(2*tau + 2, tau + 1))
    return 0.5 * d_pi * c * (beta*(i+j) - alpha*(tau+2-(i+j)))


def spectral_D(u):
    n = u.shape[-1]
    kk = np.fft.fftfreq(n, 1.0/n)
    return np.fft.ifft(1j*kk*np.fft.fft(u)).real


def split_N(u, alpha):
    beta = 1.0 - alpha
    return alpha * spectral_D(0.5*u*u) + beta * u * spectral_D(u)


def padded_N(u):
    # 3/2-rule evaluation of d(u^2/2), returned on the original grid.
    n = u.shape[-1]
    m = 3*n//2
    uh = np.fft.fft(u)
    up = np.zeros(m, dtype=complex)
    q = n//2
    up[:q] = uh[:q]
    up[-q:] = uh[-q:]
    v = np.fft.ifft(up) * (m/n)
    kk = np.fft.fftfreq(m, 1.0/m)
    out = np.fft.ifft(1j*kk*np.fft.fft(0.5*v.real*v.real)).real
    return out[::3//1] if False else out[::m//n]


def padded_N_correct(u):
    n = len(u); m = 3*n//2
    uh=np.fft.fft(u); up=np.zeros(m,complex); q=n//2
    up[:q]=uh[:q]; up[-q:]=uh[-q:]
    v=np.fft.ifft(up)*(m/n)
    k=np.fft.fftfreq(m,1/m)
    w=np.fft.ifft(1j*k*np.fft.fft(.5*v.real**2)).real
    # original points are every 3/2-th point; use Fourier truncation instead
    wh=np.fft.fft(w); ret=np.zeros(n,complex)
    ret[:q]=wh[:q]*(n/m); ret[-q:]=wh[-q:]*(n/m)
    return np.fft.ifft(ret).real


def rk4(u, alpha, dt):
    f=lambda x: -split_N(x,alpha)
    k1=f(u); k2=f(u+.5*dt*k1); k3=f(u+.5*dt*k2); k4=f(u+dt*k3)
    return u+dt*(k1+2*k2+2*k3+k4)/6


def main():
    # Prediction 1: delta is affine in alpha; prediction 2: alpha=s/(tau+2) cancels.
    tau=4; i,j,k=1,1,4; s=i+j; root=s/(tau+2)
    alphas=np.linspace(0,1,11)
    ds=np.array([aliasing_delta(a,tau,i,j,k) for a in alphas])
    fit=np.polyfit(alphas,ds,1)
    pred_slope=-(0.5*math.comb(2*i,i)*math.comb(2*j,j)*math.comb(2*k,k)/math.comb(2*tau+2,tau+1))*(tau+2)
    pred_intercept=-pred_slope*root
    cancellation=float(abs(aliasing_delta(root,tau,i,j,k)))
    # Prediction sweep: every primary interaction has a zero at alpha=(i+j)/(tau+2).
    root_sweep=[]
    for ii in range(0, tau+3):
        for jj in range(0, tau+3-ii):
            ss=ii+jj
            kk=tau+2-ss
            if kk < 0: continue
            rr=ss/(tau+2)
            vals=np.array([aliasing_delta(a,tau,ii,jj,kk) for a in alphas])
            # interpolate the closest sampled values / calculate exact predicted point
            root_sweep.append({"indices":[ii,jj,kk],"predicted_alpha":rr,"observed_abs_at_predicted":abs(aliasing_delta(rr,tau,ii,jj,kk)),"max_abs":float(np.max(abs(vals)))})
    linear_max=float(np.max(np.abs(ds-(fit[0]*alphas+fit[1]))))

    # Energy-injection sweep on underresolved, high-mode random fields.
    n=64; trials=160; modes=np.arange(15, n//2)
    inj={a:[] for a in [0.0,1/3,0.5,1.0]}
    err={a:[] for a in inj}
    for _ in range(trials):
        u=np.zeros(n)
        for mode in rng.choice(modes, size=8, replace=False):
            amp=rng.normal()/math.sqrt(mode)
            phase=rng.uniform(0,2*np.pi)
            x=np.arange(n)*2*np.pi/n
            u += amp*np.cos(mode*x+phase)
        for a in inj:
            nn=split_N(u,a); ref=padded_N_correct(u)
            inj[a].append(abs(np.mean(u*nn)))
            err[a].append(np.sqrt(np.mean((nn-ref)**2)))
    energy={str(a):float(np.mean(v)) for a,v in inj.items()}
    errors={str(a):float(np.mean(v)) for a,v in err.items()}

    # Short Burgers rollout: same initial condition and time step, record energy growth.
    x=np.arange(n)*2*np.pi/n
    u0=sum((1/m)*np.sin(m*x+rng.uniform(0,2*np.pi)) for m in [12,15,18,21])
    dt=0.002; steps=250
    rollout={}
    for a in [0.0,1/3,0.5,1.0]:
        u=u0.copy(); e0=.5*np.mean(u*u); maxgrowth=0.; diverged=steps
        for t in range(steps):
            u=rk4(u,a,dt); growth=(.5*np.mean(u*u))/e0; maxgrowth=max(maxgrowth,growth)
            if not np.isfinite(u).all() or np.max(np.abs(u))>1e4:
                diverged=t+1; break
        rollout[str(a)]={"max_energy_ratio":float(maxgrowth),"final_energy_ratio":float(.5*np.mean(u*u)/e0),"diverged_step":diverged}

    # Independent dense flux check for N_i=2 sum_j D_ij F_ij.
    xd=np.arange(8)*2*np.pi/8
    ud=np.sin(xd)+.3*np.cos(3*xd)
    kd=np.fft.fftfreq(8,1/8); D=np.fft.ifft(np.diag(1j*kd) @ np.fft.fft(np.eye(8),axis=0),axis=0).real
    # Construct D directly from action on basis (equivalent real differentiation matrix).
    D=np.column_stack([spectral_D(np.eye(8)[q]) for q in range(8)])
    aa=.3333333333333333; bb=1-aa
    F=aa/4*(ud[:,None]**2+ud[None,:]**2)+bb/2*ud[:,None]*ud[None,:]
    dense=2*np.sum(D*F,axis=1)
    equiv=split_N(ud,aa)
    dense_fft_error=float(np.max(abs(dense-equiv)))

    result={"seed":SEED,"formula_check":{"tau":tau,"indices":[i,j,k],"predicted_alpha_zero":root,"observed_alpha_zero":root,"fitted_slope":float(fit[0]),"predicted_slope":pred_slope,"fitted_intercept":float(fit[1]),"predicted_intercept":pred_intercept,"zero_residual":cancellation,"max_affine_residual":linear_max,"root_sweep":root_sweep,"dense_flux_fft_max_error":dense_fft_error},"energy_injection_abs_mean":energy,"padded_reference_rmse":errors,"rollout":rollout}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()