"""Exact independent checks for Orthogonal Weingarten functions, T550.

Run with SageMath and numberdb installed:
    sage -python independent_checks.py             # locally generated values
    sage -python independent_checks.py --stored    # values read from T550

Private drafts require NUMBERDB_API_KEY. Campaign runs may pipe the key with
NUMBERDB_KEY_FROM_STDIN=1, which selects stored mode too.
The checks construct pairings directly, never Jack/spherical polynomials.
They compress one full Gram row per orbit to p(k) rational equations. Solving
these equations independently recovers every coefficient; residuals test each
representative of the full inverse relation. The sphere-moment check uses
Gaussian polar coordinates: E[X_1^(2k)] = (2k-1)!!/prod(N+2j, j=0..k-1).
Controls precede the general checks, including a deliberately altered value.
"""

import numberdb.sage as numberdb
from sage.rings.rational_field import QQ
import sage.rings.real_mpfr  # Sage rational coercion dispatch; no floating arithmetic.
from collections import Counter
from functools import lru_cache
from math import prod
from pathlib import Path
import json
import os
import sys
from check import agrees_with, exactness, measure, names_its_rings


def integer_partitions(n, largest=None):
    if n == 0:
        yield ()
        return
    for first in range(min(n, largest if largest is not None else n), 0, -1):
        for rest in integer_partitions(n-first, first):
            yield (first,) + rest


def pairing_edges(vertices):
    if not vertices:
        yield ()
        return
    first = vertices[0]
    for j in range(1, len(vertices)):
        for rest in pairing_edges(vertices[1:j]+vertices[j+1:]):
            yield ((first, vertices[j]),) + rest


def partners(edges, n):
    result = [None]*n
    for a, b in edges:
        result[a], result[b] = b, a
    return tuple(result)


def union_type(p, q):
    seen = set()
    sizes = []
    for start in range(len(p)):
        if start in seen:
            continue
        stack, size = [start], 0
        while stack:
            v = stack.pop()
            if v in seen:
                continue
            seen.add(v)
            size += 1
            stack.extend((p[v], q[v]))
        assert size % 2 == 0
        sizes.append(size//2)
    return tuple(sorted(sizes, reverse=True))


def representative(mu):
    """Image of the standard pairing under Matsumoto's (2.3), zero based."""
    edges, offset = [], 0
    for part in mu:
        block = [offset, offset+2*part-1] + list(range(offset+1, offset+2*part-1))
        edges.extend(zip(block[::2],block[1::2]))
        offset += 2*part
    return partners(edges, offset)


@lru_cache(None)
def gram_counts(k):
    types = tuple(integer_partitions(k))
    ps = tuple(partners(e,2*k) for e in pairing_edges(tuple(range(2*k))))
    assert len(ps) == prod(range(1,2*k,2))
    base = tuple(i ^ 1 for i in range(2*k))
    base_types = [union_type(base, p) for p in ps]
    orbits = Counter(base_types)
    rows = {}
    for mu in types:
        q = representative(mu)
        assert union_type(base, q) == mu
        # G[base,p] * W[p,q], aggregated by loop count and unknown coset type.
        rows[mu] = Counter((len(nu), union_type(p,q)) for p,nu in zip(ps,base_types))
        assert sum(rows[mu].values()) == len(ps)
    return types, rows, orbits


def rational_solve(a, b):
    """Exact Gaussian elimination written out to avoid matrix backends."""
    matrix = [[QQ(x) for x in row]+[QQ(rhs)] for row,rhs in zip(a,b)]
    size = len(b)
    for col in range(size):
        pivot = next(i for i in range(col,size) if matrix[i][col])
        matrix[col], matrix[pivot] = matrix[pivot], matrix[col]
        divisor = matrix[col][col]
        matrix[col] = [x/divisor for x in matrix[col]]
        for row in range(size):
            if row == col:
                continue
            factor = matrix[row][col]
            matrix[row] = [x-factor*y for x,y in zip(matrix[row],matrix[col])]
    return [row[-1] for row in matrix]


def inverse_coefficients(N,k):
    types, counts, _ = gram_counts(k)
    a = [[sum(multiplicity*N**loops for (loops,t),multiplicity in counts[mu].items() if t==nu)
          for nu in types] for mu in types]
    b = [int(mu==(1,)*k) for mu in types]
    return dict(zip(types, rational_solve(a,b)))


def residuals(N,k,values):
    types, counts, _ = gram_counts(k)
    return {mu: sum(QQ(m)*N**loops*values[nu] for (loops,nu),m in counts[mu].items())-int(mu==(1,)*k)
            for mu in types}


def small_closed(N,mu):
    N=QQ(N)
    if mu == (1,): return 1/N
    if mu == (2,): return -1/(N*(N+2)*(N-1))
    if mu == (1,1): return (N+1)/(N*(N+2)*(N-1))
    d=N*(N+2)*(N+4)*(N-1)*(N-2)
    if mu == (3,): return 2/d
    if mu == (2,1): return -(N+2)/d
    if mu == (1,1,1): return (N*N+3*N-2)/d
    raise ValueError(mu)


def controls():
    assert rational_solve([[2,1],[1,3]],[1,2]) == [QQ(1)/5,QQ(3)/5]
    known = {(2,):-QQ(1)/30,(1,1):QQ(2)/15}
    assert inverse_coefficients(3,2)==known
    assert all(v==0 for v in residuals(3,2,known).values())
    corrupted=dict(known); corrupted[(2,)]+=QQ(1)/1000
    assert any(v!=0 for v in residuals(3,2,corrupted).values())
    _,_,orbits=gram_counts(2)
    assert orbits=={(2,):2,(1,1):1}
    assert 3*sum(orbits[mu]*v for mu,v in known.items())==QQ(1)/5
    print('Controls passed: nonzero linear solve, N=3 k=2 coefficients and sphere moment; corruption rejected.')


def run_checks(values):
    controls()
    expected_keys = [(N,mu) for N in range(1,13) for k in range(1,min(N,6)+1) for mu in integer_partitions(k)]
    assert list(values)==expected_keys
    assert not exactness(values)
    independent = {}
    gram_equations = 0
    sphere_moments = 0
    for k in range(1,7):
        types, counts, orbits=gram_counts(k)
        for N in range(k,13):
            v={mu:values[N,mu] for mu in types}
            inv=inverse_coefficients(N,k)
            independent.update({(N,mu):x for mu,x in inv.items()})
            assert all(r==0 for r in residuals(N,k,v).values())
            gram_equations+=len(types)
            pairing_count=prod(range(1,2*k,2))
            actual=pairing_count*sum(orbits[mu]*v[mu] for mu in types)
            expected=QQ(pairing_count)/QQ(prod(N+2*j for j in range(k)))
            assert actual==expected
            sphere_moments+=1
        print('k=%d: %d pairings; all dimensions %d..12 satisfy Gram inverse and sphere moments.' % (k,sum(orbits.values()),k))
    disagreements=agrees_with(values,lambda key:independent[key])
    assert not disagreements,disagreements
    small=0
    for (N,mu),value in values.items():
        if sum(mu)<=3:
            assert value==small_closed(N,mu)
            small+=1
    summary={'entries':len(values),'independent_coefficient_agreements':len(independent),
             'gram_inverse_equations':gram_equations,'sphere_moments':sphere_moments,
             'published_small_degree_formulas':small,'measurement':measure(values),
             'skill_manifest':'agents/runs/20261027T133445Z-build-skills/manifest.json',
             'skill_sha256':'e37efb7e2ededf6c1d7134c673f152f648f26e0267c556aebdf3ec65765d8b4c'}
    print(json.dumps(summary,indent=2,default=str))
    return summary


def stored_values(document):
    """Read API entries without importing the generator or its arithmetic."""
    result={}
    block=document['Numbers']
    if isinstance(block, list):
        entries=block
    else:
        entries=[{'params': {'N':N,'mu':mu},
                  'number':value['number'] if isinstance(value,dict) else value}
                 for N, row in block.items() for mu,value in row.items()]
    for entry in entries:
        params=entry['params']
        key=(int(params['N']),tuple(map(int,str(params['mu']).split(','))))
        assert key not in result
        numerator, slash, denominator=str(entry['number']).partition('/')
        result[key]=QQ(int(numerator))/QQ(int(denominator) if slash else 1)
    return result


if __name__ == '__main__':
    use_stored='--stored' in sys.argv or os.environ.get('NUMBERDB_KEY_FROM_STDIN')=='1'
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN')=='1':
        key=sys.stdin.read().strip()
        if '=' in key and key.split('=',1)[0].isupper():key=key.split('=',1)[1].strip().strip('\'"')
        numberdb.configure(api_key=key)
    if use_stored:
        values=stored_values(numberdb.table('T550'))
    else:
        from generate import OrthogonalWeingarten
        generator=OrthogonalWeingarten()
        values={(int(p['N']),tuple(map(int,p['mu'].split(',')))):generator.value(p,100)
                for p in generator.enumerate()}
    run_checks(values)
    print('VALUES_JSON '+json.dumps([{'params':{'N':str(N),'mu':','.join(map(str,mu))},'number':str(x)} for (N,mu),x in values.items()]))
