"""Coexact one-form Laplacian determinants on lens spaces -- numberdb.org/T631.

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

The radius is one, eigenvalues are m^2 for m >= 2, and both helicities
are counted using complex dimensions. Exact weight counts prove a quadratic
quasipolynomial on residue classes modulo 2p. See continuation.md for the
proof, and argument-selection.md for the enumeration. Hurwitz zeta values
and derivatives use Arb series arithmetic; no numerical fitting is used.
Working precision is numberdb.bits(digits, losing=180). checks.py measures
the achieved accuracy and tests independently computed character traces.
"""
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 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 multiplicity(p,q,u,m):
    if m < 2:
        return ZZ(0)
    period=2*p
    m0=2+(m-2)%period
    _,c0,c1,c2=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 m0 in range(2,period+2):
        a=QQ(m0)/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()
            out[m0,j]=(z[0].real(),z[1].real())
    return out


@lru_cache(None)
def 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 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


class CoexactLensDeterminants(numberdb.Generator):
    table='T631'
    parameters=('p','q','u','convention')
    type='R'
    digits=100
    rigour='proven'
    files=('generate.py','checks.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 {'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)
        return L.exp() if params['convention']=='det' else 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('\"\'')
        os.environ['NUMBERDB_API_KEY']=key
    g=CoexactLensDeterminants()
    if '--fill' in sys.argv or os.environ.get('NUMBERDB_PUBLISH')=='1':
        print(g.publish(message='Fill coexact lens-space determinant draft with certified 100-digit values',assisted_by='codex-cli'))
    else:
        report=g.verify(sample=None)
        print(report)
        sys.exit(0 if report.ok else 1)