"""Independent checks for Augmented Chow polynomials -- numberdb.org/T582.

Run: sage -python checks.py
For stored values: NUMBERDB_CHECK_STORED=1 sage -python checks.py
Private drafts require NUMBERDB_API_KEY or NUMBERDB_KEY_FROM_STDIN=1.

The generator uses positive weighted-chain sums. This checker instead uses
FMSV (2024), Theorem 3.14, on a separately enumerated subset-rank lattice
for named matroids; permutation cycles give signed Stirling coefficients for
braid matroids. Uniform cases use Theorem 3.21 with Eulerian coefficients
counted directly from permutations. Published Table 1(b,d) is transcribed
as a further external control. No generator arithmetic is reused.
"""
import numberdb.sage as numberdb
from sage.rings.integer_ring import ZZ
from sage.rings.polynomial.polynomial_ring_constructor import PolynomialRing
from sage.matroids import database_matroids as catalog
from sage.env import SAGE_VERSION
from itertools import permutations
from functools import lru_cache
from math import comb
import hashlib
import json
import os
from pathlib import Path
import sys
import time
from check import exactness, measure, agrees_with, names_its_rings
from generate import (MatroidAugmentedChow, uniform_ordinary, braid_ordinary,
                      named_intervals)

Q = PolynomialRing(ZZ, 't')
x = Q.gen()


def strip(value):
    return value['number'] if isinstance(value, dict) else value


@lru_cache(None)
def eulerian(j):
    counts = [0] * max(1, j)
    for p in permutations(range(j)):
        counts[sum(a > b for a, b in zip(p, p[1:]))] += 1
    return Q(counts)


def uniform_check(r, n):
    return Q.one() + sum((comb(n, j) * eulerian(j)
                         * Q([0] + [1] * (r - j)) for j in range(r)), Q.zero())


@lru_cache(None)
def cycle_counts(m):
    counts = [0] * (m + 1)
    for p in permutations(range(m)):
        seen = set(); cycles = 0
        for a in range(m):
            if a in seen:
                continue
            cycles += 1
            while a not in seen:
                seen.add(a); a = p[a]
        counts[cycles] += 1
    return counts


@lru_cache(None)
def braid_check(m):
    if m == 1:
        return Q.one()
    counts = cycle_counts(m)
    return -sum(((-1)**(m-k) * counts[k] * Q([1] * (m-k+1))
                 * braid_check(k) for k in range(1, m)), Q.zero())


@lru_cache(None)
def named_check(name):
    M = getattr(catalog, name)()
    elements = sorted(M.groundset(), key=str)
    n = len(elements)
    ranks = [int(M.rank([elements[j] for j in range(n) if b & (1 << j)]))
             for b in range(1 << n)]
    # Closure test from the rank axiom, independent of M.flats().
    flats = [b for b in range(1 << n)
             if all(ranks[b | (1 << j)] > ranks[b]
                    for j in range(n) if not b & (1 << j))]
    flats.sort(key=lambda b: (ranks[b], b))
    aug = {}
    for f in reversed(flats):
        upper = [g for g in flats if g != f and f & g == f]
        if not upper:
            aug[f] = Q.one()
            continue
        mu = {f: 1}
        for g in upper:
            mu[g] = -sum(a for h, a in mu.items() if h & g == h)
        aug[f] = -sum((mu[g] * Q([1] * (ranks[g]-ranks[f]+1)) * aug[g]
                       for g in upper), Q.zero())
    return aug[0], len(flats), ranks[-1], n


def partitions(n):
    if n == 0:
        yield ()
        return
    for p in partitions(n-1):
        for j in range(max(p, default=-1)+2):
            yield p+(j,)


PUBLISHED = {
    'U(3,4)': [1,11,11,1], 'U(4,5)': [1,26,66,26,1],
    'U(5,6)': [1,57,302,302,57,1],
    'U(6,7)': [1,120,1191,2416,1191,120,1],
    'U(7,8)': [1,247,4293,15619,15619,4293,247,1],
    'U(3,5)': [1,16,16,1], 'U(4,6)': [1,42,117,42,1],
    'U(5,7)': [1,99,610,610,99,1],
    'U(6,8)': [1,219,2641,5637,2641,219,1],
    'U(7,9)': [1,466,10204,40444,40444,10204,466,1],
}


def independent(key):
    if key.startswith('U('):
        return uniform_check(*map(int,key[2:-1].split(',')))
    if key.startswith('K('):
        return braid_check(int(key[2:-1]))
    return named_check(key)[0]


def main():
    start = time.monotonic()
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
        numberdb.configure(api_key=sys.stdin.read().strip())
    # Known nonzero controls precede all comparisons.
    assert eulerian(0) == 1 and eulerian(3) == Q([1,4,1])
    assert cycle_counts(3) == [0,2,3,1]
    assert braid_check(1) == 1 and braid_check(3) == Q([1,4,1])
    assert uniform_check(1,3) == 1+x
    assert uniform_check(2,3) == Q([1,4,1])
    assert named_check('Fano') == (Q([1,15,15,1]),16,3,7)
    assert len(list(partitions(3))) == 5
    gen = MatroidAugmentedChow()
    expected_keys = [p['matroid'] for p in gen.enumerate()]
    if os.environ.get('NUMBERDB_CHECK_STORED') == '1':
        raw = numberdb.table('T582')['Numbers']
        if isinstance(raw, dict):
            values = {k: Q(strip(v)) for k,v in raw.items()}
        else:
            values = {row['params']['matroid']: Q(strip(row)) for row in raw}
        assert list(values) == expected_keys, (list(values),expected_keys)
    else:
        values = {p['matroid']: strip(gen.value(p,100)) for p in gen.enumerate()}
    assert len(values) == 37
    assert not exactness(values)
    assert not names_its_rings(str(Path(__file__).with_name('generate.py')))
    for key in expected_keys:
        assert values[key] == independent(key), (key, values[key], independent(key))
    assert not agrees_with(values, independent)
    for key, coeffs in PUBLISHED.items():
        assert list(values[key]) == coeffs, (key,list(values[key]),coeffs)
    metadata = {}
    for key, p in values.items():
        if key.startswith('U('):
            r,n=map(int,key[2:-1].split(',')); flats=1+sum(comb(n,j) for j in range(r))
        elif key.startswith('K('):
            m=int(key[2:-1]); r=m-1; n=comb(m,2); flats=sum(1 for _ in partitions(m))
        else:
            _,flats,r,n=named_check(key)
        coeff=list(p)
        assert p.degree()==r and coeff[0]==coeff[-1]==1
        assert coeff==coeff[::-1] and all(c>=0 for c in coeff)
        assert all(coeff[j]<=coeff[j+1] for j in range(r//2))
        assert p[1] == flats-1, (key,p[1],flats)
        # Proposition 3.7 and Remark 3.8, using the generator's ordinary
        # upper-interval values; separate from independent coefficient checks.
        if key.startswith('U('):
            weighted = uniform_ordinary(r,n) + x**r + sum(
                (comb(n,j)*x**j*uniform_ordinary(r-j,n-j)
                 for j in range(1,r)), Q.zero())
            unweighted = 1+x*sum((comb(n,j)*uniform_ordinary(r-j,n-j)
                                 for j in range(r)), Q.zero())
        elif key.startswith('K('):
            ranks_by_blocks = {}
            for partition in partitions(m):
                k=max(partition)+1
                ranks_by_blocks[k]=ranks_by_blocks.get(k,0)+1
            weighted=sum((count*x**(m-k)*braid_ordinary(k)
                          for k,count in ranks_by_blocks.items()),Q.zero())
            unweighted=1+x*sum((count*braid_ordinary(k)
                               for k,count in ranks_by_blocks.items() if k>1),Q.zero())
        else:
            ranks,ordinary=named_intervals(key)
            weighted=sum((x**ranks[f]*ordinary[f] for f in ranks),Q.zero())
            unweighted=1+x*sum((ordinary[f] for f in ranks if ranks[f]<r),Q.zero())
        assert p==weighted==unweighted, key
        metadata[key]={'rank':r,'size':n,'flats':flats}
    print(json.dumps({'table':'T582','stored':os.environ.get('NUMBERDB_CHECK_STORED')=='1',
        'versions':{'SageMath':SAGE_VERSION,'numberdb':numberdb.__version__},
        'controls':'passed','independent_full_polynomial_comparisons':len(values),
        'published_full_polynomial_comparisons':len(PUBLISHED),
        'degree_endpoints_nonnegativity_symmetry_unimodality_linear_coefficient':len(values),
        'measure':measure(values),'seconds':round(time.monotonic()-start,3),
        'metadata':metadata,'values':{k:str(p) for k,p in values.items()}},indent=2))


if __name__ == '__main__':
    main()
