"""Estrada indices of connected graphs -- numberdb.org/T545.

Run with SageMath 10.9 and numberdb 0.1.10 (or compatible versions):
    sage -pip install numberdb
    sage -python generate.py                 # verify all stored rows
    sage -python checks.py                   # independent numerical checks
    sage -python generate.py --fill-draft    # write entries with NUMBERDB_API_KEY

The last command fills the existing draft; it never offers or publishes it.
Default selection: every connected simple graph of orders 2 through 6.
Canonical labels use algorithm='sage'; see argument-selection.md. The integer
adjacency characteristic polynomial is factored and its real algebraic roots
are isolated in AA, including multiplicities. Conversion to Arb and all
exponentials and sums are certified. The initial precision has 80 guard bits;
precision rises by 64 bits if the measured result has less than the requested
precision plus 24 reliable bits. checks.py measures every enclosure and checks
an independent exact closed-walk sum with a rigorously bounded positive tail.
"""
import os
import sys
from functools import lru_cache

import numberdb.sage as numberdb
from sage.graphs.graph import Graph
from sage.graphs.graph_generators import graphs
from sage.rings.qqbar import AA
from sage.rings.real_arb import RealBallField

WORKING_GUARD = 80
EXPECTED_COUNTS = {2: 1, 3: 2, 4: 6, 5: 21, 6: 112}


def canonical_key(graph):
    return graph.canonical_label(algorithm='sage').graph6_string()


@lru_cache(None)
def graph_keys():
    keys = []
    for n, count in EXPECTED_COUNTS.items():
        group = [canonical_key(g) for g in graphs.nauty_geng(f'{n} -c')]
        assert len(group) == len(set(group)) == count
        keys.extend(sorted(group))
    return tuple(keys)


@lru_cache(None)
def named_graphs():
    names = {}
    # First applicable standard name takes precedence for overlapping families.
    for n in range(2, 7):
        candidates = [(graphs.CompleteGraph(n), f'This is the complete graph $K_{n}$.')]
        if n >= 3:
            candidates.extend([
                (graphs.PathGraph(n), f'This is the path $P_{n}$.'),
                (graphs.CycleGraph(n), f'This is the cycle $C_{n}$.'),
            ])
        for a in range(1, n // 2 + 1):
            b = n - a
            candidates.append((graphs.CompleteBipartiteGraph(a, b),
                               f'This is the complete bipartite graph $K_{{{a},{b}}}$.'))
        for g, name in candidates:
            names.setdefault(canonical_key(g), name)
    return names


@lru_cache(None)
def algebraic_spectrum(key):
    graph = Graph(key, format='graph6')
    assert graph.is_connected() and graph.order() >= 2
    polynomial = graph.adjacency_matrix().charpoly('x')
    roots = tuple(polynomial.roots(ring=AA))
    assert sum(m for _, m in roots) == graph.order()
    return roots


def estrada(key, digits=100):
    precision = numberdb.bits(digits, losing=WORKING_GUARD)
    for _ in range(8):
        field = RealBallField(precision)
        value = sum((field(root).exp() * multiplicity
                     for root, multiplicity in algebraic_spectrum(key)), field.zero())
        if value.is_finite() and value.accuracy() >= numberdb.bits(digits, losing=24):
            return value
        precision += 64
    raise ArithmeticError('Insufficient certified precision')


class EstradaIndices(numberdb.Generator):
    table = 'T545'
    parameters = ('g',)
    type = 'R'
    digits = 100
    rigour = 'proven'
    files = ('generate.py', 'checks.py', 'check.py', 'dry_run.py',
             'argument-selection.md', 'corpus-reference.json',
             'graph2c.g6', 'graph3c.g6', 'graph4c.g6', 'graph5c.g6', 'graph6c.g6',
             'source-downloads.json', 'skill-evidence.json', 'skill-manifest.json',
             'numerical-evidence.md', 'checks-before.json')

    def enumerate(self):
        for key in graph_keys():
            yield {'g': key}

    def value(self, params, digits):
        key = str(params['g'])
        result = {'number': estrada(key, digits)}
        if key in named_graphs():
            result['comment'] = named_graphs()[key]
        return result


if __name__ == '__main__':
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
        key = sys.stdin.read().strip()
        if '=' in key and key.split('=', 1)[0].isidentifier():
            key = key.split('=', 1)[1].strip().strip('\"\'')
        numberdb.configure(api_key=key)
    generator = EstradaIndices()
    if '--fill-draft' in sys.argv or os.environ.get('NUMBERDB_PUBLISH') == '1':
        print(generator.publish(message='Fill 142 certified adjacency Estrada indices with independent checks'))
    else:
        report = generator.verify(sample=None)
        print(report)
        sys.exit(0 if report.ok else 1)