"""K3 elliptic genus coefficients -- numberdb.org/T667.

Run with SageMath (or a compatible passagemath installation):
    sage -pip install numberdb
    sage -python generate.py                 # verify every stored entry
    sage -python generate.py --publish       # fill an authorized draft

The latter command sends entries; it does not publish a website draft.
Private draft verification also requires NUMBERDB_API_KEY. For the campaign
wrapper, NUMBERDB_KEY_FROM_STDIN=1 reads the key from standard input.

Exact finite series implement Kawai, hep-th/9710016, equation (4.33).
Let a=sum_{j in Z} q^(j*j), b=2*sum_{j>=0} q^(j*(j+1)), and
P=prod_{m>=1}(1-q^m)^(-6). Then
  sum_{n>=0} c(4*n) q^n = (10*a^4*b - 2*q*b^5)*P,
  sum_{n>=0} c(4*n-1) q^n = (2*a^5 - 10*q*a*b^4)*P.
All operations are additions and multiplications of Sage integers. Products
are truncated only after terms that cannot contribute to the requested grade.
The bound is reference coverage, not part of the infinite family's definition.
See argument-selection.md and the independently implemented checks.py.
"""
import os
import sys
from functools import lru_cache
import numberdb.sage as numberdb
from sage.rings.integer_ring import ZZ

MAX_D = 1000


def multiply(a, b, n):
    out = [ZZ(0) for _ in range(n + 1)]
    for i, x in enumerate(a):
        if x:
            for j, y in enumerate(b[:n + 1 - i]):
                if y:
                    out[i + j] += x * y
    return out


def power(a, exponent, n):
    result = [ZZ(1)] + [ZZ(0)] * n
    for _ in range(exponent):
        result = multiply(result, a, n)
    return result


@lru_cache(maxsize=None)
def coefficient_sequences(n):
    if n < 0:
        raise ValueError('The Fourier grade must be nonnegative.')
    a = [ZZ(0) for _ in range(n + 1)]
    a[0] = ZZ(1)
    k = 1
    while k * k <= n:
        a[k * k] = ZZ(2)
        k += 1
    b = [ZZ(0) for _ in range(n + 1)]
    k = 0
    while k * (k + 1) <= n:
        b[k * (k + 1)] = ZZ(2)
        k += 1
    p = [ZZ(1)] + [ZZ(0)] * n
    for m in range(1, n + 1):
        for _ in range(6):
            for j in range(m, n + 1):
                p[j] += p[j - m]
    a4 = power(a, 4, n)
    b4 = power(b, 4, n)
    a4b = multiply(a4, b, n)
    b5 = multiply(b4, b, n)
    a5 = multiply(a4, a, n)
    ab4 = multiply(a, b4, n)
    even = [10 * a4b[j] - (2 * b5[j - 1] if j else 0)
            for j in range(n + 1)]
    odd = [2 * a5[j] - (10 * ab4[j - 1] if j else 0)
           for j in range(n + 1)]
    return tuple(multiply(even, p, n)), tuple(multiply(odd, p, n))


class K3EllipticGenus(numberdb.Generator):
    table = 'T667'
    parameters = ('D',)
    type = 'Z'
    rigour = 'exact'
    digits = 100  # ignored for exact integers; no rounding is performed
    files = ('generate.py', 'checks.py', 'argument-selection.md',
             'numerical-evidence.json', 'skill-evidence.json')

    def enumerate(self, max_d=MAX_D):
        for d in range(-1, max_d + 1):
            if d % 4 in (0, 3):
                yield {'D': str(d)}

    def value(self, params, digits):
        d = int(params['D'])
        if d < -1 or d % 4 not in (0, 3):
            raise ValueError('D must be at least -1 and congruent to 0 or 3 mod 4.')
        grade = (d + (d % 4 == 3)) // 4
        even, odd = coefficient_sequences(max(MAX_D // 4, grade))
        return (odd if d % 4 == 3 else even)[grade]


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('\"\'')
        numberdb.configure(api_key=key)
    generator = K3EllipticGenus()
    if '--publish' in sys.argv:
        print(generator.publish(message='Fill the K3 elliptic genus with exact coefficients.'))
    else:
        report = generator.verify(sample=None)
        print(report)
        sys.exit(0 if report.ok else 1)