#!/usr/bin/env python3
"""Exact algebra for the prescribed-jet polynomial cycle construction.

This program checks polynomial identities over the rational numbers and
symbolic parameters.  It neither integrates trajectories nor establishes a
numerical threshold for the analytic small-rho realization theorem.
"""

from __future__ import annotations

import argparse
import csv
import json
import platform
from pathlib import Path

import sympy as sp


U, H = sp.symbols("u h", real=True)
Z = sp.symbols("z", positive=True)
A, M, THETA, C, DELTA, RHO = sp.symbols(
    "a m theta c delta rho", positive=True
)
CASES = ((1, 2), (2, 2), (3, 2), (4, 2), (3, 4), (3, 5))


def polynomial_truncate(expression: sp.Expr, order: int) -> sp.Expr:
    """Retain u-powers 0 through order, using exact polynomial arithmetic."""
    return sp.Add(
        *(
            coefficient * U**power[0]
            for power, coefficient in sp.Poly(sp.expand(expression), U).terms()
            if power[0] <= order
        )
    )


def truncated_power(base: sp.Expr, exponent: int, order: int) -> sp.Expr:
    """Integer power in the quotient ring Q[u]/(u**(order+1))."""
    result = sp.Integer(1)
    for _ in range(exponent):
        result = polynomial_truncate(result * base, order)
    return result


def analytic_coordinate_factor(order: int) -> sp.Expr:
    """Taylor polynomial of 2*(u-log(1+u))/u**2, whose value at 0 is 1."""
    return sp.Add(
        *(sp.Rational(2 * (-1) ** power, power + 2) * U**power
          for power in range(order + 1))
    )


def basis_polynomial(d: int, degree: int) -> sp.Expr:
    """Return Q for m=1: T_(degree-1)[(u*sqrt(G(u)))**d].

    Construction uses a generalized binomial series.  Verification below
    independently squares this jet and compares it with the integer power
    of the logarithmic germ; the leading coefficient fixes its square root.
    """
    remainder_order = degree - 1 - d
    factor = analytic_coordinate_factor(remainder_order)
    increment = factor - 1
    power = sp.Integer(1)
    result = sp.Integer(1)
    for index in range(1, remainder_order + 1):
        power = polynomial_truncate(power * increment, remainder_order)
        result += sp.binomial(sp.Rational(d, 2), index) * power
    return sp.expand(U**d * polynomial_truncate(result, remainder_order))


def moment_coefficient(index: int) -> sp.Rational:
    """C_index divided by its common factor pi/sqrt(c)."""
    return sp.binomial(2 * index + 2, index + 1) / sp.Integer(2) ** index


def terms(expression: sp.Expr) -> list[dict[str, object]]:
    """JSON representation without rounded coefficients."""
    return [
        {"power": power[0], "coefficient": str(coefficient)}
        for power, coefficient in reversed(sp.Poly(sp.expand(expression), U).terms())
        if coefficient != 0
    ]


def record(checks: dict[str, bool], name: str, condition: bool) -> None:
    value = bool(condition)
    checks[name] = value
    if not value:
        raise AssertionError(name)


def verify_case(number: int, jet_order: int) -> dict[str, object]:
    q = (jet_order + 1) // 2
    degree = 2 * (number + q + 1)
    roots = tuple(sp.Integer(i) for i in range(1, number + 1))
    exponents = tuple(2 * q + 2 * j + 1 for j in range(number + 1))
    base = C + THETA * M * sp.Add(
        *(sp.Rational((-1) ** (n + 1), n) * U**n
          for n in range(2, degree + 1))
    )
    center = C - THETA * M * (U - sp.log(1 + U))
    target = sp.expand(
        -moment_coefficient(q + number)
        * H ** (q + 1)
        * sp.prod(H - root for root in roots)
    )
    coefficients = tuple(
        sp.factor(target.coeff(H, q + j + 1) / moment_coefficient(q + j))
        for j in range(number + 1)
    )
    unit_basis = tuple(basis_polynomial(d, degree) for d in exponents)
    general_basis = tuple(
        M ** sp.Rational(d, 2) * polynomial
        for d, polynomial in zip(exponents, unit_basis, strict=True)
    )
    perturbation = DELTA * sp.Add(
        *(coefficients[j] * RHO ** (2 * (number - j)) * general_basis[j]
          for j in range(number + 1))
    )
    family = sp.expand(base + perturbation)
    checks: dict[str, bool] = {}

    record(checks, "degree_is_even", degree % 2 == 0)
    record(checks, "exact_degree", sp.Poly(family, U).degree() == degree)
    record(checks, "leading_coefficient",
           sp.simplify(family.coeff(U, degree) + THETA * M / degree) == 0)
    record(checks, "degree_formula",
           degree == 2 * number + 2 * q + 2)
    record(checks, "linear_degree_bound", degree <= 2 * number + jet_order + 3)
    record(checks, "first_perturbation_order_above_K", exponents[0] > jet_order)
    record(checks, "last_basis_exact_monomial", unit_basis[-1] == U ** (degree - 1))
    record(checks, "b_N_equals_minus_one", coefficients[-1] == -1)
    record(checks, "b_0_nonzero", coefficients[0] != 0)
    record(checks, "b_0_parity_sign", sp.sign(coefficients[0]) == (-1) ** (number + 1))

    jet_residual = [
        sp.simplify(sp.diff(family - center, U, n).subs(U, 0))
        for n in range(jet_order + 1)
    ]
    record(checks, "K_jet_matches_center", all(value == 0 for value in jet_residual))
    base_jet_residual = [
        sp.simplify(sp.diff(base - center, U, n).subs(U, 0))
        for n in range(degree + 1)
    ]
    record(checks, "base_D_jet_matches_center",
           all(value == 0 for value in base_jet_residual))

    basis_data = []
    for j, (d, polynomial) in enumerate(zip(exponents, unit_basis, strict=True)):
        lowest_power = min(power[0] for power in sp.Poly(polynomial, U).monoms())
        record(checks, f"Q{j}_vanishing_order", lowest_power == d)
        record(checks, f"Q{j}_leading_germ_coefficient", polynomial.coeff(U, d) == 1)
        record(checks, f"Q{j}_degree_at_most_D_minus_one",
               sp.Poly(polynomial, U).degree() <= degree - 1)

        remaining = degree - 1 - d
        factor = analytic_coordinate_factor(remaining)
        normalized_q = sp.cancel(polynomial / U**d)
        squared_residual = polynomial_truncate(
            normalized_q**2 - truncated_power(factor, d, remaining), remaining
        )
        record(checks, f"Q{j}_squared_log_germ_identity", squared_residual == 0)
        basis_data.append({
            "j": j,
            "d_j": d,
            "Q_j_at_m_1": str(polynomial),
            "Q_j_general_scaling": f"m**({d}/2) times Q_j_at_m_1",
            "coefficients": terms(polynomial),
            "polynomial_degree": sp.Poly(polynomial, U).degree(),
            "vanishing_order": lowest_power,
            "squared_germ_difference_order_at_least": degree + d,
        })

    constructed_moment = sp.expand(sp.Add(
        *(coefficients[j] * moment_coefficient(q + j) * H ** (q + j + 1)
          for j in range(number + 1))
    ))
    record(checks, "harmonic_moment_target_identity", constructed_moment == target)
    root_slopes = []
    for i, root in enumerate(roots, start=1):
        slope = sp.factor(sp.diff(target, H).subs(H, root))
        record(checks, f"target_root_{i}_is_zero", target.subs(H, root) == 0)
        record(checks, f"target_root_{i}_is_simple", slope != 0)
        record(checks, f"target_root_{i}_slope_sign",
               sp.sign(slope) == (-1) ** (number - i + 1))
        root_slopes.append(str(slope))
    for index in range(q, q + number + 1):
        # Integral of cos(theta)^(2*index+2) divided by pi.
        cosine_integral = (
            2 * sp.factorial(2 * index + 2)
            / (sp.Integer(4) ** (index + 1) * sp.factorial(index + 1) ** 2)
        )
        record(checks, f"oscillator_moment_C{index}",
               sp.Integer(2) ** (index + 1) * cosine_integral
               == moment_coefficient(index))

    curvature = sp.expand(-sp.diff(base, U, 2) / (THETA * M))
    curvature_sum = sp.Add(
        *((j + 1) * (-U) ** j for j in range(degree - 1))
    )
    curvature_numerator = 1 + degree * U ** (degree - 1) + (degree - 1) * U**degree
    record(checks, "Taylor_curvature_sum", sp.expand(curvature - curvature_sum) == 0)
    record(checks, "Taylor_curvature_rational_identity",
           sp.expand((1 + U) ** 2 * curvature - curvature_numerator) == 0)
    record(checks, "Taylor_remainder_derivative_identity",
           sp.simplify((1 + U) * sp.diff(base - center, U) + THETA * M * U**degree) == 0)
    record(checks, "Taylor_remainder_zero_at_origin", (base - center).subs(U, 0) == 0)

    normalized_base = sp.expand(base.subs({C: 3, M: 1, THETA: 1}))
    normalized_family = sp.expand(family.subs({C: 3, M: 1, THETA: 1}))
    endpoint = sp.factor(normalized_base.subs(U, -1))
    endpoint_formula = 4 - sp.harmonic(degree)
    record(checks, "base_endpoint_harmonic_identity", endpoint == endpoint_formula)
    record(checks, "base_endpoint_positive_at_c3_theta1_m1", endpoint > 0)

    # The anchor b=a*(1+z) is arbitrary here.  If P_0(b)=0, the verified
    # anchor identity makes it the corrected family's fixed zero.
    limit_family = family.subs(RHO, 0)
    endpoint_error = sp.expand(family - limit_family)
    error_left = endpoint_error.subs(U, -1)
    error_right = endpoint_error.subs(U, Z)
    t1 = sp.cancel((error_left - error_right / Z ** (degree - 2)) / (1 + Z))
    t0 = sp.cancel(-(Z * error_left + error_right / Z ** (degree - 2)) / (1 + Z))
    correction = t0 * U ** (degree - 2) + t1 * U ** (degree - 1)
    corrected_family = family + correction
    correction_polynomial = sp.Poly(correction, U)
    corrected_polynomial = sp.Poly(corrected_family, U)
    determinant = sp.Matrix(((1, -1), (Z ** (degree - 2), Z ** (degree - 1)))).det()
    record(checks, "endpoint_correction_matrix_determinant",
           sp.expand(determinant - Z ** (degree - 2) * (1 + Z)) == 0)
    record(checks, "corrected_left_endpoint_identity", sp.cancel(error_left + t0 - t1) == 0)
    record(checks, "corrected_right_anchor_identity",
           sp.cancel(error_right + t0 * Z ** (degree - 2) + t1 * Z ** (degree - 1)) == 0)
    for index, coefficient in enumerate((t0, t1)):
        rho_polynomial = sp.Poly(coefficient, RHO)
        record(checks, f"t{index}_is_polynomial_in_rho", rho_polynomial.is_univariate)
        record(checks, f"t{index}_divisible_by_rho_squared",
               rho_polynomial.nth(0) == 0 and rho_polynomial.nth(1) == 0)
    record(checks, "correction_vanishes_to_order_at_least_D_minus_two",
           all(correction_polynomial.nth(power) == 0 for power in range(degree - 2)))
    record(checks, "corrected_exact_degree", corrected_polynomial.degree() == degree)
    record(checks, "corrected_leading_coefficient",
           sp.simplify(corrected_polynomial.nth(degree) + THETA * M / degree) == 0)
    record(checks, "corrected_K_jet_matches_center",
           all(correction_polynomial.nth(power) == 0 for power in range(jet_order + 1))
           and all(value == 0 for value in jet_residual))

    return {
        "N": number,
        "K": jet_order,
        "q": q,
        "D": degree,
        "d_j": list(exponents),
        "h_i": [str(root) for root in roots],
        "b_j": [str(value) for value in coefficients],
        "A_D_general": str(base),
        "A_D_at_c3_theta1_m1": str(normalized_base),
        "alpha_rho_general": str(family),
        "alpha_rho_at_c3_theta1_m1": str(normalized_family),
        "leading_coefficient_in_u": str(-THETA * M / degree),
        "leading_coefficient_in_x": str(-THETA * M / (degree * A**degree)),
        "base_endpoint_at_c3_theta1_m1": str(endpoint),
        "K_jet_derivative_residuals_in_u": [str(value) for value in jet_residual],
        "basis": basis_data,
        "moment_common_factor": "pi/sqrt(c)",
        "moment_target_divided_by_common_factor": str(target),
        "moment_target_factored_divided_by_common_factor": str(sp.factor(target)),
        "moment_target_root_slopes_divided_by_common_factor": root_slopes,
        "normalized_curvature_polynomial": str(curvature),
        "curvature_rational_numerator": str(curvature_numerator),
        "curvature_rational_denominator": "(1+u)**2",
        "endpoint_correction": {
            "anchor": "b=a*(1+z), with z>0",
            "E_rho": str(endpoint_error),
            "E_rho_at_x0": str(error_left),
            "E_rho_at_b": str(error_right),
            "t_0": str(t0),
            "t_1": str(t1),
            "correction": "t_0*u**(D-2)+t_1*u**(D-1)",
            "correction_order_at_least": degree - 2,
            "corrected_degree": corrected_polynomial.degree(),
            "fixed_left_value": str(limit_family.subs(U, -1)),
            "fixed_anchor_value": str(limit_family.subs(U, Z)),
            "anchor_value_scope": "The identity is P_hat_rho(b)=P_0(b). It equals zero when b is chosen as a root of P_0 in the analytic construction.",
        },
        "checks": checks,
        "all_checks_passed": all(checks.values()),
    }


def worked_example(case: dict[str, object]) -> dict[str, object]:
    expected_q0 = U**3 - U**4 + sp.Rational(11, 12) * U**5
    expected_q1 = U**5
    actual_q0 = sp.sympify(case["basis"][0]["Q_j_at_m_1"], locals={"u": U})
    actual_q1 = sp.sympify(case["basis"][1]["Q_j_at_m_1"], locals={"u": U})
    checks: dict[str, bool] = {}
    record(checks, "Q0_matches_requested_rational_example", actual_q0 == expected_q0)
    record(checks, "Q1_matches_requested_rational_example", actual_q1 == expected_q1)
    record(checks, "b0_equals_5_over_3", case["b_j"][0] == "5/3")
    record(checks, "b1_equals_minus_one", case["b_j"][1] == "-1")
    record(checks, "base_endpoint_equals_31_over_20",
           case["base_endpoint_at_c3_theta1_m1"] == "31/20")
    base = 3 - U**2 / 2 + U**3 / 3 - U**4 / 4 + U**5 / 5 - U**6 / 6
    family = sp.expand(base + DELTA * (sp.Rational(5, 3) * RHO**2 * expected_q0 - expected_q1))
    expected_t0 = 5 * DELTA * RHO**2 * (2 * Z - 1) / (3 * Z)
    expected_t1 = -5 * DELTA * RHO**2 * (11 * Z + 12) / (36 * Z)
    symbol_map = {"a": A, "m": M, "theta": THETA, "c": C, "delta": DELTA, "rho": RHO, "z": Z}
    actual_t0 = sp.sympify(case["endpoint_correction"]["t_0"], locals=symbol_map).subs({M: 1, C: 3, THETA: 1})
    actual_t1 = sp.sympify(case["endpoint_correction"]["t_1"], locals=symbol_map).subs({M: 1, C: 3, THETA: 1})
    record(checks, "t0_matches_rational_correction", sp.cancel(actual_t0 - expected_t0) == 0)
    record(checks, "t1_matches_rational_correction", sp.cancel(actual_t1 - expected_t1) == 0)
    corrected_left = sp.cancel((family + expected_t0 * U**4 + expected_t1 * U**5).subs(U, -1))
    record(checks, "corrected_left_endpoint_equals_31_over_20_plus_delta",
           corrected_left == sp.Rational(31, 20) + DELTA)
    return {
        "parameters": {"N": 1, "K": 2, "a": "1", "m": "1", "theta": "1", "c": "3", "h_1": "1"},
        "coordinate": "u=x-1",
        "D": 6,
        "Q_0": str(expected_q0),
        "Q_1": str(expected_q1),
        "b_0": "5/3",
        "b_1": "-1",
        "A_6": str(base),
        "alpha_rho": str(family),
        "alpha_rho_coefficients_in_u": terms(family),
        "alpha_rho_at_x0": str(sp.factor(family.subs(U, -1))),
        "endpoint_correction_t0": str(expected_t0),
        "endpoint_correction_t1": str(expected_t1),
        "corrected_alpha_rho_at_x0": str(corrected_left),
        "M_divided_by_pi_over_sqrt3": "-5*h**2*(h-1)",
        "M_prime_at_h1_divided_by_pi_over_sqrt3": "-5",
        "formal_section_location": "s=rho*sqrt(2)+O(rho**2)",
        "formal_multiplier": "1-5*pi*delta*rho**4/sqrt(3)+O(rho**5)",
        "asymptotic_scope": "The location and multiplier are analytic-theorem consequences as rho tends to zero; this script does not certify any fixed positive rho or delta.",
        "checks": checks,
        "all_checks_passed": all(checks.values()),
    }


def save_json(path: Path, value: object) -> None:
    path.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="utf-8")


def draw_degree_grid(output: Path) -> None:
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    import numpy as np
    from matplotlib.patches import Rectangle

    n_values = tuple(range(1, 7))
    k_values = tuple(range(2, 9))
    matrix = np.array([[2 * (n + (k + 1) // 2 + 1) for k in k_values] for n in n_values])
    figure, axis = plt.subplots(figsize=(7.1, 4.9), layout="constrained")
    chart = axis.imshow(matrix, cmap="Blues", vmin=4, vmax=24, aspect="auto", origin="lower")
    for row, n in enumerate(n_values):
        for column, k in enumerate(k_values):
            value = matrix[row, column]
            axis.text(column, row, str(value), ha="center", va="center",
                      color="white" if value >= 17 else "#17324d", fontsize=12)
            if (n, k) in CASES:
                axis.add_patch(Rectangle((column - 0.44, row - 0.42), 0.88, 0.84,
                                         fill=False, edgecolor="#d27617", linewidth=1.8))
    axis.set_xticks(range(len(k_values)), k_values)
    axis.set_yticks(range(len(n_values)), n_values)
    axis.set_xlabel("Preserved jet order K", labelpad=9)
    axis.set_ylabel("Prescribed cycle number N", labelpad=9)
    axis.set_title("Explicit construction degree\n$D=2(N+\\lceil K/2\\rceil+1)$", pad=13)
    for spine in axis.spines.values():
        spine.set_visible(False)
    axis.tick_params(length=0)
    colorbar = figure.colorbar(chart, ax=axis, shrink=0.82, pad=0.04)
    colorbar.set_ticks((6, 10, 14, 18, 22))
    colorbar.set_label("Polynomial degree D")
    figure.text(0.5, -0.025, "Orange boxes: cases checked by exact algebra. The displayed degrees are sufficient, not claimed minimal.",
                ha="center", fontsize=9, color="#36495c")
    figure.savefig(output / "degree_grid.pdf", bbox_inches="tight")
    figure.savefig(output / "degree_grid.png", dpi=200, bbox_inches="tight")
    plt.close(figure)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output-dir", type=Path, default=Path(__file__).resolve().parent / "outputs")
    parser.add_argument("--no-plot", action="store_true", help="Run exact algebra without importing matplotlib.")
    arguments = parser.parse_args()
    output = arguments.output_dir
    output.mkdir(parents=True, exist_ok=True)
    reports = [verify_case(n, k) for n, k in CASES]
    for report in reports:
        save_json(output / f"case_N{report['N']}_K{report['K']}.json", report)
    worked = worked_example(reports[0])
    save_json(output / "worked_degree6.json", worked)
    with (output / "degree_table.csv").open("w", newline="", encoding="utf-8") as handle:
        writer = csv.writer(handle)
        writer.writerow(("N", "K", "q", "D", "d_0", "d_N", "algebraic_checks", "all_passed"))
        for report in reports:
            writer.writerow((report["N"], report["K"], report["q"], report["D"],
                             report["d_j"][0], report["d_j"][-1], len(report["checks"]),
                             report["all_checks_passed"]))
    if not arguments.no_plot:
        draw_degree_grid(output)
    summary = {
        "schema_version": 1,
        "scope": "Exact algebraic identities; no numerical integration and no fixed-parameter cycle certification.",
        "normalization": "alpha_center=c-theta*m*(u-log(1+u))=c-theta*s**2/2; u=(x-a)/a.",
        "versions": {"python": platform.python_version(), "sympy": sp.__version__},
        "cases": [
            {key: report[key] for key in ("N", "K", "q", "D", "d_j", "b_j", "all_checks_passed")}
            | {"checks_passed": len(report["checks"])}
            for report in reports
        ],
        "worked_example_checks_passed": len(worked["checks"]),
        "total_checks_passed": sum(len(report["checks"]) for report in reports) + len(worked["checks"]),
        "all_checks_passed": all(report["all_checks_passed"] for report in reports) and worked["all_checks_passed"],
        "figure_generated": not arguments.no_plot,
    }
    save_json(output / "summary.json", summary)
    print(json.dumps(summary, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()