"""Independent exact checks for Unitary Weingarten functions, numberdb.org/T549.

Run: sage -python checks.py
With --stored, read the authenticated draft instead of computing local values.
The campaign can call main(stored=True) after configure_key().

The independent computation uses permutations, rational Gaussian elimination
and the Gram inverse relation, with Python Fraction arithmetic. It shares no
partition, rim-hook, character or hook-length code with the generator. Cycle
types and the expected finite selection are inferred from all permutations.

A second independent fact is the single-entry Haar moment. A column of a Haar
unitary matrix is uniform on the complex unit sphere. Its first squared
coordinate has Beta(1,N-1) distribution for N>1. Thus the 2k-th absolute moment
is k!/(N(N+1)...(N+k-1)); N=1 gives one. Summing the integration formula with
all matrix indices equal gives k! times the sum over permutations of Wg.

Controls are run first: a nonzero known linear-system solution, cycle types,
the N=2,k=2 coefficients 1/3 and -1/6, and Haar moments 1/2 and 1/3 for N=2.
A deliberately corrupted coefficient must fail the inverse check.
"""
import numberdb.sage as numberdb
from sage.rings.rational_field import QQ
from fractions import Fraction as F
from itertools import permutations
from collections import Counter
from functools import lru_cache
from math import factorial, prod
import json
import sys
import generate
from check import exactness, measure, agrees_with, names_its_rings


def cycle_type(p):
    remaining = set(range(len(p)))
    lengths = []
    while remaining:
        start = min(remaining)
        length = 0
        j = start
        while j in remaining:
            remaining.remove(j)
            length += 1
            j = p[j]
        lengths.append(length)
    return tuple(sorted(lengths, reverse=True))


def inverse(p):
    q = [0] * len(p)
    for i, j in enumerate(p):
        q[j] = i
    return tuple(q)


def compose(p, q):
    return tuple(p[q[i]] for i in range(len(p)))


@lru_cache(None)
def group_data(k):
    ps = list(permutations(range(k)))
    by_type = {}
    for p in ps:
        by_type.setdefault(cycle_type(p), []).append(p)
    types = sorted(by_type, reverse=True)
    reps = [by_type[t][0] for t in types]
    positions = {t:i for i,t in enumerate(types)}
    # coeff[i][j][ell] counts tau with c(tau)=ell and type(tau^-1 sigma_i)=type_j.
    coeff = [[Counter() for _ in types] for _ in types]
    inverses = [(inverse(tau), len(cycle_type(tau))) for tau in ps]
    for i, sigma in enumerate(reps):
        for inv_tau, ell in inverses:
            j = positions[cycle_type(compose(inv_tau, sigma))]
            coeff[i][j][ell] += 1
    sizes = [len(by_type[t]) for t in types]
    return types, coeff, sizes


def gram_system(N, k):
    types, coeff, sizes = group_data(k)
    matrix = [[sum(count * N**ell for ell,count in cell.items())
               for cell in row] for row in coeff]
    rhs = [int(t == (1,)*k) for t in types]
    return types, matrix, rhs, sizes


def solve(matrix, rhs):
    a = [[F(x) for x in row] + [F(y)] for row, y in zip(matrix, rhs)]
    n = len(a)
    for col in range(n):
        pivot = next(i for i in range(col,n) if a[i][col])
        a[col], a[pivot] = a[pivot], a[col]
        scale = a[col][col]
        a[col] = [x / scale for x in a[col]]
        for i in range(n):
            if i == col:
                continue
            scale = a[i][col]
            a[i] = [x - scale*y for x,y in zip(a[i], a[col])]
    return [row[-1] for row in a]


def is_inverse(matrix, values, rhs):
    return all(sum(a*x for a,x in zip(row, values)) == y
               for row,y in zip(matrix,rhs))


def controls():
    assert solve([[2,1],[1,3]], [1,0]) == [F(3,5), F(-1,5)]
    assert cycle_type((1,2,0,3)) == (3,1)
    p = (2,0,1,3)
    assert compose(inverse(p), p) == (0,1,2,3)
    types, a, b, sizes = gram_system(2,2)
    v = solve(a,b)
    assert dict(zip(types,v)) == {(2,): F(-1,6), (1,1): F(1,3)}
    assert is_inverse(a,v,b)
    v[0] += 1
    assert not is_inverse(a,v,b), 'Corruption control failed'
    assert F(factorial(1), prod(range(2,3))) == F(1,2)
    assert F(factorial(2), prod(range(2,4))) == F(1,3)


def flat_entries(document):
    # Follow the API's stored structure, independent of generator enumeration.
    block = document['Numbers']
    if isinstance(block, list):
        for row in block:
            yield row['params'], row['number']
    else:
        for N, inner in block.items():
            if isinstance(inner, dict) and 'numbers' in inner:
                inner = inner['numbers']
            for mu, entry in inner.items():
                yield {'N':N, 'mu':mu}, entry['number'] if isinstance(entry,dict) else entry


def main(stored=False):
    controls()
    generator = generate.UnitaryWeingarten()
    if stored:
        entries = list(flat_entries(numberdb.table('T549')))
    else:
        entries = [(params,generator.value(params,100)) for params in generator.enumerate()]
    values = {}
    for p, x in entries:
        rational = F(str(x))
        key = (int(p['N']), tuple(map(int, p['mu'].split(','))))
        values[key] = QQ(rational.numerator) / QQ(rational.denominator)
    assert len(values) == len(entries), 'duplicate parameters'
    assert not exactness(values)
    assert not names_its_rings(generate.__file__)
    expected = {}
    order = []
    inverse_count = 0
    moments = 0
    for N in range(1,13):
        for k in range(1,min(6,N)+1):
            types,a,b,sizes = gram_system(N,k)
            independent = solve(a,b)
            for mu, value in zip(types,independent):
                key = (N,mu)
                expected[key] = QQ(value.numerator) / QQ(value.denominator)
                order.append(key)
            observed = [F(str(values[(N,mu)])) for mu in types]
            assert is_inverse(a,observed,b), ('Gram inverse failure',N,k)
            inverse_count += len(types)
            computed_moment = factorial(k) * sum(size*x for size,x in zip(sizes,observed))
            assert computed_moment == F(factorial(k),prod(range(N,N+k))), ('Haar moment',N,k)
            moments += 1
    assert list(values) == order, 'selection or ordering mismatch'
    assert not agrees_with(values, expected.__getitem__)
    # Published low-degree formulas (Matsumoto Example 2.1), all admissible N.
    low_degree = 0
    for N in range(1,13):
        assert values[N,(1,)] == QQ(1)/QQ(N)
        low_degree += 1
        if N >= 2:
            assert values[N,(1,1)] == QQ(1)/QQ(N*N-1)
            assert values[N,(2,)] == -QQ(1)/QQ(N*(N*N-1))
            low_degree += 2
    from program import unitary_weingarten
    library_checks = 0
    for (N, mu), value in values.items():
        assert unitary_weingarten(N, mu) == value, ('Sage character expansion', N, mu)
        library_checks += 1
    assert unitary_weingarten(13, (3, 1)) == QQ(67)/QQ(149909760)
    stats = measure(values)
    stats['longest_at'] = str(stats['longest_at'])
    result = {'table':'T549','data_source':'stored' if stored else 'local',
              'controls_passed':True, 'independent_gram_solutions_compared':len(expected),
              'gram_inverse_equations_checked':inverse_count, 'haar_moments_checked':moments,
              'published_low_degree_values_checked':low_degree,
              'sage_symmetric_function_values_checked':library_checks,
              'program_extension_example_checked':True,
              'exactness_complaints':[], 'enumeration_and_order_match':True,
              'measure':stats,
              'sample_values': {str(k):str(values[k]) for k in [(2,(2,)),(3,(2,1)),(3,(3,)),(12,(6,))]}}
    print(json.dumps(result, indent=2))
    return result


if __name__ == '__main__':
    generate.configure_key()
    main(stored='--stored' in sys.argv)