#!/usr/bin/env python3
"""Exact audit of WKB spectral-cover bookkeeping.

The quadratic differential is

    phi = R(z) dz**2,

and its normalized spectral cover carries a one-form lambda satisfying
lambda**2 = phi.  This script checks only algebraic and local-order claims:

1. the normalized local models above monomial zeros and poles;
2. the coordinate change at infinity for the Airy and Weber examples;
3. discriminant, multiple-root, and coalescence data for
   R(z; E) = z**3 - 3*z - E;
4. a double zero whose singular affine model separates into two
   unramified points after normalization; and
5. branch-count parity, including infinity, for polynomial and rational
   examples, with a split global-square cover kept separate.

Every check raises an explicit RuntimeError.  Consequently, ``python -O``
performs the same audit.  No analytic continuation, WKB summability, or
global period claim is inferred from these symbolic calculations.
"""

from __future__ import annotations

import platform

import sympy as sp


MONOMIAL_ORDER_MIN = -6
MONOMIAL_ORDER_MAX = 5
MONOMIAL_ORDERS = tuple(
    range(MONOMIAL_ORDER_MIN, MONOMIAL_ORDER_MAX + 1)
)


def reduced(expression: sp.Expr) -> sp.Expr:
    """Return a stable exact form for a rational expression."""
    return sp.factor(sp.cancel(sp.together(expression)))


def require(condition: bool, label: str, detail: object = "") -> None:
    """Raise a visible failure instead of relying on an assert statement."""
    if not condition:
        suffix = f": {detail}" if detail != "" else ""
        raise RuntimeError(f"{label} failed{suffix}")


def require_zero(expression: sp.Expr, label: str) -> None:
    """Raise unless an exact symbolic expression vanishes."""
    residual = reduced(expression)
    if residual != 0:
        residual = sp.simplify(residual)
    if residual != 0:
        raise RuntimeError(f"{label} failed: residual = {residual}")


def require_nonzero(expression: sp.Expr, label: str) -> None:
    """Raise if an expression expected to be nonzero vanishes exactly."""
    value = reduced(expression)
    if value == 0:
        raise RuntimeError(f"{label} failed: expression is zero")


def polynomial_zero_order(poly: sp.Poly) -> int:
    """Return the multiplicity of zero as a root of a nonzero polynomial."""
    if poly.is_zero:
        raise RuntimeError("zero polynomial has no finite vanishing order")
    exponents = [monomial[0] for monomial, _ in poly.terms()]
    return min(exponents)


def order_at_zero(expression: sp.Expr, variable: sp.Symbol) -> int:
    """Return the exact Laurent order of a rational expression at zero."""
    numerator, denominator = sp.fraction(sp.cancel(expression))
    try:
        numerator_poly = sp.Poly(numerator, variable)
        denominator_poly = sp.Poly(denominator, variable)
    except sp.PolynomialError as exc:
        raise RuntimeError(
            f"cannot determine the rational order of {expression}"
        ) from exc
    return (
        polynomial_zero_order(numerator_poly)
        - polynomial_zero_order(denominator_poly)
    )


def pole_order(order: int) -> int:
    """Convert a divisor order into the nonnegative pole order."""
    return max(-order, 0)


def monomial_local_audit() -> list[dict[str, object]]:
    """Check normalized local covers for phi = z**m dz**2."""
    z, xi = sp.symbols("z xi")
    rows: list[dict[str, object]] = []

    require(
        MONOMIAL_ORDER_MIN <= -6 and MONOMIAL_ORDER_MAX >= 5,
        "declared monomial range includes -6..5",
        MONOMIAL_ORDERS,
    )

    for m in MONOMIAL_ORDERS:
        is_branched = bool(m % 2)
        if is_branched:
            z_of_xi = xi**2
            jacobian = sp.diff(z_of_xi, xi)
            pulled_phi = reduced(
                z_of_xi**m * jacobian**2
            )
            lambda_coefficient = 2 * xi ** (m + 1)
            require_zero(
                lambda_coefficient**2 - pulled_phi,
                f"odd monomial m={m}: lambda squared",
            )
            require(
                order_at_zero(z_of_xi, xi) == 2,
                f"odd monomial m={m}: ramification index",
            )
            require_zero(
                -lambda_coefficient.subs(xi, -xi)
                + lambda_coefficient,
                f"odd monomial m={m}: deck involution",
            )
            lambda_order = order_at_zero(lambda_coefficient, xi)
            require(
                lambda_order == m + 1,
                f"odd monomial m={m}: lambda order",
                lambda_order,
            )
            points_above = 1
        else:
            exponent = m // 2
            lambda_plus = z**exponent
            lambda_minus = -lambda_plus
            require_zero(
                lambda_plus**2 - z**m,
                f"even monomial m={m}: plus sheet",
            )
            require_zero(
                lambda_minus**2 - z**m,
                f"even monomial m={m}: minus sheet",
            )
            require_nonzero(
                lambda_plus - lambda_minus,
                f"even monomial m={m}: distinct signs",
            )
            plus_order = order_at_zero(lambda_plus, z)
            minus_order = order_at_zero(lambda_minus, z)
            require(
                plus_order == exponent and minus_order == exponent,
                f"even monomial m={m}: sheet orders",
                (plus_order, minus_order),
            )
            lambda_order = exponent
            points_above = 2

        rows.append(
            {
                "m": m,
                "branched": is_branched,
                "points_above": points_above,
                "lambda_order": lambda_order,
                "pole_order": pole_order(lambda_order),
            }
        )
    return rows


def infinity_coefficient(
    radicand: sp.Expr,
    z: sp.Symbol,
    xi: sp.Symbol,
) -> sp.Expr:
    """Return the coefficient of phi after z = 1/xi."""
    transformed = radicand.subs(z, 1 / xi) * sp.diff(1 / xi, xi) ** 2
    displayed = radicand.subs(z, 1 / xi) * xi**-4
    require_zero(
        transformed - displayed,
        "quadratic-differential transformation at infinity",
    )
    return reduced(transformed)


def connected_genus(branch_count: int, label: str) -> int:
    """Apply Riemann--Hurwitz to a connected double cover of P1."""
    require(branch_count >= 2, f"{label}: enough branch points")
    require(branch_count % 2 == 0, f"{label}: branch-count parity")
    genus_numerator = branch_count - 2
    require(
        genus_numerator % 2 == 0,
        f"{label}: integral Riemann--Hurwitz genus",
    )
    genus = genus_numerator // 2
    require(
        2 * genus - 2 == -4 + branch_count,
        f"{label}: Riemann--Hurwitz identity",
    )
    return genus


def infinity_examples_audit() -> list[dict[str, object]]:
    """Check Airy and Weber finite and infinite divisor data."""
    z, xi = sp.symbols("z xi")
    a = sp.symbols("a", nonzero=True)
    examples = (
        ("Airy", z, ((sp.Integer(0), 1),), 5, 2),
        ("Weber", z**2 - a**2, ((-a, 1), (a, 1)), 6, 2),
    )
    rows: list[dict[str, object]] = []

    require(a.is_nonzero is True, "Weber parameter assumption a != 0")
    for label, radicand, roots, expected_pole, expected_branches in examples:
        transformed = infinity_coefficient(radicand, z, xi)
        degree = sp.Poly(radicand, z).degree()
        infinity_order = order_at_zero(transformed, xi)
        require(
            infinity_order == -degree - 4,
            f"{label}: infinity order",
            infinity_order,
        )
        require(
            pole_order(infinity_order) == expected_pole,
            f"{label}: infinity pole order",
            pole_order(infinity_order),
        )

        finite_branches = 0
        for root, multiplicity in roots:
            shifted = radicand.subs(z, root + xi)
            actual_multiplicity = order_at_zero(shifted, xi)
            require(
                actual_multiplicity == multiplicity,
                f"{label}: zero multiplicity at z={root}",
                actual_multiplicity,
            )
            finite_branches += multiplicity % 2

        infinity_branch = abs(infinity_order) % 2
        branch_count = finite_branches + infinity_branch
        require(
            branch_count == expected_branches,
            f"{label}: total branch count",
            branch_count,
        )
        require(
            branch_count > 0,
            f"{label}: odd valuation proves connectedness",
        )
        genus = connected_genus(branch_count, label)
        require(genus == 0, f"{label}: genus zero", genus)
        rows.append(
            {
                "label": label,
                "radicand": radicand,
                "transformed": transformed,
                "finite_simple_zeros": len(roots),
                "infinity_order": infinity_order,
                "infinity_pole_order": pole_order(infinity_order),
                "branch_count": branch_count,
                "connected": True,
                "genus": genus,
            }
        )
    return rows


def cubic_family_audit() -> dict[str, object]:
    """Check the cubic discriminant and its two double-root fibers."""
    z, energy, u, epsilon = sp.symbols("z E u epsilon")
    potential = z**3 - 3 * z
    radicand = potential - energy
    derivative = sp.diff(potential, z)

    discriminant = sp.factor(sp.discriminant(radicand, z))
    expected_discriminant = 27 * (4 - energy**2)
    require_zero(
        discriminant - expected_discriminant,
        "cubic discriminant",
    )
    discriminant_roots = sp.Poly(discriminant, energy).all_roots()
    require(
        set(discriminant_roots) == {sp.Integer(-2), sp.Integer(2)},
        "cubic critical energies",
        discriminant_roots,
    )

    rational_function_field = sp.QQ.frac_field(energy)
    generic_poly = sp.Poly(
        radicand,
        z,
        domain=rational_function_field,
    )
    generic_derivative = sp.Poly(
        sp.diff(radicand, z),
        z,
        domain=rational_function_field,
    )
    generic_gcd = sp.gcd(generic_poly, generic_derivative).monic()
    require(
        generic_gcd.degree() == 0,
        "generic cubic is square-free over Q(E)",
        generic_gcd.as_expr(),
    )

    z_turning = sp.Function("z_t")(energy)
    implicit_equation = potential.subs(z, z_turning) - energy
    implicit_derivative = sp.diff(implicit_equation, energy)
    derivative_formula = 1 / derivative.subs(z, z_turning)
    require_zero(
        implicit_derivative.subs(
            sp.diff(z_turning, energy),
            derivative_formula,
        ),
        "implicit turning-point derivative",
    )

    fibers: list[dict[str, object]] = []
    critical_pairs = ((sp.Integer(-1), sp.Integer(2)),
                      (sp.Integer(1), sp.Integer(-2)))
    for critical_point, critical_energy in critical_pairs:
        require_zero(
            derivative.subs(z, critical_point),
            f"critical point z={critical_point}",
        )
        require_zero(
            potential.subs(z, critical_point) - critical_energy,
            f"critical value at z={critical_point}",
        )

        specialized = sp.factor(radicand.subs(energy, critical_energy))
        other_root = -2 * critical_point
        expected_factorization = (
            (z - critical_point) ** 2 * (z - other_root)
        )
        require_zero(
            specialized - expected_factorization,
            f"critical fiber E={critical_energy}: factorization",
        )
        specialized_poly = sp.Poly(specialized, z, domain=sp.QQ)
        specialized_derivative = sp.Poly(
            sp.diff(specialized, z),
            z,
            domain=sp.QQ,
        )
        fiber_gcd = sp.gcd(
            specialized_poly,
            specialized_derivative,
        ).monic()
        expected_gcd = sp.Poly(
            z - critical_point,
            z,
            domain=sp.QQ,
        )
        require(
            fiber_gcd == expected_gcd,
            f"critical fiber E={critical_energy}: gcd",
            fiber_gcd.as_expr(),
        )

        local_expression = sp.expand(
            radicand.subs(
                {
                    z: critical_point + u,
                    energy: critical_energy + epsilon,
                }
            )
        )
        expected_local = u**3 + 3 * critical_point * u**2 - epsilon
        require_zero(
            local_expression - expected_local,
            f"critical fiber E={critical_energy}: local expansion",
        )
        local_multiplicity = order_at_zero(
            local_expression.subs(epsilon, 0),
            u,
        )
        require(
            local_multiplicity == 2,
            f"critical fiber E={critical_energy}: double zero",
            local_multiplicity,
        )
        require_nonzero(
            sp.diff(local_expression, u, 2).subs(
                {u: 0, epsilon: 0}
            ),
            f"critical fiber E={critical_energy}: quadratic term",
        )
        fibers.append(
            {
                "critical_point": critical_point,
                "critical_energy": critical_energy,
                "factorization": specialized,
                "gcd": fiber_gcd.as_expr(),
                "local_expression": local_expression,
            }
        )

    return {
        "radicand": radicand,
        "discriminant": discriminant,
        "critical_energies": tuple(discriminant_roots),
        "generic_gcd": generic_gcd.as_expr(),
        "turning_derivative": derivative_formula,
        "fibers": fibers,
    }


def double_zero_normalization_audit() -> dict[str, object]:
    """Separate the two normalized points over a double zero."""
    z, y, w, u = sp.symbols("z y w u")
    radicand = z**2 * (z - 1)

    # Away from z=0 put w=y/z.  The relation w**2=z-1 extends the
    # normalization over z=0, where w=+i and w=-i are distinct points.
    z_of_w = w**2 + 1
    y_of_w = z_of_w * w
    require_zero(
        y_of_w**2 - radicand.subs(z, z_of_w),
        "double-zero normalization map",
    )
    require_zero(
        (y**2 - radicand).subs(y, z * w) / z**2 - (w**2 - z + 1),
        "double-zero normalized relation",
    )

    sheet_rows: list[dict[str, object]] = []
    for point in (-sp.I, sp.I):
        local_z = sp.expand(z_of_w.subs(w, point + u))
        local_y = sp.expand(y_of_w.subs(w, point + u))
        local_lambda = sp.expand(local_y * sp.diff(local_z, u))
        require_zero(
            local_y**2 - radicand.subs(z, local_z),
            f"double-zero sheet w={point}: curve equation",
        )
        projection_order = order_at_zero(local_z, u)
        lambda_order = order_at_zero(local_lambda, u)
        require(
            projection_order == 1,
            f"double-zero sheet w={point}: unramified projection",
            projection_order,
        )
        require(
            lambda_order == 1,
            f"double-zero sheet w={point}: lambda order",
            lambda_order,
        )
        sheet_rows.append(
            {
                "point": point,
                "ramification_index": projection_order,
                "lambda_order": lambda_order,
            }
        )

    require(
        sheet_rows[0]["point"] != sheet_rows[1]["point"],
        "double zero has two normalized points",
    )
    return {
        "radicand": radicand,
        "normalization": sp.Eq(w**2, z - 1),
        "sheets": sheet_rows,
        "branched_at_double_zero": False,
    }


def factor_orders(
    polynomial: sp.Poly,
    sign: int,
) -> list[tuple[sp.Expr, int, int]]:
    """Return irreducible factors with signed order and complex root count."""
    _, factors = sp.factor_list(polynomial)
    rows: list[tuple[sp.Expr, int, int]] = []
    for factor, multiplicity in factors:
        rows.append(
            (
                factor.as_expr(),
                sign * int(multiplicity),
                int(factor.degree()),
            )
        )
    return rows


def rational_branch_data(
    label: str,
    radicand: sp.Expr,
    z: sp.Symbol,
) -> dict[str, object]:
    """Count odd finite valuations and the quadratic-differential order at infinity."""
    numerator, denominator = sp.fraction(sp.cancel(radicand))
    numerator_poly = sp.Poly(numerator, z, domain=sp.QQ)
    denominator_poly = sp.Poly(denominator, z, domain=sp.QQ)
    require(not numerator_poly.is_zero, f"{label}: nonzero radicand")

    valuations = factor_orders(numerator_poly, 1)
    valuations.extend(factor_orders(denominator_poly, -1))
    finite_divisor_degree = sum(
        order * root_count
        for _, order, root_count in valuations
    )
    infinity_order = (
        denominator_poly.degree() - numerator_poly.degree() - 4
    )
    require(
        finite_divisor_degree + infinity_order == -4,
        f"{label}: degree of quadratic-differential divisor",
        finite_divisor_degree + infinity_order,
    )

    finite_branches = sum(
        root_count
        for _, order, root_count in valuations
        if abs(order) % 2 == 1
    )
    infinity_branch = abs(infinity_order) % 2
    total_branches = finite_branches + infinity_branch
    require(
        total_branches % 2 == 0,
        f"{label}: branch-count parity including infinity",
        total_branches,
    )

    has_odd_valuation = total_branches > 0
    if has_odd_valuation:
        genus: int | None = connected_genus(total_branches, label)
        connected = True
        split = False
    else:
        # On P1, even valuations make a rational function a square over C;
        # the two sign components are therefore separate.
        genus = None
        connected = False
        split = True

    return {
        "label": label,
        "radicand": sp.factor(radicand),
        "valuations": valuations,
        "finite_branches": finite_branches,
        "infinity_order": infinity_order,
        "infinity_branch": bool(infinity_branch),
        "total_branches": total_branches,
        "connected": connected,
        "split": split,
        "genus": genus,
    }


def branch_parity_audit() -> list[dict[str, object]]:
    """Audit branch parity for several polynomial and rational examples."""
    z, y = sp.symbols("z y")
    square_radicand = (z**2 - 1) ** 2
    examples = (
        ("linear polynomial", z),
        ("quadratic polynomial", z**2 - 1),
        (
            "quartic polynomial",
            (z - 3) * (z - 1) * (z + 1) * (z + 2),
        ),
        ("double-plus-simple polynomial", z**2 * (z - 1)),
        ("rational zero and pole", (z - 1) / (z + 1)),
        (
            "rational three finite odd valuations",
            (z - 1) / (z * (z + 1)),
        ),
        (
            "rational infinity zero",
            1 / (z * (z - 1) * (z + 1) * (z - 2) * (z + 2)),
        ),
        ("global square polynomial", square_radicand),
    )
    rows = [
        rational_branch_data(label, radicand, z)
        for label, radicand in examples
    ]

    square_row = rows[-1]
    require_zero(
        square_radicand - (z**2 - 1) ** 2,
        "global-square square root",
    )
    require_zero(
        y**2 - square_radicand
        - (y - (z**2 - 1)) * (y + (z**2 - 1)),
        "global-square cover factorization",
    )
    require(
        square_row["split"] is True,
        "global-square cover splits",
        square_row,
    )
    require(
        square_row["total_branches"] == 0,
        "global-square cover has no branch points",
        square_row["total_branches"],
    )
    return rows


def print_monomial_rows(rows: list[dict[str, object]]) -> None:
    """Print the local monomial classification."""
    print("\n1. Local monomial normalization")
    print("   m   branch  points  ord(lambda)  pole-order")
    for row in rows:
        branch = "yes" if row["branched"] else "no"
        print(
            f"  {row['m']:>3}   {branch:<3}      "
            f"{row['points_above']}         "
            f"{row['lambda_order']:>3}          {row['pole_order']:>3}"
        )
    print("   odd m: z=xi^2, lambda=2*xi^(m+1) dxi")
    print("   even m: two signs, lambda=+/-z^(m/2) dz")


def print_infinity_rows(rows: list[dict[str, object]]) -> None:
    """Print the Airy and Weber infinity checks."""
    print("\n2. Infinity transformation and genus")
    for row in rows:
        print(
            f"   {row['label']}: R(1/xi)*xi^-4 = "
            f"{row['transformed']}"
        )
        print(
            "      finite simple zeros="
            f"{row['finite_simple_zeros']}, infinity pole-order="
            f"{row['infinity_pole_order']}, branches="
            f"{row['branch_count']}, connected=yes, genus={row['genus']}"
        )


def print_cubic(data: dict[str, object]) -> None:
    """Print the cubic-family checks."""
    print("\n3. Cubic family R=z^3-3*z-E")
    print(f"   discriminant = {data['discriminant']}")
    print(f"   critical energies = {data['critical_energies']}")
    print(f"   generic gcd(R, dR/dz) = {data['generic_gcd']}")
    print(
        "   away from V'(z_t)=0: dz_t/dE = "
        f"{data['turning_derivative']}"
    )
    fibers = data["fibers"]
    if not isinstance(fibers, list):
        raise RuntimeError("internal cubic-fiber data is not a list")
    for row in fibers:
        print(
            f"   E={row['critical_energy']}, z_c={row['critical_point']}: "
            f"R={row['factorization']}, gcd={row['gcd']}"
        )
        print(
            "      with z=z_c+u and E=E_c+epsilon: "
            f"R={row['local_expression']}"
        )


def print_double_zero(data: dict[str, object]) -> None:
    """Print the normalized double-zero fiber."""
    print("\n4. Double zero without ramification")
    print(f"   R = {data['radicand']}; normalization: {data['normalization']}")
    print("   fiber over z=0:")
    sheets = data["sheets"]
    if not isinstance(sheets, list):
        raise RuntimeError("internal double-zero sheet data is not a list")
    for row in sheets:
        print(
            f"      w={row['point']}: e={row['ramification_index']}, "
            f"ord(lambda)={row['lambda_order']}"
        )
    print("   branch=no; the normalized fiber has two local sheets")


def print_branch_rows(rows: list[dict[str, object]]) -> None:
    """Print the branch-parity table."""
    print("\n5. Branch-count parity, including infinity")
    print(
        "   example                            finite  ord_inf  "
        "br_inf  total  result"
    )
    for row in rows:
        infinity_branch = "yes" if row["infinity_branch"] else "no"
        if row["split"]:
            result = "split (global square)"
        else:
            result = f"connected, genus {row['genus']}"
        print(
            f"   {row['label']:<34}"
            f"{row['finite_branches']:>3}"
            f"{row['infinity_order']:>10}"
            f"{infinity_branch:>8}"
            f"{row['total_branches']:>7}  {result}"
        )


def main() -> None:
    """Run every independent audit and print the exact results."""
    print(
        f"Python {platform.python_version()} "
        f"({platform.python_implementation()})"
    )
    print(f"SymPy {sp.__version__}")

    monomial_rows = monomial_local_audit()
    infinity_rows = infinity_examples_audit()
    cubic_data = cubic_family_audit()
    double_zero_data = double_zero_normalization_audit()
    branch_rows = branch_parity_audit()

    print_monomial_rows(monomial_rows)
    print_infinity_rows(infinity_rows)
    print_cubic(cubic_data)
    print_double_zero(double_zero_data)
    print_branch_rows(branch_rows)
    print("\nAll exact spectral-cover checks passed.")


if __name__ == "__main__":
    main()