"""Chiral unitary Weingarten functions -- numberdb.org/T552.

Run with SageMath and the numberdb client installed:
    sage -pip install numberdb
    sage -python generate.py                 # verify all stored entries
    sage -python generate.py --fill-draft    # send entries using NUMBERDB_API_KEY

The fill operation writes entries to the existing draft; it does not offer,
review or publish the table. A private draft needs its author's key to read.
The campaign wrapper can instead set NUMBERDB_KEY_FROM_STDIN=1 and
NUMBERDB_PUBLISH=1 (the latter means fill entries, not make the table public).

Exact character expansion from Matsumoto, arXiv:1301.5401, section 4.1:
    Wg_AIII(a,b,mu) = sum_lambda s_lambda(1^a,(-1)^b) chi_lambda(mu)/C_lambda(a+b).
Terms have length(lambda)<=a+b. C_lambda(N)=product_(i,j in lambda)(N+j-i).
Characters use Murnaghan-Nakayama; Schur specializations use Jacobi-Trudi.
All division is in QQ. No rounded constants or third-party data are used.
The complete finite selection is described in argument-selection.md.
"""
import numberdb.sage as numberdb
from sage.rings.rational_field import QQ
from functools import lru_cache
from itertools import permutations
from math import prod
import os
import sys


@lru_cache(None)
def partitions(n, largest=None):
    if n == 0:
        return ((),)
    top = n if largest is None else min(n, largest)
    return tuple((r,) + tail for r in range(top, 0, -1)
                 for tail in partitions(n-r, r))


@lru_cache(None)
def character(lam, mu):
    """Remove connected border strips, excluding 2-by-2 squares."""
    if not mu:
        return int(not lam)
    r, rest = mu[0], mu[1:]
    cells = {(i, j) for i, row in enumerate(lam) for j in range(row)}
    answer = 0
    for nu in partitions(sum(lam)-r):
        if len(nu) > len(lam) or any(v > lam[i] for i, v in enumerate(nu)):
            continue
        strip = cells - {(i, j) for i, row in enumerate(nu) for j in range(row)}
        if len(strip) != r:
            continue
        if any({(i,j),(i+1,j),(i,j+1),(i+1,j+1)} <= strip for i,j in strip):
            continue
        seen = {next(iter(strip))}
        frontier = list(seen)
        while frontier:
            i,j = frontier.pop()
            for neighbor in ((i-1,j),(i+1,j),(i,j-1),(i,j+1)):
                if neighbor in strip and neighbor not in seen:
                    seen.add(neighbor)
                    frontier.append(neighbor)
        if seen == strip:
            answer += (-1)**(len({i for i,j in strip})-1) * character(nu, rest)
    return answer


@lru_cache(None)
def schur_signature(lam, a, b):
    """Jacobi-Trudi evaluated from product (1-t)^(-a)(1+t)^(-b)."""
    k = sum(lam)
    h = [1] + [0]*k
    for x in (1,)*a + (-1,)*b:
        for degree in range(1, k+1):
            h[degree] += x*h[degree-1]
    def complete(degree):
        return h[degree] if degree >= 0 else 0
    ell = len(lam)
    value = 0
    for perm in permutations(range(ell)):
        inversions = sum(perm[i] > perm[j] for i in range(ell) for j in range(i+1,ell))
        value += (-1)**inversions * prod(complete(lam[i]-i+perm[i]) for i in range(ell))
    return value


@lru_cache(None)
def weingarten(a, b, mu):
    if a < 0 or b < 0 or a+b < 1 or not mu or any(x<1 for x in mu) or tuple(sorted(mu,reverse=True)) != mu:
        raise ValueError('Use nonnegative multiplicities of positive total and a nonempty partition.')
    n = a+b
    result = QQ(0)
    for lam in partitions(sum(mu)):
        if len(lam) > n:
            continue
        content_product = prod(n+j-i for i,row in enumerate(lam) for j in range(row))
        result += QQ(schur_signature(lam,a,b)*character(lam,mu))/QQ(content_product)
    return result


class ChiralUnitaryWeingarten(numberdb.Generator):
    table = 'T552'
    parameters = ('a','b','mu')
    files = ('generate.py','checks.py','argument-selection.md')
    type = 'Q'
    rigour = 'exact'
    digits = 100

    def enumerate(self, max_dimension=8, max_degree=4):
        for a in range(1,max_dimension):
            for b in range(1,min(a,max_dimension-a)+1):
                for k in range(1,min(max_degree,a+b)+1):
                    for mu in partitions(k):
                        yield {'a':a,'b':b,'mu':','.join(map(str,mu))}

    def value(self, params, digits=100):
        a,b = int(params['a']),int(params['b'])
        if not a >= b >= 1:
            raise ValueError('The table uses a >= b >= 1.')
        mu=tuple(map(int,str(params['mu']).split(',')))
        return weingarten(a,b,mu)


if __name__ == '__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
        token=sys.stdin.read().strip()
        if '=' in token and token.split('=',1)[0] in ('NUMBERDB_KEY','NUMBERDB_API_KEY'):
            token=token.split('=',1)[1].strip().strip('\"\'')
        numberdb.configure(api_key=token)
    generator=ChiralUnitaryWeingarten()
    if '--fill-draft' in sys.argv or os.environ.get('NUMBERDB_PUBLISH') == '1':
        print(generator.publish(message='Fill exact AIII coefficients after independent permutation and moment checks'))
    else:
        report=generator.verify(sample=None)
        print(report)
        raise SystemExit(0 if report.ok else 1)