"""Symplectic Weingarten functions -- numberdb.org/T551.

Run with SageMath (or a compatible passagemath installation):
    sage -pip install numberdb
    sage -python generate.py             # verify every stored coefficient
    sage -python generate.py --fill      # write entries with NUMBERDB_API_KEY

The --fill operation fills this table; it does not make a draft public.
The default enumeration is N=1..6 and all partitions 1<=|mu|<=N.
N is rank; matrices have complex dimension 2N. Values are exact rationals.

Matsumoto arXiv:1301.5401v3, (2.8), gives the twisted spherical expansion.
At his even representatives (2.3), pi^lambda = omega^(lambda transpose).
Collins--Matsumoto arXiv:0903.5143, (3.5), expresses omega in the power-sum
coefficients of the integral Jack J polynomial with parameter 2.
All divisions take place in QQ. No numerical precision is used.
"""
import numberdb.sage as numberdb
from sage.rings.rational_field import QQ
from sage.symbolic.ring import SR  # initialize symbolic conversion used by Jack
import sage.rings.polynomial.laurent_polynomial_ring
from sage.combinat.partition import Partitions
from sage.combinat.sf.sf import SymmetricFunctions
from collections import Counter
from functools import lru_cache
from math import factorial, prod
import os
import sys


@lru_cache(maxsize=None)
def partitions(k):
    return tuple(sorted((tuple(int(v) for v in lam) for lam in Partitions(k)),
                        reverse=True))


def transpose(lam):
    return tuple(sum(row > j for row in lam) for j in range(lam[0]))


def representation_dimension(lam):
    """Hook-length formula, evaluated in the rational field."""
    hooks = prod(row-j + sum(other > j for other in lam[i+1:])
                 for i, row in enumerate(lam) for j in range(row))
    result = QQ(factorial(sum(lam))) / QQ(hooks)
    assert result.denominator() == 1
    return result


@lru_cache(maxsize=None)
def spherical_data(k):
    sf = SymmetricFunctions(QQ)
    p = sf.power()
    jack = sf.jack(t=QQ(2)).J()
    return tuple((lam, representation_dimension(tuple(r for r in lam for _ in (0, 1))),
                  {tuple(mu): QQ(c) for mu, c in p(jack[transpose(lam)])})
                 for lam in partitions(k))


def symplectic_weingarten(N, mu):
    N = int(N)
    mu = tuple(int(x) for x in mu)
    if N < 1 or not mu or min(mu) < 1 or mu != tuple(sorted(mu, reverse=True)):
        raise ValueError('N must be positive and mu a nonempty decreasing partition')
    k = sum(mu)
    z_mu = prod(r**m * factorial(m) for r, m in Counter(mu).items())
    total = QQ(0)
    for lam, dimension, coefficients in spherical_data(k):
        if len(lam) > N:
            continue
        denominator = prod(2*N-2*i+j for i, row in enumerate(lam) for j in range(row))
        assert denominator != 0
        total += dimension * coefficients.get(mu, QQ(0)) / QQ(denominator)
    return QQ(2**len(mu)*z_mu) * total / QQ(factorial(2*k))


class SymplecticWeingarten(numberdb.Generator):
    table = 'T551'
    parameters = ('N', 'mu')
    type = 'Q'
    rigour = 'exact'
    digits = 100

    def enumerate(self, rank_bound=6, degree_bound=6):
        for N in range(1, rank_bound+1):
            for k in range(1, min(N, degree_bound)+1):
                for mu in partitions(k):
                    yield {'N': str(N), 'mu': ','.join(map(str, mu))}

    def value(self, params, digits):
        return symplectic_weingarten(int(params['N']),
                                    tuple(map(int, str(params['mu']).split(','))))


if __name__ == '__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
        key = sys.stdin.read().strip()
        if '=' in key:
            key = key.split('=', 1)[1].strip().strip('\"\'')
        numberdb.configure(api_key=key)
    generator = SymplecticWeingarten()
    if '--fill' in sys.argv or os.environ.get('NUMBERDB_PUBLISH') == '1':
        print(generator.publish(message='Fill exact symplectic Weingarten coefficients after independent Gram and Haar-moment checks'))
    else:
        report = generator.verify(sample=None)
        print(report)
        sys.exit(0 if report.ok else 1)