"""Check T609 against live T99, T122, T61, T397 and its Programs snippet.

    sage -python companion-checks.py

The published companions need no key. Run beside generate.py and checks.py.
The standalone program is copied from the final document.json in the build
artifacts; after filling, this script reads it directly from T609 when a key is
available. Set NUMBERDB_API_KEY or NUMBERDB_KEY_FROM_STDIN for the draft read.
"""
import json
import os
import re
from pathlib import Path
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.rings.real_arb import RealBallField
from sage.rings.number_field.number_field import CyclotomicField
from generate import QuantumDimensions, parse, authenticate_stdin
from checks import expected_character, embedding, written_bounds, R, I, bounded_agreement

def number(x):
    return str(x['number'] if isinstance(x,dict) else x)

def interval(text):
    lo,hi=written_bounds(text)
    return R(I(lo,hi))

def path_order(g):
    n=ord(g[0])-63
    if n<2 or n>7:return None
    bits=''.join(format(ord(c)-63,'06b') for c in g[1:])
    neighbours=[set() for _ in range(n)];pos=0
    for j in range(1,n):
        for i in range(j):
            if bits[pos]=='1':neighbours[i].add(j);neighbours[j].add(i)
            pos+=1
    if sorted(map(len,neighbours)) != [1,1]+[2]*(n-2):return None
    reached={0}
    while True:
        new=reached|set.union(*(neighbours[v] for v in reached))
        if new==reached:break
        reached=new
    return n if len(reached)==n else None

authenticate_stdin()
g=QuantumDimensions();params=list(g.enumerate())
t99=numberdb.table('T99')['Numbers'];t122=numberdb.table('T122')['Numbers']
t61=numberdb.table('T61')['Numbers'];t397=numberdb.table('T397')['Numbers']
counts={'T99':0,'T122':0,'T61':0,'T397':0,'program':0}
P=PolynomialRing(QQ,'x')
for p in params:
    N,k,a=parse(p);exact=expected_character(N,k,a)
    q=exact.parent().gen()
    if N==2:
        assert P(number(t99[str(a[0])]))((q+q**(-1))/2)==exact
        counts['T99']+=1
    if N<=4:
        partition=tuple(sum(a[i:]) for i in range(N-1))
        nonzero=tuple(x for x in partition if x)
        if nonzero and sum(nonzero)<=4:
            key=','.join(map(str,nonzero))
            poly=PolynomialRing(QQ,tuple('x'+str(i) for i in range(1,N+1)))(number(t122[str(N)][key]))
            assert poly(*[q**(N-1-2*j) for j in range(N)])==exact
            counts['T122']+=1
for k in range(1,21):
    key=str(QQ(1)/(k+2))
    bounded_agreement(2*interval(number(t61[key])),embedding(expected_character(2,k,(1,)),k+2))
    counts['T61']+=1
path_keys={}
for key,record in t397.items():
    n=path_order(key)
    if n:
        k=n-1
        bounded_agreement(interval(number(record)),embedding(expected_character(2,k,(1,)),k+2))
        path_keys[str(k)]=key
        counts['T397']+=1
assert counts['T397']==6
if os.environ.get('NUMBERDB_API_KEY'):
    doc=numberdb.table('T609')
else:
    doc=json.loads(Path(__file__).with_name('document.json').read_text())
exec(doc['Programs']['Sage']['code'],globals())
for p in params:
    N,k,a=parse(p)
    value=qdim(N,k,list(a))
    bounded_agreement(value,embedding(expected_character(N,k,a),N+k))
    counts['program']+=1
print(json.dumps({'table':'T609','companion_checks':counts,'path_entry_keys':path_keys,'source':'live NumberDB companion documents'},sort_keys=True))