#!/usr/bin/env python3
"""Derive the polynomial kernels and contact coefficients of the article.

Run ``python code/coefficients.py`` from the project directory. All arithmetic
in this file is exact. The output expressions use SymPy notation.
"""
from __future__ import annotations

import json
from collections import Counter
from pathlib import Path
import sympy as s

OUT = Path(__file__).resolve().parents[1] / "data"
OUT.mkdir(parents=True, exist_ok=True)
CHECKS: Counter[str] = Counter()

def check(group: str, actual, expected=0) -> None:
    difference = s.cancel(s.expand(actual - expected))
    if difference != 0:
        raise AssertionError(f"{group}: nonzero difference {difference}")
    CHECKS[group] += 1

def boolean(group: str, condition: bool) -> None:
    if not condition:
        raise AssertionError(group)
    CHECKS[group] += 1

Z, Y, v, g, h, delta = s.symbols("Z Y v g h delta")
T = s.symbols("T", positive=True)
w = v*(1-v)
AA = Z+2*v*Y+v**2
BB = Z+2*(v-1)*Y+(v-1)**2
CC = Z+2*(v-2)*Y+(v-2)**2
UU = Z+(2*v-1)*Y-w
VV = Z+2*(v-1)*Y+v*(v-2)
NUM = {
    "alternating": s.expand(16*UU**4-16*AA*BB*UU**2+3*AA**2*BB**2),
    "contact": s.expand(BB**2*(4*VV**2-AA*CC)),
}

def lap(poly):
    return s.expand(4*Z*s.diff(poly,Z,2)+8*s.diff(poly,Z)
                    +4*Y*s.diff(poly,Z,Y)+s.diff(poly,Y,2))

def differential_kernel(poly):
    terms=[]
    q=poly
    for j in range(5):
        for (m,l), c in s.Poly(q,Z,Y).terms():
            if c:
                boolean("degree bound j+m+l<=4", j+m+l<=4)
                boolean("heat derivative order<=8", 2*m+l<=8)
        terms.append(s.expand(q.subs({Z:delta*h*h/(4*T*T),Y:h/(2*T)})
                              /(4**j*s.factorial(j)*T**j)))
        q=lap(q)
    check("fifth Laplacian vanishes", q)
    return s.expand(sum(terms))

PK={}
PNULL={}
COEFF={}
for name,num in NUM.items():
    PK[name]=differential_kernel(num)
    PNULL[name]=s.expand(PK[name].subs(delta,0))
    COEFF[name]=[s.expand(PNULL[name]).coeff(T,-j) for j in range(5)]
    check("null decomposition",PNULL[name],sum(COEFF[name][j]/T**j for j in range(5)))
    boolean("null derivative bound", all(s.degree(COEFF[name][j],h)<=j for j in range(5)))

pa=COEFF["alternating"]
ra=s.expand(w*(pa[2]-w*pa[3]+w**2*pa[4]/2))
a=[s.factor(ra.coeff(h,i)) for i in range(5)]
a_expected=[
    6*w*(1+20*w+210*w*w),
    30*w*w*(2*v-1)*(1+21*w),
    w*w*(4+141*w-378*w*w),
    7*w**3*(2*v-1)*(2-3*w),
    s.Rational(1,2)*w**3*(1-w)*(1-3*w),
]
for lhs,rhs in zip(a,a_expected):
    check("alternating coefficient derived from numerator",lhs,rhs)
check("alternating interior adjoint",sum((-1)**i*s.diff(a[i],v,i) for i in range(5)),8)

def boundary_jet(coefs,point,derivative):
    return s.expand(sum((-1)**(r-1-derivative)
                        *s.diff(coefs[r],v,r-1-derivative).subs(v,point)
                        for r in range(derivative+1,len(coefs))))
for d in range(4):
    check("alternating left endpoint jet",-boundary_jet(a,0,d),3 if d==0 else 0)
    check("alternating right endpoint jet",boundary_jet(a,1,d),3 if d==0 else 0)

pc=COEFF["contact"]
Fc=v+2*g-v*v
rc=s.expand((v-2*g)*(pc[2]-Fc*pc[3]+Fc**2*pc[4]/2))
b_left=[s.factor(s.integrate(rc.coeff(h,i),(g,0,v/2))) for i in range(5)]
b_right=[s.factor(s.integrate(rc.coeff(h,i),(g,v-1,v/2))) for i in range(5)]
b_expected=[
 s.Rational(3,2)*v**2*(210*v**4-720*v**3+775*v**2-264*v+19),
 s.Rational(1,2)*v**2*(v-1)*(630*v**4-2040*v**3+1980*v**2-508*v+21),
 s.Rational(1,12)*v**2*(1134*v**6-5724*v**5+11115*v**4-10176*v**3+4230*v**2-586*v+12),
 s.Rational(1,24)*v**3*(v-1)*(252*v**5-1224*v**4+2292*v**3-2004*v**2+753*v-64),
 s.Rational(1,48)*v**4*(v-1)**2*(3*v**2-6*v+4)*(6*v**2-16*v+11),
]
for i in range(5):
    check("contact coefficient derived from numerator",b_left[i],b_expected[i])
    check("contact reflection",b_right[i],(-1)**i*b_left[i].subs(v,2-v))
for bb in [b_left,b_right]:
    check("contact interior adjoint",sum((-1)**i*s.diff(bb[i],v,i) for i in range(5)),8)
for d in range(4):
    check("contact outer left jet",boundary_jet(b_left,0,d))
    check("contact outer right jet",boundary_jet(b_right,2,d))
    check("contact left joining jet",boundary_jet(b_left,1,d),s.Rational(1,4) if d==1 else 0)
    check("contact right joining jet",-boundary_jet(b_right,1,d),-s.Rational(1,4) if d==1 else 0)

moments=[]
for m in range(11):
    ma=s.integrate(s.expand(sum(a[i]*s.diff(v**m,v,i) for i in range(5))),(v,0,1))
    mc=sum(s.integrate(s.expand(sum(bb[i]*s.diff(v**m,v,i) for i in range(5))),(v,lo,hi))
           for bb,lo,hi in [(b_left,0,1),(b_right,1,2)])
    check("alternating polynomial moment",ma,s.Rational(8,m+1)+3+(3 if m==0 else 0))
    check("contact polynomial moment",mc,s.Rational(8*2**(m+1),m+1))
    moments.append({"degree":m,"alternating":str(ma),"contact":str(mc),"sum_three":str(ma+2*mc)})

# The exactly delta-linear quadratic marking difference, before time integration.
def act_quadratic(kernel,heat_shift=0):
    test=v*v+heat_shift
    return s.expand(sum(kernel.coeff(h,r)*s.diff(test,v,r) for r in range(3)))
quad=s.expand(act_quadratic(PK["alternating"],delta/(2*T))
              -act_quadratic(PNULL["alternating"]))
boolean("quadratic correction exactly linear in tilt",s.degree(quad,delta)==1)
D2=s.cancel(quad/delta)
D2_expected=(3*T**4*w**4-T**3*(72*w**3+16*w*w)
             +T*T*(540*w*w+160*w+16)-T*(1440*w+320)+1080)/(2*T**5)
check("quadratic full differential formula",D2,D2_expected)
small=s.expand(T*w*D2*sum((-T*w)**j/s.factorial(j) for j in range(5)))
uv_coeffs={}
for power,expected in [(-4,90),(-3,-s.Rational(206,3)),(-2,s.Rational(55,3)),(-1,-s.Rational(8,3))]:
    val=s.integrate(s.expand(small).coeff(T,power),(v,0,1))
    check("quadratic UV singular time coefficient",val,expected)
    uv_coeffs[str(power)]=str(val)

unmarked=[s.integrate(s.expand(w*pp.subs(h,0)),(v,0,1)) for pp in pa]
A0=unmarked[4]
B0=s.integrate(s.expand(w*(pa[3]-w*pa[4]).subs(h,0)),(v,0,1))
check("unmarked leading numerator",A0,60)
check("unmarked subleading coefficient",B0,-48)
check("unmarked logarithm",s.Rational(moments[0]["alternating"]),14)

# Conditional moments in three real Gaussian variables.
x1,x2,nu,u=s.symbols("x1 x2 nu u",real=True)
def gaussian_monomial_power(power,var):
    if power%2:
        return s.S.Zero
    return s.factorial2(power-1)*var**(power//2) if power else s.S.One

def conditional(num,dd,vv,tt):
    expr=s.Poly(s.expand(num.subs(v,vv).subs({
        Z:x1*x1+x2*x2+dd*nu*nu+u*u/dd,
        Y:x1+s.I*nu+u/dd,
    })),x1,x2,nu)
    out=0
    for powers,c in expr.terms():
        out+=c*gaussian_monomial_power(powers[0],1/(2*tt))*gaussian_monomial_power(powers[1],1/(2*tt))*gaussian_monomial_power(powers[2],1/(2*dd*tt))
    return s.expand(out)

for name in NUM:
    for dd,vv,tt in [(s.Rational(1,4),s.Rational(1,3),s.Rational(2,3)),
                      (s.Rational(1,9),s.Rational(2,3),s.Rational(3,4))]:
        conditional_value=conditional(NUM[name],dd,vv,tt)
        # (-d/du)^r exp(-T u^2/delta), divided by that Gaussian.
        hermite=[s.S.One]
        for r in range(8):
            hermite.append(s.expand(-s.diff(hermite[-1],u)+2*tt*u*hermite[-1]/dd))
        pp=PK[name].subs({delta:dd,v:vv,T:tt})
        operator_value=sum(s.expand(pp).coeff(h,r)*hermite[r] for r in range(9))
        check("independent conditional Gaussian versus differential density",conditional_value,operator_value)

art={
 "alternating_numerator":str(NUM["alternating"]),
 "contact_numerator":str(NUM["contact"]),
 "alternating_residue_coefficients":[str(p) for p in a],
 "contact_left_coefficients":[str(p) for p in b_left],
 "contact_right_coefficients":[str(p) for p in b_right],
 "finite_tilt_operators":{name:str(poly) for name,poly in PK.items()},
 "null_operators":{name:[str(p) for p in pp] for name,pp in COEFF.items()},
 "quadratic_tilt_difference":str(s.factor(D2)),
 "quadratic_UV_time_coefficients":uv_coeffs,
 "residue_moments":moments,
}
(OUT/"coefficients.json").write_text(json.dumps(art,indent=2)+"\n")
print("Polynomial identities and Gaussian moments passed.")
print("Wrote data/coefficients.json.")