# SPDX-License-Identifier: GPL-2.0-or-later
#
# Derived from SageMath (https://www.sagemath.org),
# sage/schemes/hyperelliptic_curves/mestre.py (Florian Bouyer, Marco Streng; Copyright (C)
# 2011-2013, and 2025 Sabrina Kunzweiler, Gareth Ma, Giacomo Pope) and invariants.py (Nick
# Alexander; Copyright (C) 2008), licensed under the GNU General Public License v2.0 or
# later. This Python transliteration of Mestre's conic-and-cubic construction keeps
# SageMath's conventions and normalizations. It is Copyright (c) 2026 Anthropic, PBC;
# created by Matthew D. Schwartz, code written by Claude (Anthropic) under his supervision,
# and is distributed under the same GNU General Public License v2.0 or later. The license
# text is LICENSE-GPL-2.0 beside this file; the rest of the bundle is MIT (see NOTICE).

"""WatsonW(w, a, b, c): the anisotropic simple-cubic Watson integral in closed form.

    W_S(alpha_1, alpha_2, alpha_3; w) = (1/pi^3) int_{[0,pi]^3} d^3k / (w - alpha_1 cos k_1 - alpha_2 cos k_2 - alpha_3 cos k_3)
                                      = |Delta M| / (4 pi^2 w)

for rates a, b, c = alpha_1, alpha_2, alpha_3 and spectral variable w > a + b + c: the constant law
(Theorem 2) of "The anisotropic Watson integral" (Noam Elkies, Thomas W. Grimm, Matthew D. Schwartz),
in the seven steps of the paper's section on the explicit evaluation. Delta is the determinant of the
marked quadratic partition of a genus-two sextic y^2 = f(t) carrying the four Igusa-Clebsch invariants
of the tuple, and M is the period minor selected by the cyclic-adjacency rule.

Inputs: w, a, b, c (numbers, or strings mpmath can read, such as '27/2') and the digit count dps
(default 30). Output: WatsonW(w, a, b, c, dps) returns W_S as an mpmath mpf at about dps digits; the
sextic is built at dps + 30 digits and the period integrals at dps. No other file is read and no
digit strings are stored: every value is computed from exact integer and rational coefficients.

Example:
    python3 WatsonW.py           # prints WatsonW(17, 1, 5, 7) = 0.069964276874484951...
    >>> from WatsonW import WatsonW
    >>> WatsonW(22, 5, 6, 9)     # the second tuple of Table 2 of the paper: 0.05723479465070378...

The script takes no command-line options (-h or --help prints this text). Requires: Python >= 3.9
and mpmath (tested with mpmath 1.3.0); nothing else. Without mpmath it stops with one line naming
the package (exit 2).
"""
import sys
if __name__ == "__main__" and len(sys.argv) > 1 and sys.argv[1] in ("-h", "--help"):
    print(__doc__.strip())
    sys.exit(0)
try:
    from mpmath import mp, mpf, sqrt, pi, matrix, eigsy, polyroots, fprod, det, quad, sin, fabs, im, re
except ImportError:
    sys.stderr.write("WatsonW.py needs the mpmath package: python3 -m pip install mpmath\n")
    sys.exit(2)

# Polynomials in one variable are coefficient lists, highest power first (the convention of polyroots).
def polmul(p, q):
    r = [0]*(len(p) + len(q) - 1)
    for i, u in enumerate(p):
        for j, v in enumerate(q): r[i + j] += u*v
    return r
def poladd(p, q):
    n = max(len(p), len(q)); p, q = [0]*(n - len(p)) + p, [0]*(n - len(q)) + q
    return [u + v for u, v in zip(p, q)]
def polscale(k, p): return [k*u for u in p]

# MestreSextic(i2, i4, i6, i10, sgn) returns the coefficients of a sextic f(t) whose genus-two curve
# y^2 = f(t) has Igusa-Clebsch invariants (i2 : i4 : i6 : i10): Mestre's conic-and-cubic construction,
# with sgn selecting between the two real points used on the conic. WatsonW tries sgn = 1 and falls
# back to sgn = -1 so that all six roots of f are real.
def MestreSextic(i2, i4, i6, i10, sgn=1):
    x = 8*(1 + 20*i4/i2**2)/225
    y = 16*(1 + 80*i4/i2**2 - 600*i6/i2**3)/3375
    z = -64*(-10800000*i10/i2**5 - 9 - 700*i4/i2**2 + 3600*i6/i2**3 + 12400*i4**2/i2**4 - 48000*i4*i6/i2**5)/253125
    L = matrix([[x + 6*y, 6*x**2 + 2*y, 2*z],
                [6*x**2 + 2*y, 2*z, 9*x**3 + 4*x*y + 6*y**2],
                [2*z, 9*x**3 + 4*x*y + 6*y**2, 6*x**2*y + 2*y**2 + 3*x*z]])
    E, Q = eigsy(L)                      # eigenvalues ascending, orthonormal eigenvectors in the columns
    P = [Q[k, 2]/sqrt(E[2]) + sgn*Q[k, 0]/sqrt(-E[0]) for k in range(3)]
    V = [[0], [1], [1, 0]]               # V = (0, 1, s), each entry a polynomial in s
    VLV = [L[2, 2], 2*L[1, 2], L[1, 1]]  # V.L.V
    LP = [sum(L[i, k]*P[k] for k in range(3)) for i in range(3)]
    PLV = [LP[2], LP[1]]                 # P.L.V
    U = [poladd(polscale(P[i], VLV), polscale(-2, polmul(PLV, V[i]))) for i in range(3)]
    c = {(1, 1, 1): 12*x*y - 2*y/3 - 4*z,
         (1, 1, 2): -18*x**3 - 12*x*y - 36*y**2 - 2*z,
         (1, 1, 3): -9*x**3 - 36*x**2*y - 4*x*y - 6*x*z - 18*y**2,
         (1, 2, 2): -9*x**3 - 36*x**2*y - 4*x*y - 6*x*z - 18*y**2,
         (1, 2, 3): -54*x**4 - 36*x**2*y - 36*x*y**2 - 6*x*z - 4*y**2 - 24*y*z,
         (1, 3, 3): -27*x**4/2 - 72*x**3*y - 6*x**2*y - 9*x**2*z - 39*x*y**2 - 36*y**3 - 2*y*z,
         (2, 2, 2): -27*x**4 - 18*x**2*y - 6*x*y**2 - 8*y**2/3 + 2*y*z,
         (2, 2, 3): 9*x**3*y - 27*x**2*z + 6*x*y**2 + 18*y**3 - 8*y*z,
         (2, 3, 3): -81*x**5/2 - 27*x**3*y - 9*x**2*y**2 - 4*x*y**2 + 3*x*y*z - 6*z**2,
         (3, 3, 3): 27*x**4*y/2 - 27*x**3*z/2 + 9*x**2*y**2 + 3*x*y**3 - 6*x*y*z + 4*y**3/3 - 10*y**2*z}
    f = [0]
    for (i, j, k), cijk in c.items():
        f = poladd(f, polscale(cijk, polmul(polmul(U[i - 1], U[j - 1]), U[k - 1])))
    return f

# WatsonW(w, alpha_1, alpha_2, alpha_3), the seven steps of the closed form. The sextic is built at dps + 30 digits
# (Mestre's construction loses up to about twenty-five digits at these tuples), the period integrals at dps.
def WatsonW(w, a, b, c, dps=30):
    mp.dps = dps + 30
    w, a, b, c = mpf(w), mpf(a), mpf(b), mpf(c)
    # (i) elementary symmetric functions e1, e2, e3 of the squared rates and x = 1/w^2
    e1, e2, e3 = a**2 + b**2 + c**2, a**2*b**2 + a**2*c**2 + b**2*c**2, a**2*b**2*c**2
    x = 1/w**2
    # (ii) the threshold quartic Q4 and the marked branch R = -sqrt(Q4)
    q4 = (e1**4*x**4 - 4*e1**3*x**3 - 8*e1**2*e2*x**4 + 6*e1**2*x**2 + 16*e1*e2*x**3
          + 16*e2**2*x**4 - 8*e2*x**2 - 64*e3*x**3 - 4*e1*x + 1)
    rr = -sqrt(q4)
    # (iii) the four Igusa-Clebsch invariants in closed form, polynomials in e1, e2, e3, x, R
    i2 = 48*e1**2*x**2 - 32*e1*x - 192*e2*x**2 + 80 - 48*rr
    i4 = 544*e1**2*x**2 - 1088*e1*x - 1152*e2*x**2 + 544 - 480*rr
    i6 = (16384*e1**4*x**4 - 48384*e1**3*x**3 - 114688*e1**2*e2*x**4 + 64256*e1**2*x**2
          + 177152*e1*e2*x**3 - 48896*e1*x + 196608*e2**2*x**4 - 99328*e2*x**2
          - 417792*e3*x**3 + 16640 + (-16384*e1**2*x**2 + 17152*e1*x + 49152*e2*x**2 - 16128)*rr)
    i10 = 32768*e3*x**3*(e1**2*x**2 - 2*e1*x - 4*e2*x**2 + 1 + rr)
    # (iv) a real sextic model from MestreSextic, sgn = 1 then -1 so that all six roots are real
    sgn = 1; f = MestreSextic(i2, i4, i6, i10, sgn)
    r = polyroots(f, maxsteps=200, extraprec=200)
    if max(fabs(im(z)) for z in r) > mpf(10)**-10:
        sgn = -1; f = MestreSextic(i2, i4, i6, i10, sgn)
        r = polyroots(f, maxsteps=200, extraprec=200)
    r = sorted(re(z) for z in r)
    lc = f[0]
    # (v) the marked partition into three real quadratics, by the disc-resultant normalization test
    #     against x {a^2, b^2, c^2}, and its determinant Delta (quadratics monic, from pairs of roots)
    prs = [(1,2,3,4,5,6), (1,3,2,4,5,6), (1,4,2,3,5,6), (1,2,3,5,4,6), (1,2,3,6,4,5),
           (1,3,2,5,4,6), (1,3,2,6,4,5), (1,4,2,5,3,6), (1,4,2,6,3,5), (1,5,2,3,4,6),
           (1,5,2,4,3,6), (1,5,2,6,3,4), (1,6,2,3,4,5), (1,6,2,4,3,5), (1,6,2,5,3,4)]
    disc = lambda i, j: (r[i-1] - r[j-1])**2
    res = lambda i, j, k, l: (r[i-1] - r[k-1])*(r[i-1] - r[l-1])*(r[j-1] - r[k-1])*(r[j-1] - r[l-1])
    def score(pr):
        i, j, k, l, m, n = pr
        v = sorted(fabs(u) for u in [disc(i, j)*res(k, l, m, n), disc(k, l)*res(m, n, i, j), disc(m, n)*res(i, j, k, l)])
        tgt = sorted([x*a**2, x*b**2, x*c**2])
        q = [v[s]/tgt[s] for s in range(3)]
        return max(q)/min(q)
    best = min(prs, key=score)
    g = [[1, -(r[i-1] + r[j-1]), r[i-1]*r[j-1]] for i, j in (best[0:2], best[2:4], best[4:6])]
    dd = lc*det(matrix(g))
    # (vi) the two integration segments by the cyclic-adjacency rule and the period minor M
    cutQ = best[1] == best[0] + 1 and best[3] == best[2] + 1
    seg = [(r[1], r[2]), (r[3], r[4])] if cutQ else [(r[0], r[1]), (r[2], r[3])]
    mp.dps = dps
    def p(gg, u, v):   # p_g(J) = 2 int_J t^g dt/sqrt|f|, with t = u + (v - u) sin^2(theta): the square-root
        others = [ri for ri in r if ri != u and ri != v]   # endpoint factors cancel against dt and the
        tt = lambda th: u + (v - u)*sin(th)**2            # quadrature converges to the working precision
        return 4*quad(lambda th: tt(th)**gg/sqrt(fabs(lc*fprod(tt(th) - ri for ri in others))), [0, pi/2])
    mm = p(0, *seg[0])*p(1, *seg[1]) - p(1, *seg[0])*p(0, *seg[1])
    # (vii) W_S = |Delta M| / (4 pi^2 w)
    return fabs(dd*mm)/(4*pi**2*w)

if __name__ == "__main__":
    print("WatsonW(17, 1, 5, 7) =", mp.nstr(WatsonW(17, 1, 5, 7), 20))