"""One-form Laplacian determinants on lens spaces -- numberdb.org/T633.

Run with SageMath:
    sage -pip install numberdb
    sage -python generate.py          # verify every stored entry
    sage -python generate.py --fill   # fill the draft with NUMBERDB_API_KEY

The full one-form spectrum is the disjoint union of exact eigenvalues
m^2-1 and coexact eigenvalues m^2, m>=2, on the radius-one quotient.
The scalar generator of T632 and coexact generator of T631 are combined here.
No companion decimals are used as numerical inputs. Exact SU(2) weight counts
give quadratic multiplicity polynomials on residue
classes modulo 2p. The shifted factors m-1 and m+1 have no multiplicative
anomaly here, because the multiplicity Dirichlet series has no positive
even pole. See continuation.md for the proof. Finite Hurwitz zeta sums and
first derivatives are evaluated with Arb enclosures, then exponentiated.
The 180 extra working bits are measured by checks.py across the full range.
"""
import os
import sys
from collections import Counter
from functools import lru_cache
from math import gcd
import numberdb.sage as numberdb
from sage.rings.integer_ring import ZZ
from sage.rings.rational_field import QQ
from sage.rings.complex_arb import ComplexBallField
from sage.rings.real_arb import RealBallField
from sage.rings.polynomial.polynomial_ring_constructor import PolynomialRing

WORKING_GUARD = 180


def presentations(limit=20):
    for p in range(2, limit + 1):
        for q in range(1, p):
            if gcd(p, q) == 1:
                r = pow(q, -1, p)
                if q == min(q, p-q, r, p-r):
                    yield p, q


def cases(limit=20):
    for p, q in presentations(limit):
        for u in range(p//2 + 1):
            yield p, q, u


@lru_cache(None)
def scalar_coefficients(p, q, u):
    """Exact (m0,c0,c1,c2) with d(m0+2p*k)=c0+c1*k+c2*k^2."""
    assert p >= 2 and 0 < q < p and gcd(p, q) == 1 and 0 <= u < p
    period = 2*p
    rows = []
    for m0 in range(2, period+2):
        weights = Counter(a % period for a in range(1-m0, m0, 2))
        c0 = c1 = c2 = 0
        residues = range((m0-1) % 2, period, 2)
        for a in residues:
            for b in residues:
                exponent = a*(1+q) + b*(1-q)
                assert exponent % 2 == 0
                if (exponent//2-u) % p == 0:
                    c0 += weights[a]*weights[b]
                    c1 += 2*(weights[a]+weights[b])
                    c2 += 4
        rows.append((m0, ZZ(c0), ZZ(c1), ZZ(c2)))
    # Residue of F(z)=sum d(m)m^(-z) at z=2 vanishes exactly.
    assert sum(c1*period-2*c2*m0 for m0,c0,c1,c2 in rows) == 0
    return tuple(rows)


def scalar_multiplicity(p, q, u, m):
    if m == 1:
        return ZZ(u == 0)
    assert m >= 2
    period = 2*p
    m0 = 2+(m-2) % period
    _, c0, c1, c2 = scalar_coefficients(p,q,u)[m0-2]
    k = (m-m0)//period
    return c0+c1*k+c2*k*k


@lru_cache(None)
def hurwitz_data(p, precision):
    C = ComplexBallField(precision)
    P = PolynomialRing(C, 's')
    s = P.gen()
    period = 2*p
    out = {}
    for numerator in range(1, period+3):
        a = QQ(numerator)/period
        for j in range(3):
            z = (s-j)._zeta_series(2, a=C(a))
            assert z[0].imag().contains_zero() and z[1].imag().contains_zero()
            assert all(v.real().is_finite() and v.imag().is_finite() for v in (z[0],z[1]))
            out[numerator,j] = z[0].real(), z[1].real()
    return out


@lru_cache(None)
def scalar_logdet(p, q, u, digits=100):
    precision = numberdb.bits(digits, losing=WORKING_GUARD)
    R = RealBallField(precision)
    period = 2*p
    data = hurwitz_data(p, precision)
    lp = R(period).log()
    answer = R(0)
    for m0, c0, c1, c2 in scalar_coefficients(p,q,u):
        for shift in (-1, 1):
            a = QQ(m0+shift)/period
            alpha = (c0-c1*a+c2*a*a, c1-2*c2*a, c2)
            for j in range(3):
                z, dz = data[m0+shift,j]
                answer += alpha[j]*(lp*z-dz)
    assert answer.is_finite()
    return answer


@lru_cache(None)
def coexact_coefficients(p, q, u):
    """(m0, c0,c1,c2) with d(m0+2p*k)=c0+c1*k+c2*k^2, k>=0.

    Weight exponents are [a(1+q)+b(1-q)]/2. Both weights have the
    same parity, so the division by 2 is exact. Adding 2p to m adds
    two copies of each admissible residue to each SU(2) weight string.
    """
    assert p >= 2 and 0 < q < p and gcd(p,q)==1 and 0 <= u < p
    period = 2*p
    rows = []
    for m0 in range(2, period+2):
        large = Counter(a % period for a in range(-m0,m0+1,2))
        small = Counter(b % period for b in range(2-m0,m0-1,2))
        c0=c1=c2=0
        residues=range(m0 % 2,period,2)
        for a in residues:
            for b in residues:
                e=a*(1+q)+b*(1-q)
                assert e % 2 == 0
                if (e//2-u) % p == 0:
                    c0 += large[a]*small[b]+small[a]*large[b]
                    c1 += 2*(large[a]+small[b]+small[a]+large[b])
                    c2 += 8
        rows.append((m0,ZZ(c0),ZZ(c1),ZZ(c2)))
    return tuple(rows)


def coexact_multiplicity(p,q,u,m):
    if m < 2:
        return ZZ(0)
    period=2*p
    m0=2+(m-2)%period
    _,c0,c1,c2=coexact_coefficients(p,q,u)[m0-2]
    k=(m-m0)//period
    return c0+c1*k+c2*k*k



@lru_cache(None)
def coexact_logdet(p, q, u, digits=100):
    precision = numberdb.bits(digits, losing=WORKING_GUARD)
    R = RealBallField(precision)
    period = 2*p
    data = hurwitz_data(p, precision)
    lp = R(period).log()
    answer = R(0)
    for m0, c0, c1, c2 in coexact_coefficients(p,q,u):
        a = QQ(m0)/period
        alpha = (c0-c1*a+c2*a*a, c1-2*c2*a, c2)
        for j in range(3):
            z, dz = data[m0,j]
            answer += 2*alpha[j]*(lp*z-dz)
    assert answer.is_finite()
    return answer


@lru_cache(None)
def logdet(p, q, u, digits=100):
    """Full one-form log determinant from the two disjoint spectra."""
    return scalar_logdet(p,q,u,digits)+coexact_logdet(p,q,u,digits)


class OneFormLensDeterminants(numberdb.Generator):
    table = 'T633'
    parameters = ('p','q','u','convention')
    type = 'R'
    digits = 100
    rigour = 'proven'
    files = ('generate.py','checks.py','independent.py','check.py',
             'argument-selection.md','continuation.md','numerical-evidence.json',
             'skill-evidence.json','source-evidence.json')

    def enumerate(self, limit=20):
        for p,q,u in cases(limit):
            for convention in ('det','logdet'):
                yield dict(p=str(p),q=str(q),u=str(u),convention=convention)

    def value(self, params, digits):
        p,q,u = (int(params[k]) for k in ('p','q','u'))
        L = logdet(p,q,u,digits)
        if params['convention'] == 'det':
            return L.exp()
        assert params['convention'] == 'logdet'
        return L


if __name__ == '__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
        key = sys.stdin.read().strip()
        if '=' in key and key.split('=',1)[0].replace('_','').isalnum():
            key = key.split('=',1)[1].strip().strip("\"'")
        numberdb.configure(api_key=key)
    generator = OneFormLensDeterminants()
    if '--fill' in sys.argv or os.environ.get('NUMBERDB_PUBLISH') == '1':
        print(generator.publish(message='Fill full one-form lens-space draft with certified 100-digit values'))
    else:
        report = generator.verify(sample=None)
        print(report)
        sys.exit(0 if report.ok else 1)
