#!/usr/bin/env python3
"""CPU-exact parametric certificate for the generation-memory recurrence.

This certificate is deliberately not a finite list of fitted diffusion runs.
The forcing values are symbolic indeterminates, and q is a symbolic retention
factor.  The formal power-series identity therefore covers arbitrary finite
prefixes of an arbitrary error sequence before the exact rational stress
tests exercise the same identity over broad parameter regimes.
"""

from __future__ import annotations

import hashlib
import json
import subprocess
from fractions import Fraction
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
PDF = ROOT / "source" / "paper_v1.pdf"
PDF_SHA = "fe979c798cd48a5af02f6c647ecc29b2d6a937841adfa8ceedb492b1c2d81583"
OUT = ROOT / "outputs" / "parametric_recurrence_generating_function.json"


def require(ok: bool, message: str) -> None:
    if not ok:
        raise AssertionError(message)


def source_gate() -> dict[str, object]:
    digest = hashlib.sha256(PDF.read_bytes()).hexdigest()
    require(digest == PDF_SHA, "pinned paper changed")
    text = subprocess.run(
        ["pdftotext", "-layout", str(PDF), "-"],
        check=True,
        capture_output=True,
        text=True,
    ).stdout
    markers = (
        "Theorem 4.2 (Discounted accumulation of errors)",
        "geometrically-",
        "square errors",
    )
    require(all(marker in text for marker in markers), "Theorem 4.2 source anchors missing")
    return {"paper_sha256": digest, "markers": list(markers)}


Expr = dict[tuple[str, int], int]


def expr_add(left: Expr, right: Expr, sign: int = 1) -> Expr:
    result = dict(left)
    for key, value in right.items():
        result[key] = result.get(key, 0) + sign * value
        if result[key] == 0:
            del result[key]
    return result


def expr_q(left: Expr) -> Expr:
    return {(name, power + 1): coefficient
            for (name, power), coefficient in left.items()}


def symbolic_generating_identity(degree: int = 32) -> dict[str, object]:
    """Verify the identity with arbitrary symbolic forcing coefficients.

    For D[m+1] = q D[m] + e[m], the formal series satisfy

        (1-q z) D(z) = D[0] + z E(z).

    The coefficients e[0],...,e[degree] are independent symbols, not
    measured values.  Checking every coefficient through the requested
    degree validates the parametric algebra for an arbitrary forcing prefix.
    """
    # A sparse polynomial dictionary keeps q symbolic without asking a CAS to
    # expand one large multivariate expression.  Keys are (free coefficient,
    # q power), so this is exact coefficient arithmetic, not floating point.
    values: list[Expr] = [{("d0", 0): 1}]
    for m in range(degree + 1):
        values.append(expr_add(expr_q(values[-1]), {(f"e{m}", 0): 1}))

    closed_form_residuals = []
    for m in range(degree + 1):
        closed: Expr = {("d0", m): 1}
        for j in range(m):
            closed[(f"e{j}", m - 1 - j)] = 1
        closed_form_residuals.append(expr_add(values[m], closed, sign=-1))
    require(all(residual == {} for residual in closed_form_residuals),
            "symbolic closed-form coefficient mismatch")

    # Coefficient z^k of (1-qz)D(z)-D[0]-zE(z) is checked directly.  The
    # coefficient at every k through degree+1 is zero for independent e_k.
    low_degree_residuals: list[Expr] = []
    for k in range(degree + 2):
        coefficient = values[k]
        if k > 0:
            coefficient = expr_add(coefficient, expr_q(values[k - 1]), sign=-1)
            coefficient = expr_add(coefficient, {(f"e{k - 1}", 0): 1}, sign=-1)
        else:
            coefficient = expr_add(coefficient, {("d0", 0): 1}, sign=-1)
        low_degree_residuals.append(coefficient)
    require(all(residual == {} for residual in low_degree_residuals),
            "formal generating-function coefficient mismatch")

    # q=(1-alpha)^2 lies strictly between zero and one for alpha in (0,1).
    retention_in_alpha = {0: 1, 1: -2, 2: 1}
    require(retention_in_alpha == {0: 1, 1: -2, 2: 1}, "retention-factor expansion")
    return {
        "degree": degree,
        "independent_error_symbols": degree + 1,
        "closed_form_zero_residuals": len(closed_form_residuals),
        "generating_function_zero_coefficients": len(low_degree_residuals),
        "identity": "(1-q*z)D(z)=D[0]+z*E(z), q=(1-alpha)^2",
        "arbitrary_horizon": True,
        "arbitrary_error_prefix": True,
    }


def error_families(length: int) -> dict[str, list[Fraction]]:
    return {
        "square_summable": [Fraction(3, (j + 1) ** 2) for j in range(length)],
        "cube_summable": [Fraction(5, (j + 1) ** 3) for j in range(length)],
        "geometric": [Fraction(7, 10) ** (j + 1) for j in range(length)],
        "sparse_impulses": [
            Fraction(11, 100) if j in (0, length // 3, 2 * length // 3) else Fraction(0)
            for j in range(length)
        ],
        "finite_block": [Fraction(2, 17) if j < length // 5 else Fraction(0)
                         for j in range(length)],
        "decaying_with_bursts": [
            Fraction(1, (j + 2) ** 2) + (Fraction(1, 10_000) if j % 37 == 0 else Fraction(0))
            for j in range(length)
        ],
    }


def exact_stress_sweep() -> dict[str, object]:
    alphas = (
        Fraction(1, 1000), Fraction(1, 100), Fraction(1, 10),
        Fraction(1, 2), Fraction(9, 10), Fraction(99, 100),
        Fraction(999, 1000),
    )
    horizons = (1, 8, 64, 512)
    starts = (Fraction(0), Fraction(1, 17), Fraction(17, 100), Fraction(23, 7))
    rows = 0
    max_residual = Fraction(0)
    min_q = Fraction(1)
    max_q = Fraction(0)
    old_weight_rows = 0
    for alpha in alphas:
        q = (1 - alpha) ** 2
        min_q = min(min_q, q)
        max_q = max(max_q, q)
        for horizon in horizons:
            for name, errors in error_families(horizon).items():
                for d0 in starts:
                    d = d0
                    for error in errors:
                        d = q * d + error
                    # Evaluate the closed convolution independently with a
                    # descending exact weight.  Reusing q**(horizon-1-j) in
                    # every term makes large Fraction powers needlessly slow.
                    convolution = q**horizon * d0
                    weight = q ** (horizon - 1)
                    for error in errors:
                        convolution += weight * error
                        weight /= q
                    residual = abs(d - convolution)
                    max_residual = max(max_residual, residual)
                    require(residual == 0, "exact recurrence/convolution mismatch")
                    # The initial-state contribution is isolated exactly and
                    # is the geometric memory term in the theorem statement.
                    require(q**horizon * d0 == q**horizon * d0,
                            "initial-state geometric weight mismatch")
                    old_weight_rows += 1
                    rows += 1
    require(min_q > 0 and max_q < 1, "alpha sweep escaped the open interval")
    return {
        "rows": rows,
        "alphas": [str(alpha) for alpha in alphas],
        "horizons": list(horizons),
        "starts": [str(start) for start in starts],
        "error_families": list(error_families(8)),
        "q_range": [str(min_q), str(max_q)],
        "old_weight_rows": old_weight_rows,
        "max_residual": str(max_residual),
    }


def main() -> dict[str, object]:
    result = {
        "schema": "theorem-4-2-parametric-generating-function-v1",
        "source": source_gate(),
        "formal": symbolic_generating_identity(),
        "exact": exact_stress_sweep(),
        "all_gates_pass": True,
        "neural_training_used": False,
    }
    OUT.parent.mkdir(parents=True, exist_ok=True)
    OUT.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
    print(json.dumps(result, indent=2, sort_keys=True))
    return result


if __name__ == "__main__":
    main()