"""Feynman periods of primitive log-divergent phi4 graphs -- numberdb.org/T464.

Run with SageMath (PARI >= 2.17) and numberdb installed:
    sage -pip install numberdb
    sage -python generate.py                 # verify all stored entries
    sage -python checks.py                   # compute and check without writing

The attached generator recomputes every non-FAIL period in the version-2
Panzer--Schnetz ancillary file, whose SHA-256 is pinned below. Download
Periods beside this file, or let the generator fetch the pinned public URL.
The Maple data is parsed with a restricted expression evaluator, never exec.
Rationals stay exact. PARI zetamult uses descending indices (unlike Sage's
Multizeta wrapper); polylogmult([s1,...],[xi,1,...]) handles the colored sums.
Every value is computed at 150 and 200 decimal digits, compared with the
source's 100 significant digits, then written to 100 significant digits.
Agreement is heuristic: no interval proof of these evaluations is asserted.
The nonreal root-of-unity arguments are displaced inward by 16 binary units
at each working precision to avoid a spurious domain refusal from rounding.
This is included in the two-precision agreement check; it is not an error proof.
The one-loop integral is exactly 1 and is returned as a rational.
Duplicate exact source expressions retain all census labels and carry equals.
"""
import ast
from functools import lru_cache
import hashlib
from pathlib import Path
import re
import sys
import urllib.request

import numberdb.sage as numberdb
from sage.libs.pari import pari
from sage.rings.complex_mpfr import ComplexField
from sage.rings.rational_field import QQ

SOURCE_URL = 'https://arxiv.org/src/1603.04289v2/anc/Periods'
SOURCE_SHA256 = '093d458994903b3b397f3ce3cd5cc9578221c065e5a6b3fb0947abf1b2ddce19'
WORKING_DIGITS = (150, 200)
# Fifty guard digits at the lower precision are measured in checks.py.
BASE = Path(__file__).resolve().parent
sys.setrecursionlimit(max(sys.getrecursionlimit(), 20000))


def split_fields(text):
    depth, start, out = 0, 0, []
    for i, char in enumerate(text):
        if char in '[{(':
            depth += 1
        elif char in ']})':
            depth -= 1
        elif char == ',' and depth == 0:
            out.append(text[start:i].strip())
            start = i + 1
        if depth < 0:
            raise ValueError('Unbalanced source record')
    if depth:
        raise ValueError('Unbalanced source record')
    out.append(text[start:].strip())
    return out


@lru_cache(None)
def records():
    path = BASE / 'Periods'
    raw = path.read_bytes() if path.exists() else urllib.request.urlopen(SOURCE_URL, timeout=120).read()
    if hashlib.sha256(raw).hexdigest() != SOURCE_SHA256:
        raise ValueError('Ancillary source has changed; inspect it before proceeding')
    out = []
    pattern = r'Period\[(\d+),(\d+)\]:=\[(.*?)\]:'
    for match in re.finditer(pattern, raw.decode(), re.S):
        fields = split_fields(match[3])
        if len(fields) != 7:
            raise ValueError('Expected seven fields')
        loops, index = int(match[1]), int(match[2])
        out.append({'loops': loops, 'index': index, 'edges': fields[0],
                    'expression': fields[2], 'source_number': fields[3],
                    'ancestor': fields[5]})
    if len(out) != 10119:
        raise ValueError('Unexpected census size')
    return out


@lru_cache(None)
def known_records():
    return tuple(r for r in records() if r['source_number'] != 'FAIL')


@lru_cache(None)
def record_map():
    return {(r['loops'], r['index']): r for r in known_records()}


@lru_cache(None)
def expression_tree(text):
    return ast.parse(text.replace('^', '**'), mode='eval').body


class Evaluator:
    """Only integers, rational arithmetic, Pi, I, and zeta subscripts."""
    def __init__(self, digits):
        self.bits = numberdb.bits(digits)
        self.C = ComplexField(self.bits)
        self.atoms = {}

    def evaluate(self, node):
        if isinstance(node, ast.Constant) and type(node.value) is int:
            return QQ(node.value)
        if isinstance(node, ast.Name) and node.id in ('Pi', 'I'):
            return self.C.pi() if node.id == 'Pi' else self.C.gen()
        if isinstance(node, ast.UnaryOp) and isinstance(node.op, (ast.UAdd, ast.USub)):
            v = self.evaluate(node.operand)
            return v if isinstance(node.op, ast.UAdd) else -v
        if isinstance(node, ast.BinOp):
            a, b = self.evaluate(node.left), self.evaluate(node.right)
            if isinstance(node.op, ast.Add): return a + b
            if isinstance(node.op, ast.Sub): return a - b
            if isinstance(node.op, ast.Mult): return a * b
            if isinstance(node.op, ast.Div): return a / b
            if isinstance(node.op, ast.Pow):
                if b in QQ and QQ(b).denominator() == 1:
                    return a ** int(b)
                return self.C(a) ** self.C(b)
        if isinstance(node, ast.Subscript) and isinstance(node.value, ast.Name) and node.value.id == 'zeta':
            key = ast.dump(node)
            if key in self.atoms:
                return self.atoms[key]
            args = node.slice.elts if isinstance(node.slice, ast.Tuple) else [node.slice]
            if isinstance(args[0], ast.List):
                pair = args[0].elts
                if len(pair) != 2: raise ValueError('Unknown colored zeta')
                xi = self.evaluate(pair[0])
                if self.C(xi).imag() != 0:
                    # PARI rejects a unit root rounded just outside its domain.
                    # Approach its continuous boundary value from inside; the
                    # 16-ulp displacement vanishes as working precision grows.
                    xi = self.C(xi) * (1 - QQ(2)**(-self.bits + 4))
                    if QQ(xi.real())**2 + QQ(xi.imag())**2 >= 1:
                        raise ValueError('Colored argument was not rounded inward')
                weights = [int(self.evaluate(pair[1]))] + [int(self.evaluate(n)) for n in args[1:]]
                value = self.C(pari.polylogmult(weights, [xi] + [1] * (len(weights)-1), precision=self.bits))
            else:
                weights = [int(self.evaluate(n)) for n in args]
                if weights[0] < 2 or min(weights) < 1:
                    raise ValueError('Divergent MZV')
                value = self.C(pari.zetamult(weights, precision=self.bits))
            self.atoms[key] = value
            return value
        raise ValueError('Unsupported expression node: ' + ast.dump(node)[:120])


@lru_cache(None)
def evaluator(digits):
    return Evaluator(digits)


@lru_cache(None)
def evaluate_expression(text, digits):
    ev = evaluator(digits)
    result = ev.C(ev.evaluate(expression_tree(text)))
    if any(x.is_NaN() or x.is_infinity() for x in (result.real(), result.imag())):
        raise ValueError('Nonfinite period')
    # Imaginary cancellation is expected for conjugate root-of-unity sums.
    if abs(result.imag()) > max(ev.C(1).real(), abs(result.real())) * QQ(10)**(-digits+25):
        raise ValueError('Imaginary remainder exceeds the guard')
    return result.real()


def decimal_rational(text):
    """The exact rational at the source's written decimal midpoint."""
    if '.' not in text:
        return QQ(text), QQ(0)
    whole, fraction = text.split('.')
    return QQ(int(whole + fraction)) / QQ(10)**len(fraction), QQ(10)**(-len(fraction))


@lru_cache(None)
def first_labels():
    first = {}
    for row in known_records():
        first.setdefault(row['expression'], (row['loops'], row['index']))
    return first


class FeynmanPhi4Periods(numberdb.Generator):
    table = 'T464'
    parameters = ('loops', 'index')
    type = 'R'
    digits = 100
    rigour = 'heuristic (agreement-checked)'

    def enumerate(self):
        for row in known_records():
            yield {'loops': row['loops'], 'index': row['index']}

    def value(self, params, digits):
        if digits != 100:
            raise ValueError('This source is checked at 100 significant digits')
        key = (int(params['loops']), int(params['index']))
        row = record_map()[key]
        if key == (1, 1):
            if row['expression'] != '1': raise ValueError('One-loop control changed')
            return {'number': QQ(1), 'comment': r'The one-loop period is exactly $1$, since $\int_0^\infty (1+x)^{-2}\,dx=1$.'}
        low, high = [evaluate_expression(row['expression'], d) for d in WORKING_DIGITS]
        source, unit = decimal_rational(row['source_number'])
        if abs(high - source) > unit:
            raise ValueError('Source disagrees at %r' % (key,))
        if abs(high - low) > abs(high) * QQ(10)**(-120):
            raise ValueError('Fewer than 120 stable digits at %r' % (key,))
        value = numberdb.agreeing(lambda d: evaluate_expression(row['expression'], d), at=WORKING_DIGITS)
        entry = {'number': value}
        first = first_labels()[row['expression']]
        if first != key:
            entry['equals'] = 'HREF{#%d,%d}' % first
        if key == (3, 1):
            entry['comment'] = r'The wheel with three spokes has period $6\zeta(3)$; its completion is $K_5$.'
        return entry


if __name__ == '__main__':
    import os
    if os.environ.get('NUMBERDB_KEY_FROM_STDIN') == '1':
        numberdb.configure(api_key=sys.stdin.read().strip())
    report = FeynmanPhi4Periods().verify(sample=None)
    print(report)
    sys.exit(0 if report.ok else 1)
