"""Power sum of character degrees of abelian wreath products -- numberdb.org/T661.

Run with SageMath (or a sufficiently complete passagemath installation):
    sage -pip install numberdb
    sage -python generate.py                 # verify every stored entry
    sage -python generate.py --fill          # write entries to this draft

Set NUMBERDB_API_KEY for authenticated draft reads and writes. --fill writes
values, not publication or review status. Download argument-selection.md and
checks.py alongside this file for the selection rationale and independent checks.

Every colored multipartition is counted, including multiplicities of equal
degrees. Degrees use integer hook products and exact quotients. At n=2 the
spectrum is 2q linear characters and q(q-1)/2 characters of degree 2, so large
base orders do not require explicit multipartition enumeration. For any abelian
base of order q the degree multiset agrees with that of the cyclic base.

The default slices are (1,3..20), (2,2..12), (3,2..10), and (m^3,2) for
m=2..40, each at s=9/4 and s=3. They are 77 pairs and 154 entries. The
complete domains are q>=1, n>=2, rational s>=2. See argument-selection.md.

Nonintegral powers use positive real Arb balls with 64 extra working bits.
At s=9/4, d^(9/4) is d^2*sqrt(sqrt(d)); cubic sums are exact integers.
The default selection retains at least 391 bits of relative accuracy.
checks.py measures the retained accuracy and independently encloses roots by
integer fourth-root isolation. No rounded input is used.
"""
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.combinat.partition import Partitions
from collections import Counter
from functools import lru_cache
from math import factorial, isqrt
import os
import sys

WORKING_GUARD = 64


def selected_pairs():
    for q, first, last in ((1, 3, 20), (2, 2, 12), (3, 2, 10)):
        for n in range(first, last + 1):
            yield q, n
    for m in range(2, 41):
        yield m**3, 2


@lru_cache(maxsize=None)
def hook_products(n):
    """Hook products of ordinary partitions, with multiplicities."""
    out = Counter()
    for partition in Partitions(n):
        product = 1
        for i, row in enumerate(partition):
            for j in range(row):
                leg = sum(partition[k] > j for k in range(i + 1, len(partition)))
                product *= row - j + leg
        out[product] += 1
    return out


@lru_cache(maxsize=None)
def degree_histogram(q, n):
    """Map each exact degree to its character multiplicity."""
    q, n = int(q), int(n)
    if q < 1 or n < 2:
        raise ValueError('q>=1 and n>=2 are required')
    if n == 2:
        out = {1: 2*q}
        if q > 1:
            out[2] = q*(q-1)//2
        return out
    # Convolve the hook-product distributions of the q ordered components.
    state = {(0, 1): 1}
    for color in range(q):
        nxt = Counter()
        for (size, denominator), multiplicity in state.items():
            for extra in range(n-size+1):
                for hooks, count in hook_products(extra).items():
                    nxt[size+extra, denominator*hooks] += multiplicity*count
        state = nxt
    out = Counter()
    numerator = factorial(n)
    for (size, denominator), multiplicity in state.items():
        if size == n:
            degree, remainder = divmod(numerator, denominator)
            assert remainder == 0 and degree > 0
            out[degree] += multiplicity
    return dict(sorted(out.items()))


def moment(q, n, s, digits=100):
    s = QQ(s)
    if s < 2:
        raise ValueError('s>=2 is required')
    spectrum = degree_histogram(int(q), int(n))
    if s.denominator() == 1:
        return sum((ZZ(count)*ZZ(d)**int(s) for d, count in spectrum.items()), ZZ(0))
    if s == QQ(9)/QQ(4):
        roots = {d: isqrt(isqrt(d)) for d in spectrum}
        if all(roots[d]**4 == d for d in spectrum):
            return sum((ZZ(count)*ZZ(d)**2*ZZ(roots[d]) for d, count in spectrum.items()), ZZ(0))
    field = RealBallField(numberdb.bits(digits, losing=WORKING_GUARD))
    if s == QQ(9)/QQ(4):
        value = sum((field(count)*field(d)**2*field(d).sqrt().sqrt()
                     for d, count in spectrum.items()), field(0))
    else:
        value = sum((field(count)*field(d)**field(s) for d, count in spectrum.items()), field(0))
    assert value.is_finite() and value > 0
    return value


class CharacterDegreePowerSum(numberdb.Generator):
    table = 'T661'
    parameters = ('q', 'n', 's')
    type = 'R'
    digits = 100
    rigour = 'proven'
    files = ('generate.py', 'checks.py', 'check.py', 'argument-selection.md',
             'numerical-evidence.json', 'skill-evidence.json',
             'program.py', 'program_check.py', 'flatten.py', 'degree-histograms.json',
             'sources-manifest.json')

    def enumerate(self):
        for q, n in selected_pairs():
            for s in ('9/4', '3'):
                yield {'q': str(q), 'n': str(n), 's': s}

    def value(self, params, digits):
        return moment(int(params['q']), int(params['n']), params['s'], digits)


if __name__ == '__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
        os.environ['NUMBERDB_API_KEY'] = sys.stdin.read().strip()
    generator = CharacterDegreePowerSum()
    if '--fill' in sys.argv:
        print(generator.publish(message='Fill the checked character-degree power sums'))
    else:
        report = generator.verify(sample=None)
        print(report)
        sys.exit(0 if report.ok else 1)