#!/usr/bin/env python3
"""Certify the global FKS--Yang absolute-psi decay constant 0.88178."""

from __future__ import annotations

from dataclasses import dataclass
from fractions import Fraction
from hashlib import sha256
from pathlib import Path

from flint import arb, ctx


if not __debug__:
    raise SystemExit("verify.py requires assertions")


K = 100_000
SOURCE_RADIUS = Fraction("4.862")
INTERIOR_RADIUS = Fraction("4.86201")
SOURCE_FKS_RADIUS = Fraction("5.5666305")
LOG_HANDOFF = Fraction(2008)
C1 = Fraction("17.4194")
C2 = Fraction("2.9089")
REMAINDER = Fraction(689_992_998, 100_000_000_000_000)
PUBLISHED_AMPLITUDE = Fraction("9.22022")
TARGET_AMPLITUDE = Fraction("9.2202181")
PUBLISHED_DECAY = Fraction("0.8476836")
TARGET_DECAY = Fraction("0.88178")

SOURCE_PINS = {
    "scratch/adjacent-unconditional--fks-corollary14-theta-amplitude-121-096/"
    "sources/fiori-kadiri-swidinsky-2204.02588v3.pdf":
        "15eff02042f2c1ab4878d029ce52a0abd6c95d66995b183a2dbc689ce9af49a4",
    "scratch/adjacent-unconditional--fks-corollary14-theta-amplitude-121-096/"
    "sources/fiori-kadiri-swidinsky-2204.02588v3.txt":
        "50fb049f04153165808d025db502c88679b853adabf889d8b03061336cd4afe5",
    "scratch/proof-system--c0060-yang-thesis-prior-art-repair/source/"
    "yang-thesis-2025.pdf":
        "b67b44e9608e3b73f261c4a4bd1ca70cc6000baee2aa43a4d749e44c5bd14d39",
    "scratch/proof-system--c0060-yang-thesis-prior-art-repair/source/"
    "yang-thesis-2025.txt":
        "189572ed43e57219c62499e65ef094efcb3034a02e542e1ee8a9529e5bfa51cc",
    "canon/claims/C-0053-fks-global-psi-amplitude-9-2202181.md":
        "fd3fd2d295477ba881a3d03db2c3fc5f7f867fdb161ab5aff7954427d1a85a30",
    "canon/witnesses/C-0053/verify.py":
        "d994f512fc20e7d754a24b2f8d094c12b42a8b1b4b1cf2369ad813053289b9c7",
}

FKS_TABLE = {
    (Fraction("5.573412"), 2500): Fraction("245.3670"),
    (Fraction("5.573412"), 5000): Fraction("204.0403"),
    (Fraction("5.573412"), 50_000): Fraction("138.1490"),
    (SOURCE_FKS_RADIUS, 2500): Fraction("245.6773"),
    (SOURCE_FKS_RADIUS, 5000): Fraction("204.2929"),
    (SOURCE_FKS_RADIUS, 50_000): Fraction("138.3136"),
    (Fraction("5.5"), 2500): Fraction("248.7907"),
    (Fraction("5.5"), 5000): Fraction("206.8264"),
    (Fraction("5.5"), 50_000): Fraction("139.9638"),
}


@dataclass(frozen=True)
class Certificate:
    table_gaps: tuple[tuple[Fraction, int, arb], ...]
    radius_padding: Fraction
    sigma2_at_handoff: arb
    asymptotic_amplitude: arb
    normalized_amplitude: arb
    asymptotic_decay: arb
    low_bridge: arb
    low_margin: arb
    tail_bridge: arb
    tail_margin: arb
    decay_improvement: Fraction
    amplitude_improvement: Fraction


def repo_root() -> Path:
    for parent in Path(__file__).resolve().parents:
        if (parent / "pyproject.toml").is_file() and (parent / "AGENTS.md").is_file():
            return parent
    raise AssertionError("repository root not found")


def pin_sources(root: Path) -> None:
    for relative, expected in SOURCE_PINS.items():
        path = root / relative
        assert path.is_file(), path
        assert sha256(path.read_bytes()).hexdigest() == expected, path


def as_arb(value: Fraction | int) -> arb:
    value = Fraction(value)
    return arb(value.numerator) / value.denominator


def fks_amplitude(log_x0: Fraction | int, radius: Fraction) -> arb:
    """Evaluate FKS (3.33) and (5.2) with directed Arb arithmetic."""
    log_x0_ball = as_arb(log_x0)
    radius_ball = as_arb(radius)
    c1_ball = as_arb(C1)
    c2_ball = as_arb(C2)

    ell = (log_x0_ball / radius_ball).sqrt()
    delta = 4 / (radius_ball * ell)
    exponent = 3 + delta
    ell_delta = ell**delta

    w = arb(K + 3) / K
    f_w = (16 * w / (3 * radius_ball)).exp() * w**exponent
    total = arb(0)

    for k in range(1, K):
        next_w = arb(K + 3 * (k + 1)) / K
        f_next = (
            (16 * next_w / (3 * radius_ball)).exp()
            * next_w**exponent
        )
        c1_step = c1_ball * (f_next - f_w)
        c2_step = c2_ball * (next_w**2 - w**2)
        weight = (-ell * (w - 1) ** 2 / w).exp()
        total += weight * (c1_step * ell_delta + c2_step / ell)
        w = next_w
        f_w = f_next

    w1 = arb(K + 3) / K
    initial = (
        c1_ball
        * ell_delta
        * (16 * w1 / (3 * radius_ball)).exp()
        * w1**exponent
        + c2_ball * w1**2 / ell
    )
    return 2 * (total + initial) + as_arb(REMAINDER)


def certify() -> Certificate:
    ctx.prec = 192
    pin_sources(repo_root())

    assert INTERIOR_RADIUS > SOURCE_RADIUS
    assert INTERIOR_RADIUS < Fraction("5.573412")
    assert LOG_HANDOFF > 1000
    radius_padding = INTERIOR_RADIUS - SOURCE_RADIUS
    assert radius_padding == Fraction(1, 100_000)

    table_gaps = []
    for (radius, log_x0), printed in FKS_TABLE.items():
        computed = fks_amplitude(log_x0, radius)
        gap = as_arb(printed) - computed
        assert gap > 0
        assert gap < arb("0.0011")
        table_gaps.append((radius, log_x0, gap))

    radius = as_arb(INTERIOR_RADIUS)
    log_handoff = as_arb(LOG_HANDOFF)
    sigma2_at_handoff = 1 - 2 / (radius * log_handoff).sqrt()
    assert sigma2_at_handoff > arb("0.9")
    assert sigma2_at_handoff < 1

    asymptotic_amplitude = fks_amplitude(LOG_HANDOFF, INTERIOR_RADIUS)
    assert asymptotic_amplitude > arb("306.40840")
    assert asymptotic_amplitude < arb("306.40841")

    normalized_amplitude = asymptotic_amplitude / (
        radius * radius.sqrt()
    )
    asymptotic_decay = 2 / radius.sqrt()
    assert normalized_amplitude > arb("28.58096")
    assert normalized_amplitude < arb("28.58097")
    assert asymptotic_decay > arb("0.9070308")
    assert asymptotic_decay < arb("0.9070310")

    target_amplitude = as_arb(TARGET_AMPLITUDE)
    target_decay = as_arb(TARGET_DECAY)
    published_decay = as_arb(PUBLISHED_DECAY)
    assert target_decay > published_decay
    assert asymptotic_decay > target_decay

    square_root_handoff = log_handoff.sqrt()
    low_bridge = 2 * (
        (target_decay - published_decay) * square_root_handoff
    ).exp()
    low_margin = target_amplitude - low_bridge
    assert low_margin > arb("0.0033")

    tail_bridge = normalized_amplitude * (
        -(asymptotic_decay - target_decay) * square_root_handoff
    ).exp()
    tail_margin = target_amplitude - tail_bridge
    assert tail_margin > arb("0.0015")

    decay_improvement = TARGET_DECAY - PUBLISHED_DECAY
    amplitude_improvement = PUBLISHED_AMPLITUDE - TARGET_AMPLITUDE
    assert decay_improvement == Fraction(85_241, 2_500_000)
    assert amplitude_improvement == Fraction(19, 10_000_000)

    return Certificate(
        table_gaps=tuple(table_gaps),
        radius_padding=radius_padding,
        sigma2_at_handoff=sigma2_at_handoff,
        asymptotic_amplitude=asymptotic_amplitude,
        normalized_amplitude=normalized_amplitude,
        asymptotic_decay=asymptotic_decay,
        low_bridge=low_bridge,
        low_margin=low_margin,
        tail_bridge=tail_bridge,
        tail_margin=tail_margin,
        decay_improvement=decay_improvement,
        amplitude_improvement=amplitude_improvement,
    )


def main() -> None:
    certificate = certify()
    print("source R:", float(SOURCE_RADIUS))
    print("interiorized R:", float(INTERIOR_RADIUS))
    print("R padding:", certificate.radius_padding)
    print("sigma2 at handoff:", certificate.sigma2_at_handoff)
    for radius, log_x0, gap in certificate.table_gaps:
        print(f"FKS table gap R={float(radius)} L={log_x0}:", gap)
    print("A(2008, 4.86201):", certificate.asymptotic_amplitude)
    print("A/R^(3/2):", certificate.normalized_amplitude)
    print("2/sqrt(R):", certificate.asymptotic_decay)
    print("low bridge coefficient:", certificate.low_bridge)
    print("low bridge margin:", certificate.low_margin)
    print("tail bridge coefficient:", certificate.tail_bridge)
    print("tail bridge margin:", certificate.tail_margin)
    print("published decay:", float(PUBLISHED_DECAY))
    print("improved decay:", float(TARGET_DECAY))
    print("exact decay increase:", certificate.decay_improvement)
    print("exact amplitude decrease:", certificate.amplitude_improvement)
    print("PASS: the global absolute-psi decay is 0.88178")


if __name__ == "__main__":
    main()
