#!/usr/bin/env python3
"""Write operators.txt: the four Picard-Fuchs operators of the period geometries through 5PM in theta form
(as printed in the paper) and in d/dz form (exact conversion), their variables, finite singular
points with local exponents, the MUM Frobenius convention used for the bound-arc value files, and
the continuation path classes.  Every d/dz form is checked by applying it to the independently
generated holomorphic period series through z^30 (exact rational arithmetic).
usage: build_operators.py OUT.txt"""
import sys
from fractions import Fraction as F
from math import comb, factorial
import sympy as sp

out = sys.argv[1]
z, r = sp.symbols('z rho')

# theta rows: L = sum_s z^s R_s(theta), R_s given by ascending coefficients in theta
OPS = {
 'L_K3':  {'rows': None, 'paper': "(2theta-1)^3 + z^2 (2theta+1)^3 - 4 z theta (4theta^2+1)", 'var': "z = x^2",
           'note': "4PM K3 (symmetric square of the Legendre operator).  The value files use the equivalent operator "
                   "(1/8) z^(-1/2) L_K3 z^(1/2) = theta^3 - z (2theta^3 + 3theta^2 + 2theta + 1/2) + z^2 (theta+1)^3, which is MUM at z = 0 with "
                   "all local exponents 0 and holomorphic solution varpi_0 = 2F1(1/2,1/2;1;z)^2 = (2/pi)^2 K(z)^2."},
 'L_A':   {'rows': [[0,0,0,1],[-5,-27,-51,-34],[1,3,3,1]], 'paper': "theta^3 - z (2theta+1)(17theta^2+17theta+5) + z^2 (theta+1)^3", 'var': "z = x^2",
           'note': "5PM-2SF topology 40, Apery / Beukers-Peters K3 surface K3'.  varpi_0 = sum_n A_n z^n, A_n = sum_k C(n,k)^2 C(n+k,k)^2 = 1, 5, 73, 1445, ..."},
 'L_CY3': {'rows': [[0,0,0,0,1],[-16,-128,-384,-512,-256]], 'paper': "theta^4 - 2^8 z (theta+1/2)^4", 'var': "z = 2^-8 x^4",
           'note': "5PM-1SF topology 3, Calabi-Yau threefold CY3.  varpi_0 = 4F3(1/2,1/2,1/2,1/2;1,1,1;2^8 z)."},
 "L_CY3'":{'rows': [[0,0,0,0,1],[-16*7,-16*48,-16*112,-16*128,-16*192],[2**14*7,2**14*64,2**14*208,2**14*256,2**14*192],
                    [-(2**26),-(2**29),-3*2**29,-(2**31),-(2**30)]],
           'paper': "theta^4 - 2^4 z (192theta^4+128theta^3+112theta^2+48theta+7) + 2^14 z^2 (192theta^4+256theta^3+208theta^2+64theta+7) - 2^30 z^3 (theta+1/2)^4",
           'var': "z = (1 - gamma^2)/2^10", 'note': "5PM-2SF topology 37, Calabi-Yau threefold CY3' (Hadamard type, chi = 80).  varpi_0 = 1 + 112 z + 47376 z^2 + ..."},
}
# K3: build the paper-form rows and the engine-form rows exactly
th = sp.symbols('theta')
def rows_from_poly_in_theta_z(expr):
    P = sp.Poly(sp.expand(expr), z, th)
    S = P.degree(z); n = P.degree(th)
    rows = [[F(0)]*(n+1) for _ in range(S+1)]
    for (sz, jt), c in P.terms():
        rows[sz][jt] = F(int(sp.numer(c)), int(sp.denom(c)))
    return [[x for x in row] for row in rows]
K3_paper = (2*th-1)**3 + z**2*(2*th+1)**3 - 4*z*th*(4*th**2+1)
K3_engine = th**3 - z*(2*th**3+3*th**2+2*th+sp.Rational(1,2)) + z**2*(th+1)**3
# check the conjugation identity: z^(-1/2) L(theta) z^(1/2) = L(theta+1/2)
assert sp.expand(K3_paper.subs(th, th+sp.Rational(1,2))/8 - K3_engine) == 0
OPS['L_K3']['rows'] = rows_from_poly_in_theta_z(K3_paper)
OPS['L_K3']['rows_engine'] = rows_from_poly_in_theta_z(K3_engine)
for k in ('L_A','L_CY3',"L_CY3'"):
    OPS[k]['rows'] = [[F(c) for c in row] for row in OPS[k]['rows']]
    # cross-check typed rows against the paper's displayed polynomial
    disp = {'L_A': th**3 - z*(2*th+1)*(17*th**2+17*th+5) + z**2*(th+1)**3,
            'L_CY3': th**4 - 2**8*z*(th+sp.Rational(1,2))**4,
            "L_CY3'": th**4 - 2**4*z*(192*th**4+128*th**3+112*th**2+48*th+7) + 2**14*z**2*(192*th**4+256*th**3+208*th**2+64*th+7) - 2**30*z**3*(th+sp.Rational(1,2))**4}[k]
    assert rows_from_poly_in_theta_z(disp) == OPS[k]['rows'], k

def stirling2(n):
    S = [[0]*(n+1) for _ in range(n+1)]; S[0][0] = 1
    for k in range(1, n+1):
        for m in range(1, k+1):
            S[k][m] = S[k-1][m-1] + m*S[k-1][m]
    return S
def theta_to_dz(rows):
    """L = sum_s z^s sum_j R[s][j] theta^j  ->  list p[i] (sympy poly in z) with L = sum_i p_i(z) (d/dz)^i."""
    n = max(len(rw) for rw in rows) - 1
    S2 = stirling2(n)
    p = [sp.Integer(0)]*(n+1)
    for s, row in enumerate(rows):
        for j, c in enumerate(row):
            if c == 0: continue
            for i in range(0, j+1):
                if S2[j][i]:
                    p[i] += sp.Rational(c.numerator, c.denominator)*S2[j][i]*z**(s+i)
    return [sp.expand(x) for x in p]
def apply_theta_rows_to_series(rows, a, N):
    """coefficients of L(sum a_n z^n) through z^N, exact."""
    outc = []
    for m in range(N+1):
        tot = F(0)
        for s, row in enumerate(rows):
            n = m - s
            if n < 0 or n >= len(a): continue
            val = sum(F(c)*F(n)**j for j, c in enumerate(row))
            tot += val*a[n]
        outc.append(tot)
    return outc
def apply_dz_to_series(p, a, N):
    f = sum(sp.Rational(a[n].numerator, a[n].denominator)*z**n for n in range(len(a)))
    Lf = sum(p[i]*sp.diff(f, z, i) for i in range(len(p)))
    P = sp.Poly(sp.expand(Lf), z)
    return [P.coeff_monomial(z**m) for m in range(N+1)]
# independent holomorphic series
NS = 36
poch = lambda n: F(factorial(2*n), 4**n*factorial(n)**2)   # (1/2)_n/n! = C(2n,n)/4^n
f21 = [poch(n)**2 for n in range(NS)]
k3sq = [sum(f21[i]*f21[n-i] for i in range(n+1)) for n in range(NS)]           # (2K/pi)^2
apery = [F(sum(comb(n,k)**2*comb(n+k,k)**2 for k in range(n+1))) for n in range(NS)]
f43 = [poch(n)**4*F(256)**n for n in range(NS)]
def holo_from_rows(rows, N):  # MUM recurrence (used for CY3' only; first terms checked against the paper)
    a = [F(1)]
    for m in range(1, N):
        acc = F(0)
        for s in range(1, len(rows)):
            n = m - s
            if n < 0: continue
            acc += sum(F(c)*F(n)**j for j, c in enumerate(rows[s]))*a[n]
        lead = sum(F(c)*F(m)**j for j, c in enumerate(rows[0]))
        a.append(-acc/lead)
    return a
had = holo_from_rows(OPS["L_CY3'"]['rows'], NS)
assert had[:3] == [1, 112, 47376]
HOLO = {'L_K3': ('engine form, varpi_0 = (2/pi)^2 K(z)^2 generated from 2F1(1/2,1/2;1;z)^2', k3sq, 'rows_engine'),
        'L_A': ('varpi_0 generated from the Apery numbers', apery, 'rows'),
        'L_CY3': ('varpi_0 generated from the 4F3 coefficients ((1/2)_n/n!)^4 2^(8n)', f43, 'rows'),
        "L_CY3'": ('varpi_0 generated by the MUM recurrence (first terms 1, 112, 47376 as in the paper)', had, 'rows')}

def indicial(p, z0):
    """local exponents of sum_i p_i (d/dz)^i at a finite point z0 (roots of the indicial polynomial, with multiplicity)."""
    t = sp.symbols('t')
    terms = []
    for i, pi in enumerate(p):
        q = sp.Poly(sp.expand(pi.subs(z, t+z0)), t)
        for (m,), c in q.terms():
            terms.append((m - i, i, sp.nsimplify(c)))
    lo = min(x[0] for x in terms)
    ind = sum(c*sp.ff(r, i) for (d, i, c) in terms if d == lo)
    ind = sp.factor(sp.expand(ind))
    roots = sp.roots(sp.Poly(ind, r))
    return sorted([rt for rt, mult in roots.items() for _ in range(mult)], key=lambda x: sp.re(sp.N(x))), ind
def indicial_inf(rows):
    """exponents at z = infinity: with w = 1/z, theta_z = -theta_w, the indicial polynomial at w = 0 is R_S(-rho)."""
    RS = rows[-1]
    ind = sum(sp.Rational(c.numerator, c.denominator)*(-r)**j for j, c in enumerate(RS))
    roots = sp.roots(sp.Poly(sp.expand(ind), r))
    return sorted([rt for rt, mult in roots.items() for _ in range(mult)], key=lambda x: sp.re(sp.N(x)))

def fmt_rows(rows):
    s = []
    for si, row in enumerate(rows):
        poly = ' + '.join(f'({c})*theta^{j}' for j, c in enumerate(row) if c != 0)
        s.append(f'    z^{si}: [{", ".join(str(c) for c in row)}]   i.e. z^{si} * ({poly})')
    return '\n'.join(s)

L = []
L.append("Picard-Fuchs operators of the four non-polylogarithmic period geometries through 5PM")
L.append("(normalization of Klemm, Nega, Sauer and Plefka, arXiv:2401.07899; theta = z d/dz).")
L.append("This file accompanies the value files bound_frobenius/values_*.txt and Table 6 of the paper.")
L.append("")
L.append("CONVENTIONS")
L.append("  theta form: L = sum_s z^s R_s(theta); each R_s is listed by its coefficients in ascending powers of theta.")
L.append("  d/dz form:  L = sum_i p_i(z) (d/dz)^i, obtained from the theta form with theta^j = sum_i S2(j,i) z^i (d/dz)^i")
L.append("              (Stirling numbers of the second kind); exact integer coefficients.")
L.append("  Frobenius basis at the MUM point z = 0 (all four operators, after the z^(1/2) conjugation of L_K3 noted below):")
L.append("     varpi_k(z) = sum_{j=0..k} (ln z)^j / j! * h_{k-j}(z),  k = 0..r-1,  h_0(0) = 1,  h_j(0) = 0 for j >= 1,")
L.append("  so varpi_0 is the power-series solution, varpi_1 = varpi_0 ln z + h_1, etc.  The value files list the Wronskian")
L.append("  matrix W[i,k] = (d/dz)^(i-1) varpi_(k-1)(z), i,k = 1..r, continued along the stated path (ln z continued along the same path).")
L.append("  In this frame the unipotent MUM monodromy acts by ln z -> ln z + 2 pi i, i.e. varpi_k -> sum_j (2 pi i)^(k-j)/(k-j)! varpi_j.")
L.append("")
for name in ('L_K3','L_A','L_CY3',"L_CY3'"):
    o = OPS[name]
    L.append("="*100)
    L.append(f"{name}:  {o['paper']}     variable {o['var']}")
    L.append(f"  {o['note']}")
    L.append("  theta rows (paper form):")
    L.append(fmt_rows(o['rows']))
    key = HOLO[name][2]
    rows_used = o[key]
    if name == 'L_K3':
        L.append("  theta rows (form used for the value files, (1/8) z^(-1/2) L_K3 z^(1/2)):")
        L.append(fmt_rows(rows_used))
    p = theta_to_dz(rows_used)
    L.append("  d/dz form" + (" (of the form used for the value files)" if name == 'L_K3' else "") + ":")
    for i, pi in enumerate(p):
        L.append(f"    p_{i}(z) = {sp.factor(pi)}")
    # series check
    desc, a, _ = HOLO[name]
    res_theta = apply_theta_rows_to_series(rows_used, a[:32], 30)
    res_dz = apply_dz_to_series(p, a[:34], 30)
    ok = all(x == 0 for x in res_theta) and all(x == 0 for x in res_dz)
    L.append(f"  check: operator applied to the holomorphic series ({desc}) vanishes identically through z^30 "
             f"in both forms: {'yes' if ok else 'NO'}")
    assert ok, name
    L.append(f"  first coefficients of varpi_0: {', '.join(str(x) for x in a[:6])}, ...")
    # singular points
    lead = sp.Poly(p[-1], z)
    L.append(f"  leading coefficient p_{len(p)-1}(z) = {sp.factor(p[-1])}")
    fin = sp.roots(lead)
    pts = sorted(fin.keys(), key=lambda x: sp.re(sp.N(x)))
    L.append("  local exponents (roots of the indicial equation):")
    for z0 in pts:
        ex, ind = indicial(p, z0)
        L.append(f"    z = {z0} (~{sp.N(z0, 8)}): {{{', '.join(str(e) for e in ex)}}}")
    exinf = indicial_inf(rows_used)
    L.append(f"    z = infinity (in w = 1/z): {{{', '.join(str(e) for e in exinf)}}}")
    L.append("")
L.append("="*100)
L.append("CONTINUATION PATHS used for bound_frobenius/values_*.txt (gamma = cos(theta) in (0,1), x = exp(-i theta) on the physical branch):")
L.append("  K3 and K3':  z = x^2 on |z| = 1, lower arc.  Two homotopic polygonal chains in the region Im z < 0, starting at z = 1/200:")
L.append("     path A: 1/200 -> 3/100 - 3i/100 -> 0.55 z_t -> z_t ;   path B: 1/200 -> 8/1000 - 6i/100 -> 0.8 exp(-i/5) z_t -> z_t .")
L.append("     All finite singular points of both operators lie on the positive real axis, so the region Im z < 0 contains none of them")
L.append("     and any two paths in it are homotopic; for K3' the physical path from the scattering segment (z_+, 1) through z - i0 encloses no singular point.")
L.append("  CY3 (4F3):   z = x^4/2^8 lies ON the singular circle |z| = 2^-8 at arg z = -4 theta in (-2 pi, 0).  Physical class: from z = 2^-8/50,")
L.append("     clockwise along a spiral inside the disk (radius 0.70 * 2^-8, angular steps 1/2, path A; radius 0.85 * 2^-8, steps 0.35, path B),")
L.append("     monotone in angle and never crossing arg z = 0, ending at z_t.  This is the class of z - i0 continued past z = 2^-8 (gamma = 1).")
L.append("     The short counterclockwise route inside the disk to the same endpoint differs from the physical class by one turn around")
L.append("     z = 0 alone, so the two Frobenius bases differ by the unipotent MUM monodromy (the 4F3 path check of the paper at gamma = 1/2).")
L.append("  CY3' (Hadamard):  z = (1-gamma^2)/2^10 is real in (0, 2^-10); path A is the real segment from z = 2^-13, path B a detour through")
L.append("     Im z < 0 (homotopic; no singular point in between).  All entries are real.")
L.append("  Working precision 600 bits; the two paths agree and the rigorous ball radii are as listed per point in README_bound.md.")
open(out, 'w').write('\n'.join(L) + '\n')
print('wrote', out.split('/')[-1], '; all four series checks passed; K3 conjugation identity verified')