"""Reproduce all T92 arguments and certify the written Sobolev constants.

    python -m pip install numberdb python-flint==0.6.0 mpmath==1.3.0 pyyaml
    python generate.py --self-test
    python generate.py --compute
    python generate.py --document table.yaml

These commands are local and never access the website. The default selection
is the original reduced p,q numerator/denominator bound 24, with integer
n=p*q/(q-p), together with p=2 in every dimension 3..20. Four rows are added.
The old generate.sage remains attached as the historical 1024-row source.

The producer evaluates the table's gamma formula in Arb using exact rational
inputs and 160 guard bits. Every stored interval is checked against a fresh
enclosure and against mpmath at two precisions. The four added odd-dimensional
p=2 rows also have an independent, gamma-free ball certificate from the
radial integrals of u(r)=(1+r*r)^(-(n-2)/2). Substitution r=tan(theta) gives
integrals of even powers of sine and cosine, evaluated by rational reduction
to pi/2. Numerical quadrature of these integrals is a supplementary check,
not a proof of their errors. No stored value is used as a computation input.
"""

import argparse
import json
import math
from contextlib import contextmanager
from decimal import Decimal
from fractions import Fraction
from pathlib import Path

import mpmath as mp
import numberdb
from flint import arb, ctx
from numberdb._write import to_text


GUARD_BITS = 160
ADDED_DIMENSIONS = (13, 15, 17, 19)
PARAMETERS = ("n", "p", "q")


@contextmanager
def precision(digits):
    original = ctx.prec
    ctx.prec = numberdb.bits(digits, losing=GUARD_BITS)
    try:
        yield
    finally:
        ctx.prec = original


def rational(value):
    return arb(str(Fraction(value)))


def input_values(params):
    dimension = Fraction(params["n"])
    exponent = Fraction(params["p"])
    conjugate = Fraction(params["q"])
    if dimension.denominator != 1 or not 1 <= exponent < dimension:
        raise ValueError("require integer n and 1 <= p < n")
    if conjugate != dimension * exponent / (dimension - exponent):
        raise ValueError("q must be n*p/(n-p), exactly")
    return int(dimension), exponent, conjugate


def selection():
    exponents = {Fraction(numerator, denominator)
                 for numerator in range(1, 25)
                 for denominator in range(1, numerator + 1)}
    identities = set()
    for exponent in exponents:
        for conjugate in exponents:
            if conjugate <= exponent:
                continue
            dimension = exponent * conjugate / (conjugate - exponent)
            if dimension.denominator == 1:
                identities.add((int(dimension), exponent, conjugate))
    for dimension in range(3, 21):
        identities.add((dimension, Fraction(2), Fraction(2 * dimension, dimension - 2)))
    return [dict(zip(PARAMETERS, map(str, identity))) for identity in sorted(identities)]


def formula(params):
    dimension, exponent, conjugate = input_values(params)
    dimension_ball = arb(dimension)
    exponent_ball = rational(exponent)
    logarithm = arb.pi().log() / 2
    if exponent == 1:
        logarithm += dimension_ball.log() - (1 + dimension_ball / 2).lgamma() / dimension
    else:
        logarithm += dimension_ball.log() / exponent_ball
        logarithm += (1 - 1 / exponent_ball) * rational((dimension - exponent) / (exponent - 1)).log()
        logarithm += ((dimension_ball / exponent_ball).lgamma()
                      + (dimension_ball + 1 - dimension_ball / exponent_ball).lgamma()
                      - dimension_ball.lgamma() - (1 + dimension_ball / 2).lgamma()) / dimension
    return logarithm.exp()


def endpoint(value):
    mantissa, exponent = value.man_exp()
    return Fraction(int(mantissa)) * Fraction(2) ** int(exponent)


def wire(value):
    return numberdb.RealInterval(endpoint(value.lower()), endpoint(value.upper()))


def ball_text(value, digits=100):
    return to_text(wire(value), digits)


def written_bounds(text):
    midpoint = Fraction(text)
    if "." not in text and "e" not in text.lower():
        return midpoint, midpoint, Fraction(0)
    unit = Fraction(10) ** Decimal(text).as_tuple().exponent
    return midpoint - unit, midpoint + unit, unit


def certify(text, value, digits=100):
    if not value.is_finite() or value.contains(0):
        raise AssertionError("nonfinite or nonpositive-informative result")
    if value.rel_accuracy_bits() < numberdb.bits(digits + 10, losing=0):
        raise AssertionError("insufficient enclosure accuracy")
    lower, upper, unit = written_bounds(text)
    if not unit:
        raise AssertionError("no exact rational Sobolev result is established here")
    left, right = endpoint(value.lower()), endpoint(value.upper())
    if not lower < left <= right < upper:
        raise AssertionError("written interval does not contain the enclosure")
    if lower + 10 * unit < left <= right < upper + 10 * unit:
        raise AssertionError("shifted interval unexpectedly accepted")
    return value.rel_accuracy_bits()


def mp_rational(value):
    value = Fraction(value)
    return mp.mpf(value.numerator) / value.denominator


def independent_formula(params):
    dimension, exponent, conjugate = input_values(params)
    dimension_mp, exponent_mp = mp.mpf(dimension), mp_rational(exponent)
    if exponent == 1:
        return mp.sqrt(mp.pi) * dimension_mp / mp.gamma(1 + dimension_mp / 2) ** (1 / dimension_mp)
    return (mp.sqrt(mp.pi) * dimension_mp ** (1 / exponent_mp)
            * ((dimension_mp - exponent_mp) / (exponent_mp - 1)) ** (1 - 1 / exponent_mp)
            * (mp.gamma(dimension_mp / exponent_mp)
               * mp.gamma(dimension_mp + 1 - dimension_mp / exponent_mp)
               / mp.gamma(dimension_mp) / mp.gamma(1 + dimension_mp / 2)) ** (1 / dimension_mp))


def even_trig_integral_factor(sine_power, cosine_power):
    if min(sine_power, cosine_power) < 0 or sine_power % 2 or cosine_power % 2:
        raise ValueError("require nonnegative even powers")
    coefficient = Fraction(1, 2)
    while sine_power:
        coefficient *= Fraction(sine_power - 1, sine_power + cosine_power)
        sine_power -= 2
    while cosine_power:
        coefficient *= Fraction(cosine_power - 1, cosine_power)
        cosine_power -= 2
    return coefficient


def p2_radial_certificate(dimension):
    if dimension < 3 or dimension % 2 != 1:
        raise ValueError("odd dimensions at least 3 only")
    half_order = (dimension - 1) // 2
    sphere = rational(Fraction(2 * 4 ** half_order * math.factorial(half_order), math.factorial(2 * half_order))) * arb.pi() ** half_order
    mass = sphere * rational(even_trig_integral_factor(dimension - 1, dimension - 1)) * arb.pi()
    energy = sphere * (dimension - 2) ** 2 * rational(even_trig_integral_factor(dimension + 1, dimension - 3)) * arb.pi()
    return energy.sqrt() / mass ** rational(Fraction(dimension - 2, 2 * dimension))


def radial_quadrature(dimension):
    half_order = (dimension - 1) // 2
    sphere = mp_rational(Fraction(2 * 4 ** half_order * math.factorial(half_order), math.factorial(2 * half_order))) * mp.pi ** half_order
    mass = mp.quad(lambda theta: mp.sin(theta) ** (dimension - 1) * mp.cos(theta) ** (dimension - 1), [0, mp.pi / 4, mp.pi / 2])
    energy = mp.quad(lambda theta: mp.sin(theta) ** (dimension + 1) * mp.cos(theta) ** (dimension - 3), [0, mp.pi / 4, mp.pi / 2])
    return mp.sqrt(sphere * (dimension - 2) ** 2 * energy) / (sphere * mass) ** mp_rational(Fraction(dimension - 2, 2 * dimension))


def controls():
    with precision(100):
        cases = [({"n": "2", "p": "1", "q": "2"}, 2 * arb.pi().sqrt()),
                 ({"n": "3", "p": "2", "q": "6"}, arb(3).sqrt() * (arb.pi() / 2) ** rational("2/3")),
                 ({"n": "4", "p": "2", "q": "4"}, 2 * rational("2/3") ** rational("1/4") * arb.pi().sqrt())]
        for params, known in cases:
            value = formula(params)
            certify(ball_text(known), value)
        certify(ball_text(cases[1][1]), p2_radial_certificate(3))
        positive = formula(cases[1][0])
        written = ball_text(positive)
        for invalid in (arb("nan"), arb(1, 100), positive + 1):
            try:
                certify(written, invalid)
            except AssertionError:
                pass
            else:
                raise AssertionError("negative enclosure control was not rejected")
        lower, upper, unit = written_bounds(written)
        shifted = rational(Fraction(written) + 10 * unit)
        try:
            certify(written, shifted)
        except AssertionError:
            pass
        else:
            raise AssertionError("last-place shift control was not rejected")
    for params in ({"n": "3", "p": "2", "q": "2"}, {"n": "3", "p": "3", "q": "6"}):
        try:
            input_values(params)
        except ValueError:
            pass
        else:
            raise AssertionError("invalid parameter control was not rejected")
    identities = selection()
    if len(identities) != 1028 or len({tuple(params.values()) for params in identities}) != 1028:
        raise AssertionError("default selection count or uniqueness failed")
    if even_trig_integral_factor(0, 0) != Fraction(1, 2) or even_trig_integral_factor(2, 2) != Fraction(1, 16):
        raise AssertionError("trigonometric integral controls failed")
    with mp.workdps(140):
        known = mp.sqrt(3) * (mp.pi / 2) ** (mp.mpf(2) / 3)
        if abs(radial_quadrature(3) / known - 1) > mp.mpf("1e-125"):
            raise AssertionError("nonzero independent quadrature control failed")
    return {"positive_closed_forms": ["S_2_1=2*sqrt(pi)", "S_3_2=sqrt(3)*(pi/2)^(2/3)", "S_4_2=2*(2/3)^(1/4)*sqrt(pi)"],
            "gamma_free_radial_control": True, "quadrature_control": True,
            "nan_wide_wrong_and_shifted_rejected": True, "invalid_parameters_rejected": True,
            "default_selection_count": 1028, "source_only": True}


def flatten(document):
    result = {}
    for dimension, by_exponent in document["Numbers"].items():
        for exponent, by_conjugate in by_exponent.items():
            for conjugate, entry in by_conjugate.items():
                result[(dimension, exponent, conjugate)] = entry
    return result


def check_document(document):
    report = {"tid": "T92", "controls": controls(), "records": [], "working_guard_bits": GUARD_BITS}
    rows = flatten(document)
    parameters = selection()
    if set(rows) != {tuple(params[name] for name in PARAMETERS) for params in parameters}:
        raise AssertionError("full stored and default selections differ")
    accuracies = []
    for params in parameters:
        identity = tuple(params[name] for name in PARAMETERS)
        leaf = rows[identity]
        written = str(leaf["number"] if isinstance(leaf, dict) else leaf)
        digits = max(100, len(Decimal(written).as_tuple().digits))
        with precision(digits):
            bits = certify(written, formula(params), digits)
        accuracies.append(bits)
        lower, upper, unit = written_bounds(written)
        for working in (digits + 50, digits + 100):
            with mp.workdps(working):
                value = independent_formula(params)
                if not mp.isfinite(value) or not mp_rational(lower) < value < mp_rational(upper):
                    raise AssertionError("independent formula disagrees: " + str(params))
        item = {"parameters": params, "number": written, "certified_accuracy_bits": bits, "independent_formula_dps": [digits + 50, digits + 100]}
        if params["p"] == "2" and int(params["n"]) in ADDED_DIMENSIONS:
            dimension = int(params["n"])
            with precision(digits + 30):
                item["gamma_free_radial_accuracy_bits"] = certify(written, p2_radial_certificate(dimension), digits)
            for working in (150, 200):
                with mp.workdps(working):
                    value = radial_quadrature(dimension)
                    if not mp.isfinite(value) or not mp_rational(lower) < value < mp_rational(upper):
                        raise AssertionError("independent radial quadrature disagrees")
            item["quadrature_dps"] = [150, 200]
        report["records"].append(item)
    report.update(entries=len(rows), minimum_certified_accuracy_bits=min(accuracies),
                  full_selection_matches=True, shifted_last_place_controls=True,
                  retained_precision_lowered=False, added_radial_certificates=4,
                  rigour="proven by Arb enclosure; mpmath agreement is supplementary")
    return report


class SobolevConstants(numberdb.Generator):
    table = "T92"
    parameters = PARAMETERS
    type = "R"
    digits = 100
    rigour = "proven"

    def enumerate(self):
        yield from selection()

    def value(self, params, digits):
        with precision(digits):
            value = formula(params)
            certify(ball_text(value, digits), value, digits)
            return wire(value)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--self-test", action="store_true")
    parser.add_argument("--compute", action="store_true")
    parser.add_argument("--document", type=Path, default=Path("table.yaml"))
    options = parser.parse_args()
    if options.self_test:
        print(json.dumps(controls(), indent=2))
    elif options.compute:
        generator = SobolevConstants()
        for params in generator.enumerate():
            print(json.dumps({"params": params, "number": to_text(generator.value(params, 100), 100)}))
    else:
        if options.document.suffix == ".json":
            document = json.loads(options.document.read_text())
        else:
            import yaml
            document = yaml.safe_load(options.document.read_text())
            if isinstance(document["Numbers"], list):
                nested = {}
                for row in document["Numbers"]:
                    params = row["params"]
                    leaf = {key: value for key, value in row.items() if key != "params"}
                    nested.setdefault(str(params["n"]), {}).setdefault(str(params["p"]), {})[str(params["q"])] = leaf
                document["Numbers"] = nested
        print(json.dumps(check_document(document), indent=2))


if __name__ == "__main__":
    main()