"""Independent checks for numberdb.org/T496.

sage -python checks.py checks locally. With NUMBERDB_CHECK_STORED=1 it also
reads the authenticated draft and verifies its entries independently.
The checks deliberately do not call the generator's order_of or weyl_order
when building the independent answers. Weyl orders use Sage's positive-root
heights, quotient orders use Chenevier's separately transcribed group
identifications, and the mass uses Bernoulli numbers.
"""
import json
import os
import re
import sys
from functools import lru_cache
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.arith.misc import bernoulli
from sage.combinat.root_system.root_system import RootSystem
from sage.rings.polynomial.polynomial_ring_constructor import PolynomialRing

from generate import NiemeierAutomorphismOrders
from check import exactness, measure, agrees_with, names_its_rings

HERE = Path(__file__).resolve().parent
QUOTIENTS = json.loads((HERE/'independent-quotients.json').read_text())


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


@lru_cache(None)
def root_height_order(letter, rank):
    result = QQ(1)
    for root in RootSystem([letter, rank]).root_lattice().positive_roots():
        h = root.height()
        result *= QQ(h+1)/QQ(h)
    assert result.denominator() == 1
    return ZZ(result)


def independent_order(key):
    a, b, _ = QUOTIENTS[key]
    result = ZZ(a)*ZZ(b)
    if key != 'Leech':
        for component in re.finditer(r'([ADE])(\d+)(?:\^(\d+))?', key):
            t, n, m = component.groups()
            result *= root_height_order(t, int(n))**int(m or 1)
    return result


def mass(dimension):
    result = abs(QQ(bernoulli(dimension//2))) / ZZ(dimension)
    for j in range(1, dimension//2):
        result *= abs(QQ(bernoulli(2*j))) / ZZ(4*j)
    return result


def validate(values):
    # Known controls are run before this method is trusted for rank 24.
    assert root_height_order('A', 2) == 6
    assert root_height_order('D', 4) == 192
    assert mass(8) == QQ(1)/ZZ(696729600)
    assert mass(16) == QQ(1)/(ZZ(2)*ZZ(696729600)**2) + QQ(1)/(ZZ(2)**15*ZZ(20922789888000))
    assert set(values) == set(QUOTIENTS)
    assert not exactness(values)
    assert not agrees_with(values, independent_order)
    total = sum((QQ(1)/v for v in values.values()), QQ(0))
    assert total == mass(24)
    # The mass check must reject one miscopied order and an omitted row.
    wrong = dict(values)
    wrong['A1^24'] += 1
    assert sum((QQ(1)/v for v in wrong.values()), QQ(0)) != mass(24)
    assert total - QQ(1)/values['Leech'] != mass(24)
    assert values['E8^3'] == ZZ(6)*ZZ(696729600)**3
    # Independently enumerate the allowable nonempty ADE root systems.
    by_h = {}
    for t, ranks in [('A',range(1,25)),('D',range(4,25)),('E',[6,7,8])]:
        for n in ranks:
            h = n+1 if t=='A' else 2*n-2 if t=='D' else {6:12,7:18,8:30}[n]
            by_h.setdefault(h,[]).append((t,n))
    possibilities = set()
    def enumerate_rank(items, index, remaining, parts):
        if index == len(items):
            if remaining == 0:
                possibilities.add(''.join(t+str(n)+('^'+str(m) if m>1 else '') for t,n,m in parts))
            return
        t,n = items[index]
        for multiplicity in range(remaining//n+1):
            enumerate_rank(items,index+1,remaining-n*multiplicity,
                           parts+([(t,n,multiplicity)] if multiplicity else []))
    for components in by_h.values():
        enumerate_rank(components,0,24,[])
    assert possibilities == set(values)-{'Leech'}
    return {'entries':len(values),'mass':str(total),'mass_formula':str(mass(24)),
            'independent_orders_matched':24,'ADE_root_systems':len(possibilities),
            'positive_controls':['A2 order 6','D4 order 192','rank-8 mass','rank-16 mass'],
            'negative_controls':['one order increased by 1','Leech row omitted'],
            'measurement':measure(values)}


def check_corpus(values, documents):
    sporadic = documents['T77']['Numbers']
    assert values['Leech'] == 2*ZZ(number_only(sporadic['Co1']))
    assert values['A1^24'] == ZZ(2)**24*ZZ(number_only(sporadic['M24']))
    assert values['A2^12'] == ZZ(6)**12*2*ZZ(number_only(sporadic['M12']))
    assert values['Leech'] == ZZ(number_only(documents['T495']['Numbers']['Lambda']['24']))
    ring = PolynomialRing(ZZ,'t')
    compared = []
    for key, polynomial in documents['T257']['Numbers'].items():
        m=re.fullmatch(r'([ADE])(\d+)',key)
        if m:
            assert ring(number_only(polynomial))(1) == root_height_order(m[1],int(m[2]))
            compared.append(key)
    return {'sporadic_relations':3,'classical_Leech_overlap':True,'stored_Coxeter_polynomials':compared}


def main():
    assert not names_its_rings(str(HERE/'generate.py'))
    gen = NiemeierAutomorphismOrders()
    values = {p['root_system']:number_only(gen.value(p,100)) for p in gen.enumerate()}
    report = {'local':validate(values)}
    if (HERE/'corpus.json').exists():
        report['corpus']=check_corpus(values,json.loads((HERE/'corpus.json').read_text()))
    if os.environ.get('NUMBERDB_CHECK_STORED') == '1':
        if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
            token=sys.stdin.read().strip()
            if '=' in token and token.split('=',1)[0].isupper():
                token=token.split('=',1)[1].strip().strip('\"\'')
            os.environ['NUMBERDB_API_KEY']=token
        saved=numberdb.table('T496')
        stored={k:ZZ(number_only(v)) for k,v in saved['Numbers'].items()}
        report['stored']=validate(stored)
        documents={tid:numberdb.table(tid) for tid in ['T77','T495','T257']}
        report['stored_corpus']=check_corpus(stored,documents)
        verified=gen.verify(sample=None)
        assert verified.ok
        report['verify']=str(verified)
    print(json.dumps(report,indent=2))
    return report

if __name__=='__main__':
    main()