#!/usr/bin/env python3
"""Reproduce the Chapters 7 and 11 general-Heun connection-matrix audits.

The script compares

    C_{t0} = Phi_t(z_*)^{-1} Phi_0(z_*)

computed from the two local Frobenius series with the classical-block
factorization

    C_{t0} = D_t^{-1} M_fr D_0.

The default matching point z=t/2 lies in both convergence disks for
0<t<2/3. The optional second-point audit repeats the direct match at
z=t/3. Only mpmath is required.

Chapter 11 additionally uses the optional endpoint-factor ablation and
the redundant reverse extraction of F_0 and F_t.
"""

from __future__ import annotations

import argparse

import mpmath as mp


def heun_g_series(
    singularity: mp.mpf,
    accessory: mp.mpf,
    alpha: mp.mpf,
    beta: mp.mpf,
    gamma: mp.mpf,
    delta: mp.mpf,
    z: mp.mpf,
    terms: int,
) -> tuple[mp.mpf, mp.mpf]:
    """Return the unit-leading HeunG series and its z derivative."""
    epsilon = alpha + beta - gamma - delta + 1
    coefficients = [mp.mpf(1)]
    previous = mp.mpf(0)
    current = mp.mpf(1)

    for n in range(terms - 1):
        p_n = (n - 1 + alpha) * (n - 1 + beta)
        q_n = n * (
            (n - 1 + gamma) * (1 + singularity)
            + singularity * delta
            + epsilon
        )
        r_n = singularity * (n + 1) * (n + gamma)
        following = ((q_n + accessory) * current - p_n * previous) / r_n
        coefficients.append(following)
        previous, current = current, following

    value = mp.fsum(
        coefficient * z**n
        for n, coefficient in enumerate(coefficients)
    )
    derivative = mp.fsum(
        n * coefficients[n] * z ** (n - 1)
        for n in range(1, terms)
    )
    return value, derivative


def rational_slice_data(
    t: mp.mpf,
    order: int,
) -> tuple[mp.mpf, mp.mpf, mp.mpf]:
    """Return q_H, partial_{a0} f-hat, and partial_{at} f-hat."""
    q_0 = -mp.mpf(14) / 225
    q_1 = mp.mpf(14) / 75
    q_2 = mp.mpf(119) / 5850

    f_0 = mp.mpf(0)
    f_t = mp.mpf(0)
    accessory = q_0

    if order >= 1:
        accessory += q_1 * t
        f_0 += t / 6
        f_t -= t / 6
    if order >= 2:
        accessory += q_2 * t**2
        f_0 += mp.mpf(11) * t**2 / 156
        f_t -= mp.mpf(5) * t**2 / 52

    return accessory, f_0, f_t


def direct_connection(
    t: mp.mpf,
    order: int,
    terms: int,
    match_fraction: mp.mpf = mp.mpf("0.5"),
) -> mp.matrix:
    """Match the two ordered Frobenius frames at z=t*match_fraction."""
    gamma = mp.mpf(2) / 3
    epsilon = mp.mpf(2) / 3
    delta = mp.mpf(4) / 5
    alpha = mp.mpf(2) / 3
    beta = mp.mpf(7) / 15
    theta_0 = 1 - gamma
    theta_t = 1 - epsilon
    accessory, _, _ = rational_slice_data(t, order)
    z_star = t * match_fraction

    h_0_minus, dh_0_minus = heun_g_series(
        t,
        accessory,
        alpha,
        beta,
        gamma,
        delta,
        z_star,
        terms,
    )

    q_0_plus = accessory - (gamma - 1) * (t * delta + epsilon)
    analytic_0_plus, d_analytic_0_plus = heun_g_series(
        t,
        q_0_plus,
        alpha + theta_0,
        beta + theta_0,
        2 - gamma,
        delta,
        z_star,
        terms,
    )
    h_0_plus = z_star**theta_0 * analytic_0_plus
    dh_0_plus = (
        theta_0 * z_star ** (theta_0 - 1) * analytic_0_plus
        + z_star**theta_0 * d_analytic_0_plus
    )

    x_star = (z_star - t) / (1 - t)
    transformed_t = t / (t - 1)
    q_t = (accessory - t * alpha * beta) / (1 - t)
    h_t_minus, d_h_t_minus_dx = heun_g_series(
        transformed_t,
        q_t,
        alpha,
        beta,
        epsilon,
        delta,
        x_star,
        terms,
    )
    dh_t_minus = d_h_t_minus_dx / (1 - t)

    q_t_plus = q_t - (epsilon - 1) * (
        transformed_t * delta + gamma
    )
    analytic_t_plus, d_analytic_t_plus_dx = heun_g_series(
        transformed_t,
        q_t_plus,
        alpha + theta_t,
        beta + theta_t,
        2 - epsilon,
        delta,
        x_star,
        terms,
    )
    local_t = t - z_star
    h_t_plus = local_t**theta_t * analytic_t_plus
    dh_t_plus = (
        -theta_t * local_t ** (theta_t - 1) * analytic_t_plus
        + local_t**theta_t * d_analytic_t_plus_dx / (1 - t)
    )

    phi_0 = mp.matrix(
        [
            [h_0_minus, h_0_plus],
            [dh_0_minus, dh_0_plus],
        ]
    )
    phi_t = mp.matrix(
        [
            [h_t_minus, h_t_plus],
            [dh_t_minus, dh_t_plus],
        ]
    )
    return phi_t**-1 * phi_0


def fusion_coefficient(
    source_sign: int,
    target_sign: int,
    internal: mp.mpf = None,
) -> mp.mpf:
    """Finite degenerate-fusion coefficient on the rational slice."""
    a_0 = mp.mpf(1) / 6
    a_t = mp.mpf(1) / 6
    if internal is None:
        internal = mp.mpf(3) / 10
    numerator = (
        mp.gamma(-2 * target_sign * a_t)
        * mp.gamma(1 + 2 * source_sign * a_0)
    )
    denominator = (
        mp.gamma(
            mp.mpf("0.5")
            + source_sign * a_0
            - target_sign * a_t
            + internal
        )
        * mp.gamma(
            mp.mpf("0.5")
            + source_sign * a_0
            - target_sign * a_t
            - internal
        )
    )
    return numerator / denominator


def block_connection(
    t: mp.mpf,
    order: int,
    include_derivatives: bool = True,
) -> mp.matrix:
    """Evaluate D_t^{-1} M_fr D_0 on the rational slice."""
    a_0 = mp.mpf(1) / 6
    a_t = mp.mpf(1) / 6
    a_1 = mp.mpf(1) / 10
    _, f_0, f_t = rational_slice_data(t, order)
    if not include_derivatives:
        f_0 = mp.mpf(0)
        f_t = mp.mpf(0)

    def entry(target_sign: int, source_sign: int) -> mp.mpf:
        t_power = (
            a_0 * (source_sign + 1)
            - a_t * (target_sign + 1)
        )
        return (
            fusion_coefficient(source_sign, target_sign)
            * t**t_power
            * (1 - t) ** (a_1 - mp.mpf("0.5"))
            * mp.exp(
                source_sign * f_0 / 2
                - target_sign * f_t / 2
            )
        )

    return mp.matrix(
        [
            [entry(-1, -1), entry(-1, +1)],
            [entry(+1, -1), entry(+1, +1)],
        ]
    )


def expected_determinant(t: mp.mpf) -> mp.mpf:
    """Abel-Wronskian determinant for the rational slice."""
    gamma = mp.mpf(2) / 3
    epsilon = mp.mpf(2) / 3
    delta = mp.mpf(4) / 5
    return (
        (1 - gamma)
        / (epsilon - 1)
        * t ** (epsilon - gamma)
        * (1 - t) ** (-delta)
    )


def recover_endpoint_derivatives(
    t: mp.mpf,
    connection: mp.matrix,
    internal: mp.mpf = None,
) -> tuple[tuple[mp.mpf, mp.mpf], tuple[mp.mpf, mp.mpf]]:
    """Recover F_0 and F_t twice from a normalized connection matrix."""
    gamma = mp.mpf(2) / 3
    epsilon = mp.mpf(2) / 3
    m_mm = fusion_coefficient(-1, -1, internal)
    m_pm = fusion_coefficient(+1, -1, internal)
    m_mp = fusion_coefficient(-1, +1, internal)
    m_pp = fusion_coefficient(+1, +1, internal)

    c_mm = connection[0, 0]
    c_mp = connection[0, 1]
    c_pm = connection[1, 0]
    c_pp = connection[1, 1]

    f_0_target_minus = mp.log(
        c_mp * m_mm
        / (t ** (1 - gamma) * c_mm * m_pm)
    )
    f_0_target_plus = mp.log(
        c_pp * m_mp
        / (t ** (1 - gamma) * c_pm * m_pp)
    )
    f_t_source_minus = -mp.log(
        t ** (1 - epsilon) * c_pm * m_mm
        / (c_mm * m_mp)
    )
    f_t_source_plus = -mp.log(
        t ** (1 - epsilon) * c_pp * m_pm
        / (c_mp * m_pp)
    )
    return (
        (f_0_target_minus, f_0_target_plus),
        (f_t_source_minus, f_t_source_plus),
    )


def effective_internal_lift(
    connection: mp.matrix,
) -> tuple[mp.mpf, mp.mpf]:
    """Recover the nearby composite trace and principal internal lift."""
    theta = mp.mpf(1) / 3
    prefactor = -mp.exp(-mp.pi * 1j * theta)
    local = mp.matrix(
        [
            [prefactor, 0],
            [0, prefactor * mp.exp(2 * mp.pi * 1j * theta)],
        ]
    )
    composite = local * connection**-1 * local * connection
    trace = composite[0, 0] + composite[1, 1]
    trace = mp.re(trace)
    theta_effective = mp.acos(-trace / 2) / mp.pi
    return trace, theta_effective / 2


def maximum_relative_error(
    approximation: mp.matrix,
    reference: mp.matrix,
) -> mp.mpf:
    return max(
        abs((approximation[row, column] - reference[row, column])
            / reference[row, column])
        for row in range(2)
        for column in range(2)
    )


def print_matrix(label: str, matrix: mp.matrix, digits: int = 16) -> None:
    print(label)
    for row in range(2):
        print(
            "  "
            + "  ".join(
                mp.nstr(matrix[row, column], digits)
                for column in range(2)
            )
        )


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--t", type=str, default="0.04")
    parser.add_argument("--order", type=int, choices=(0, 1, 2), default=2)
    parser.add_argument("--terms", type=int, default=120)
    parser.add_argument("--dps", type=int, default=50)
    parser.add_argument(
        "--no-table",
        action="store_true",
        help="omit the four-t, three-order convergence table",
    )
    parser.add_argument(
        "--show-ablation",
        action="store_true",
        help="also omit both endpoint-derivative factors and report the error",
    )
    parser.add_argument(
        "--check-second-point",
        action="store_true",
        help="repeat direct matching at z=t/3 and compare the matrices",
    )
    parser.add_argument(
        "--show-reverse",
        action="store_true",
        help="recover both endpoint derivatives twice from the direct matrix",
    )
    args = parser.parse_args()

    if args.dps < 30:
        parser.error("--dps must be at least 30 for this numerical audit")
    if args.terms < 24:
        parser.error("--terms must be at least 24")

    mp.mp.dps = args.dps
    try:
        t = mp.mpf(args.t)
    except (TypeError, ValueError):
        parser.error("--t must be a finite real number")

    if not mp.isfinite(t):
        parser.error("--t must be a finite real number")
    if not 0 < t < mp.mpf(2) / 3:
        parser.error(
            "matching at z=t/2 requires 0 < t < 2/3 so that both "
            "Frobenius series converge"
        )

    direct = direct_connection(t, args.order, args.terms)
    block = block_connection(t, args.order)
    determinant = expected_determinant(t)

    print(f"t={mp.nstr(t, 12)}, order={args.order}, terms={args.terms}")
    print_matrix("direct Frobenius matrix", direct)
    print_matrix("classical-block matrix", block)
    print(
        "maximum relative entry error = "
        + mp.nstr(maximum_relative_error(block, direct), 14)
    )
    print(
        "determinants (direct, block, Abel) = "
        + ", ".join(
            mp.nstr(value, 18)
            for value in (mp.det(direct), mp.det(block), determinant)
        )
    )
    direct_determinant_residual = abs(mp.det(direct) - determinant)
    convergence_tolerance = mp.power(
        10,
        -min(12, max(6, args.dps // 4)),
    )
    if direct_determinant_residual > convergence_tolerance:
        print(
            "WARNING: direct-series determinant residual "
            + mp.nstr(direct_determinant_residual, 8)
            + " exceeds "
            + mp.nstr(convergence_tolerance, 4)
            + "; increase --terms or move farther from the convergence boundary"
        )

    if args.show_ablation:
        ablated = block_connection(
            t,
            args.order,
            include_derivatives=False,
        )
        print()
        print_matrix("matrix with endpoint derivatives omitted", ablated)
        print(
            "ablated maximum relative entry error = "
            + mp.nstr(maximum_relative_error(ablated, direct), 14)
        )
        print(
            "ablated determinant = "
            + mp.nstr(mp.det(ablated), 18)
            + "  (still constrained by Abel)"
        )

    if args.check_second_point:
        direct_second = direct_connection(
            t,
            args.order,
            args.terms,
            match_fraction=mp.mpf(1) / 3,
        )
        print()
        print(
            "maximum relative change between z=t/2 and z=t/3 = "
            + mp.nstr(
                maximum_relative_error(direct_second, direct),
                14,
            )
        )

    if args.show_reverse:
        recovered_f_0, recovered_f_t = recover_endpoint_derivatives(
            t,
            direct,
        )
        _, predicted_f_0, predicted_f_t = rational_slice_data(
            t,
            args.order,
        )
        print()
        print(
            "recovered F_0 (target minus, target plus) = "
            + ", ".join(mp.nstr(value, 16) for value in recovered_f_0)
        )
        print(
            "recovered F_t (source minus, source plus) = "
            + ", ".join(mp.nstr(value, 16) for value in recovered_f_t)
        )
        print(
            "truncated predictions (F_0, F_t) = "
            + mp.nstr(predicted_f_0, 16)
            + ", "
            + mp.nstr(predicted_f_t, 16)
        )
        effective_trace, effective_internal = effective_internal_lift(
            direct,
        )
        effective_f_0, effective_f_t = recover_endpoint_derivatives(
            t,
            direct,
            internal=effective_internal,
        )
        print(
            "effective composite trace = "
            + mp.nstr(effective_trace, 18)
        )
        print(
            "effective (theta_0t, a) = "
            + mp.nstr(2 * effective_internal, 18)
            + ", "
            + mp.nstr(effective_internal, 18)
        )
        print(
            "effective-core F_0 pair = "
            + ", ".join(mp.nstr(value, 16) for value in effective_f_0)
        )
        print(
            "effective-core F_t pair = "
            + ", ".join(mp.nstr(value, 16) for value in effective_f_t)
        )

    if not args.no_table:
        print("\nmaximum relative entry error")
        print("t          N=0              N=1              N=2")
        for table_t_text in ("0.08", "0.04", "0.02", "0.01"):
            table_t = mp.mpf(table_t_text)
            errors = []
            for order in range(3):
                reference = direct_connection(table_t, order, args.terms)
                approximation = block_connection(table_t, order)
                errors.append(
                    maximum_relative_error(approximation, reference)
                )
            print(
                f"{table_t_text:<10} "
                + "  ".join(f"{mp.nstr(error, 12):<16}" for error in errors)
            )


if __name__ == "__main__":
    main()