"""Independent checks for Chow polynomials of matroids, numberdb.org/T581.

    sage -python checks.py
    NUMBERDB_CHECK_STORED=1 sage -python checks.py

The checker computes the characteristic convolution (FMSV Theorem 1.4),
not the generator's weighted-chain recursion (Proposition 3.5). Uniform
restrictions use Boolean characteristic polynomials. Braid restrictions
come from explicitly enumerated set partitions. Named objects have every
subset ranked and every flat found by the rank-increase test, without
calling the generator's lattice functions or Sage's flats(). Subset sums
compute characteristic polynomials of every restriction of a contraction.
Uniform entries also agree with the derangement formula (Theorem 3.21),
with permutations counted directly. Known-answer controls run first.
"""
import os
import json
import re
from functools import lru_cache
from itertools import permutations
from math import comb
from time import perf_counter
import numberdb.sage as numberdb
from sage.rings.integer_ring import ZZ
from sage.rings.rational_field import QQ
from sage.rings.polynomial.polynomial_ring_constructor import PolynomialRing
from sage.matroids import catalog

P = PolynomialRing(ZZ, 't')
z = P.gen()


def reduced(poly):
    quotient, remainder = P(poly).quo_rem(z-1)
    assert remainder == 0
    return P(quotient)


def uniform_characteristic(r, n):
    return sum(((-1)**j * ZZ(comb(n,j)) * z**(r-j) for j in range(r)), P.zero()) + (-1)**r * ZZ(comb(n-1,r-1))

@lru_cache(None)
def uniform_check(r, n):
    if r == 0:
        assert n == 0
        return P.one()
    answer = reduced(uniform_characteristic(r,n))
    for j in range(1,r):
        answer += ZZ(comb(n,j)) * (z-1)**(j-1) * uniform_check(r-j,n-j)
    return P(answer)


def partitions(n):
    """Set partitions by adding the next label to a block or a new block."""
    if n == 0:
        yield ()
        return
    for blocks in partitions(n-1):
        for i in range(len(blocks)):
            yield blocks[:i]+(blocks[i]+(n-1,),)+blocks[i+1:]
        yield blocks+((n-1,),)

@lru_cache(None)
def braid_check(n):
    if n == 1:
        return P.one()
    answer = P.zero()
    for blocks in partitions(n):
        if len(blocks) == n:
            continue
        chi = P.one()
        for block in blocks:
            for i in range(1,len(block)):
                chi *= z-i
        answer += reduced(chi)*braid_check(len(blocks))
    return answer

@lru_cache(None)
def subset_lattice(name):
    M = getattr(catalog,name)()
    E = sorted(M.groundset(),key=repr)
    n = len(E)
    ranks = [int(M.rank([E[i] for i in range(n) if mask & (1 << i)]))
             for mask in range(1 << n)]
    flats = [mask for mask in range(1 << n)
             if all(ranks[mask | (1 << i)] > ranks[mask]
                    for i in range(n) if not mask & (1 << i))]
    assert 0 in flats
    return ranks, sorted(flats, key=lambda mask:(ranks[mask],mask),reverse=True), n

@lru_cache(None)
def named_check(name):
    ranks, flats, n = subset_lattice(name)
    H = {}
    for A in flats:
        if A == (1 << n)-1:
            H[A] = P.one()
            continue
        H[A] = P.zero()
        for B in flats:
            if A == B or A & B != A:
                continue
            difference = B ^ A
            coefs = [0]*(ranks[B]-ranks[A]+1)
            S = difference
            while True:
                coefs[ranks[B]-ranks[A | S]] += (-1)**S.bit_count()
                if S == 0:
                    break
                S = (S-1) & difference
            H[A] += reduced(P(coefs))*H[B]
    return H[0]

@lru_cache(None)
def derangement_poly(n):
    coefs = [0]*max(1,n)
    for perm in permutations(range(n)):
        if all(perm[i] != i for i in range(n)):
            coefs[sum(perm[i] > i for i in range(n))] += 1
    return P(coefs)


def derangement_check(r,n):
    return sum((ZZ(comb(n,j))*derangement_poly(j)*P([1]*(r-j))
                for j in range(r)), P.zero())


def independent_value(key):
    if key.startswith('U('):
        r,n=map(int,re.fullmatch(r'U\((\d+),(\d+)\)',key).groups())
        return uniform_check(r,n)
    if key.startswith('K('):
        return braid_check(int(key[2:-1]))
    return named_check(key)


def controls():
    assert uniform_check(0,0) == 1
    assert uniform_check(1,3) == 1
    assert uniform_check(2,5) == 1+z
    assert uniform_check(3,3) == 1+4*z+z**2
    assert uniform_check(3,4) == 1+7*z+z**2
    assert braid_check(4) == 1+8*z+z**2
    assert named_check('Fano') == 1+8*z+z**2
    assert derangement_poly(0) == 1
    assert derangement_poly(1) == 0
    assert derangement_poly(4) == z+7*z**2+z**3
    # Different family: descent enumerators of permutations for Boolean matroids.
    for n in range(1,8):
        counts=[0]*n
        for perm in permutations(range(n)):
            counts[sum(perm[i] > perm[i+1] for i in range(n-1))] += 1
        assert uniform_check(n,n) == P(counts)
    assert len(list(partitions(8))) == 4140
    print('Controls: 10 known answers, 7 Boolean/Eulerian comparisons, Bell(8)=4140 passed',flush=True)


def check_values(values):
    """Usable on generated or freshly read-back values; no generator computations."""
    controls()
    from check import exactness, agrees_with, measure
    assert not exactness(values)
    disagreements=agrees_with(values,independent_value)
    assert not disagreements, disagreements
    n_uniform=0
    for key,pol in values.items():
        pol=P(pol)
        if key.startswith('U('):
            r,n=map(int,re.fullmatch(r'U\((\d+),(\d+)\)',key).groups())
            assert pol == derangement_check(r,n)
            flat_count=1+sum(comb(n,j) for j in range(r))
            n_uniform+=1
        elif key.startswith('K('):
            m=int(key[2:-1]);r=m-1;n=comb(m,2)
            flat_count=len(list(partitions(m)))
        else:
            ranks,flats,n=subset_lattice(key);r=max(ranks);flat_count=len(flats)
        assert pol.degree()==r-1 and pol[0]==1 and pol[r-1]==1
        assert list(pol)==list(reversed(list(pol)))
        assert all(c>=0 for c in pol)
        assert all(pol[i]<=pol[i+1] for i in range((r-1)//2))
        assert pol[1] == flat_count-n-1
        print(key,str(pol),flush=True)
    print('Full-coefficient characteristic-convolution comparisons:',len(values))
    print('Independent counted-derangement comparisons:',n_uniform)
    print('Degree, endpoints, positivity, symmetry, unimodality and linear coefficient:',len(values))
    print('Measure:',json.dumps(measure(values)))
    return {k:str(v) for k,v in values.items()}


def read_stored():
    from generate import authenticate_from_stdin
    authenticate_from_stdin()
    doc=numberdb.table('T581')
    numbers=doc['Numbers']
    if isinstance(numbers,dict):
        return {str(k):P(v['number']) for k,v in numbers.items()}
    return {str(v['params']['matroid']):P(v['number']) for v in numbers}


if __name__ == '__main__':
    started=perf_counter()
    if os.environ.get('NUMBERDB_CHECK_STORED') == '1':
        values=read_stored()
    else:
        from generate import MatroidChow
        gen=MatroidChow()
        values={p['matroid']:gen.value(p,100) for p in gen.enumerate()}
    check_values(values)
    print('Elapsed seconds:',round(perf_counter()-started,3))
