"""Genus-character twisted traces of singular moduli -- numberdb.org/T671.

Run with SageMath (install numberdb once with sage -pip install numberdb):
    sage -python generate.py                 # verify all stored entries
    sage -python generate.py --fill          # fill the existing private draft
    sage -python checks.py                   # offline independent checks

The trace uses all reduced positive definite integral forms, including
imprimitive forms. J_m has principal part q^(-m), with no constant term.
Jenkins, Theorem 5.1, identifies t_m(D,d)/sqrt(D) with an integer coefficient,
for positive fundamental D and every negative order discriminant -d.
We isolate that integer using rigorous Arb CM sums; we do not infer
integrality from near-integer decimal digits. Exact zero is retained.

The default selection is all D*d <= 300 and 1 <= m <= 4, in both invertible
normalizations. See argument-selection.md for each parameter's rationale.
The 768-bit working precision leaves over 100 absolute decimal places in
all coefficient-isolation balls at the default range. Every call checks
finiteness and a radius below 10^(-110), and raises instead of guessing.
Returned irrational traces use numberdb.bits(digits, losing=96) bits, so
rounding error is carried in a ball throughout. The attached checks measure
the actual accuracy, test known answers, and use independent Hecke sums.
"""
import os
import sys
from functools import lru_cache
from math import gcd, isqrt
import numberdb.sage as numberdb
from sage.rings.integer_ring import ZZ
from sage.rings.rational_field import QQ
from sage.rings.real_arb import RealBallField
from sage.rings.complex_arb import ComplexBallField

WORKING_BITS = 768
OUTPUT_GUARD_BITS = 96


def pairs(bound=300):
    for D in range(5, bound // 3 + 1):
        if not ZZ(D).is_fundamental_discriminant():
            continue
        for d in range(3, bound // D + 1):
            if d % 4 in (0, 3):
                yield D, d


@lru_cache(None)
def reduced_forms(N):
    """Unique PSL2(Z) representatives of discriminant -N, no primitivity filter."""
    out = []
    for a in range(1, isqrt(N // 3) + 1):
        for b in range(-a, a + 1):
            if (b*b + N) % (4*a):
                continue
            c = (b*b + N) // (4*a)
            if c < a or (b < 0 and (abs(b) == a or a == c)):
                continue
            out.append((a, b, c))
    return tuple(out)


def weight(Q):
    a,b,c = Q
    return 3 if a == b == c else 2 if a == c and b == 0 else 1


@lru_cache(None)
def character(D, Q):
    a,b,c = Q
    if gcd(gcd(gcd(a,b),c),D) > 1:
        return ZZ(0)
    for n in (a,c,a+b+c):
        if gcd(n,D) == 1:
            return ZZ(D).kronecker(n)
    # Residue representatives modulo D suffice, by the local definition and CRT.
    for x in range(D):
        for y in range(D):
            n = a*x*x+b*x*y+c*y*y
            if n > 0 and gcd(n,D) == 1:
                return ZZ(D).kronecker(n)
    raise ArithmeticError('No represented integer coprime to D')


def faber(m, J):
    if m == 1: return J
    if m == 2: return J*J - 393768
    if m == 3: return J**3 - 590652*J - 64481280
    if m == 4: return J**4 - 787536*J**2 - 85975040*J + 74069419032
    raise ValueError('This generator implements the first four Hecke orders')


@lru_cache(None)
def cm_j(N, Q, bits=WORKING_BITS):
    a,b,c = Q
    C = ComplexBallField(bits)
    tau = (C(-b) + C(0,1)*C(N).sqrt())/(2*a)
    value = tau.modular_j()
    if not (value.real().is_finite() and value.imag().is_finite()):
        raise ArithmeticError('Nonfinite CM value')
    return value


@lru_cache(None)
def coefficient_ball(m,D,d,bits=WORKING_BITS):
    C = ComplexBallField(bits)
    total = C(0)
    for Q in reduced_forms(D*d):
        chi = character(D,Q)
        if chi:
            total += chi*faber(m,cm_j(D*d,Q,bits)-744)/weight(Q)
    result = total/C(D).sqrt()
    if not (result.real().is_finite() and result.imag().is_finite()) or not result.imag().contains_zero():
        raise ArithmeticError('Trace is not enclosed as a finite real')
    limit = QQ(1)/ZZ(10)**110
    if result.real().rad() >= limit or result.imag().rad() >= limit:
        raise ArithmeticError('CM enclosure is too wide for certification')
    return result.real()


@lru_cache(None)
def coefficient(m,D,d):
    b = coefficient_ball(m,D,d)
    n = ZZ(b.mid().round())
    # The theorem supplies integrality; this interval identifies its unique integer.
    if n not in b or n-1 in b or n+1 in b:
        raise ArithmeticError('Normalized trace does not isolate one integer')
    return n


class TwistedTraces(numberdb.Generator):
    table = 'T671'
    parameters = ('m','D','d','normalisation')
    type = 'R'
    digits = 100
    rigour = 'proven'
    files = ('generate.py','checks.py','argument-selection.md','numerical-evidence.md',
             'skill-evidence.json','companions.json')

    def enumerate(self, bound=300, max_m=4):
        if max_m > 4:
            raise ValueError('Extend the Faber implementation before raising max_m')
        for m in range(1,max_m+1):
            for D,d in pairs(bound):
                for norm in ('trace','coefficient'):
                    yield dict(m=m,D=D,d=d,normalisation=norm)

    def value(self,params,digits):
        m,D,d = (int(params[k]) for k in ('m','D','d'))
        if m < 1 or D <= 1 or not ZZ(D).is_fundamental_discriminant() or d < 3 or d%4 not in (0,3):
            raise ValueError('Arguments are outside the mathematical domain')
        A = coefficient(m,D,d)
        norm = params['normalisation']
        if norm == 'coefficient' or A == 0:
            value = A
        elif norm == 'trace':
            R = RealBallField(numberdb.bits(digits,losing=OUTPUT_GUARD_BITS))
            value = A*R(D).sqrt()
        else:
            raise ValueError('Unknown normalization')
        return {'number':value,'param-latex':'$t_m$' if norm=='trace' else '$A_m$'}


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].isidentifier():
            key = key.split('=',1)[1].strip().strip('\"\'')
        os.environ['NUMBERDB_API_KEY'] = key
    gen = TwistedTraces()
    if '--fill' in sys.argv or os.environ.get('NUMBERDB_PUBLISH') == '1':
        print(gen.publish(message='Add certified genus-character traces and exact coefficients'))
    else:
        report = gen.verify(sample=None)
        print(report)
        sys.exit(0 if report.ok else 1)