#!/usr/bin/env python3
"""Exact Kummer/Whittaker audit of a sectorial confluent-Heun connection.

The parameter choice

    (gamma, delta, epsilon, alpha, q) = (4/3, 0, -1, -5/12, -5/12)

removes z=1 and reduces the confluent-Heun equation to Kummer's
equation.  The script checks the lower-lateral regular-to-infinity
connection matrix, its derivatives, and its Abel-Wronskian determinant.

Requires only mpmath.
"""

from __future__ import annotations

import argparse

import mpmath as mp


def format_complex(value: mp.mpc, digits: int = 18) -> str:
    """Return a compact complex number with a stable explicit sign."""

    real = mp.nstr(mp.re(value), digits)
    imag = mp.nstr(abs(mp.im(value)), digits)
    sign = "+" if mp.im(value) >= 0 else "-"
    return f"{real} {sign} {imag}j"


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Audit the exact lower-lateral Kummer/Whittaker connection."
    )
    parser.add_argument("--dps", type=int, default=60, help="mpmath precision")
    parser.add_argument(
        "--radius",
        type=str,
        default="3",
        help="radius of the matching point",
    )
    parser.add_argument(
        "--angle",
        type=str,
        default="0.25",
        help="matching angle as a multiple of pi",
    )
    args = parser.parse_args()

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

    mp.mp.dps = args.dps

    kappa = mp.mpf(1) / 4
    mu = mp.mpf(1) / 6
    try:
        radius = mp.mpf(args.radius)
        angle_multiple = mp.mpf(args.angle)
    except (TypeError, ValueError):
        parser.error("--radius and --angle must be finite real numbers")

    if not mp.isfinite(radius) or radius <= 0:
        parser.error("--radius must be a finite positive number")
    if not mp.isfinite(angle_multiple) or not 0 < angle_multiple < 1:
        parser.error(
            "--angle must satisfy 0 < angle < 1 on the declared upper-half-plane branch"
        )

    angle = angle_multiple * mp.pi
    z_star = radius * mp.e ** (1j * angle)

    def gauge(z: mp.mpc) -> mp.mpc:
        return mp.e ** (z / 2) * z ** (-mp.mpf(2) / 3)

    def regular_plus(z: mp.mpc) -> mp.mpc:
        return gauge(z) * mp.whitm(kappa, mu, z)

    def regular_minus(z: mp.mpc) -> mp.mpc:
        return gauge(z) * mp.whitm(kappa, -mu, z)

    def infinity_algebraic(z: mp.mpc) -> mp.mpc:
        return gauge(z) * mp.whitw(kappa, mu, z)

    def infinity_exponential_lower(z: mp.mpc) -> mp.mpc:
        rotated = mp.e ** (-1j * mp.pi) * z
        return (
            mp.e ** (-1j * mp.pi * kappa)
            * gauge(z)
            * mp.whitw(-kappa, mu, rotated)
        )

    c00 = (
        mp.gamma(mp.mpf(4) / 3)
        * mp.e ** (5j * mp.pi / 12)
        / mp.gamma(mp.mpf(11) / 12)
    )
    c01 = (
        mp.gamma(mp.mpf(2) / 3)
        * mp.e ** (1j * mp.pi / 12)
        / mp.gamma(mp.mpf(7) / 12)
    )
    c10 = mp.gamma(mp.mpf(4) / 3) / mp.gamma(mp.mpf(5) / 12)
    c11 = mp.gamma(mp.mpf(2) / 3) / mp.gamma(mp.mpf(1) / 12)

    upper_c00 = (
        mp.gamma(mp.mpf(4) / 3)
        * mp.e ** (-5j * mp.pi / 12)
        / mp.gamma(mp.mpf(11) / 12)
    )
    upper_c01 = (
        mp.gamma(mp.mpf(2) / 3)
        * mp.e ** (-1j * mp.pi / 12)
        / mp.gamma(mp.mpf(7) / 12)
    )

    d0 = infinity_algebraic(z_star)
    d1 = infinity_exponential_lower(z_star)
    r0 = regular_plus(z_star)
    r1 = regular_minus(z_star)

    reconstructed_r0 = d0 * c00 + d1 * c10
    reconstructed_r1 = d0 * c01 + d1 * c11

    d0_prime = mp.diff(infinity_algebraic, z_star)
    d1_prime = mp.diff(infinity_exponential_lower, z_star)
    r0_prime = mp.diff(regular_plus, z_star)
    r1_prime = mp.diff(regular_minus, z_star)

    reconstructed_r0_prime = d0_prime * c00 + d1_prime * c10
    reconstructed_r1_prime = d0_prime * c01 + d1_prime * c11

    value_residual = max(
        abs(reconstructed_r0 - r0),
        abs(reconstructed_r1 - r1),
    )
    derivative_residual = max(
        abs(reconstructed_r0_prime - r0_prime),
        abs(reconstructed_r1_prime - r1_prime),
    )

    determinant = c00 * c11 - c01 * c10
    determinant_error = abs(determinant + mp.mpf(1) / 3)
    stokes_entry = (
        2j
        * mp.pi
        / (
            mp.gamma(mp.mpf(11) / 12)
            * mp.gamma(mp.mpf(7) / 12)
        )
    )
    stokes_matrix_residual = max(
        abs(upper_c00 + stokes_entry * c10 - c00),
        abs(upper_c01 + stokes_entry * c11 - c01),
    )

    print("Exact Kummer/Whittaker confluent-Heun audit")
    print(f"dps = {args.dps}")
    print(f"z*  = {format_complex(z_star)}")
    print()
    print("C_lower, with (R_+, R_-) = (D_0, D_1^lower) C_lower")
    print(f"  [{format_complex(c00)}, {format_complex(c01)}]")
    print(f"  [{format_complex(c10)}, {format_complex(c11)}]")
    print()
    print(f"D_0(z*)       = {format_complex(d0)}")
    print(f"D_1^lower(z*) = {format_complex(d1)}")
    print(f"R_+(z*)       = {format_complex(r0)}")
    print(f"R_-(z*)       = {format_complex(r1)}")
    print()
    print(f"max value residual      = {mp.nstr(value_residual, 8)}")
    print(f"max derivative residual = {mp.nstr(derivative_residual, 8)}")
    print(f"det C_lower             = {mp.nstr(determinant, 24)}")
    print(f"|det C_lower + 1/3|     = {mp.nstr(determinant_error, 8)}")
    print(f"upper/lower Stokes entry = {format_complex(stokes_entry)}")
    print(
        "max |S C_upper - C_lower| = "
        f"{mp.nstr(stokes_matrix_residual, 8)}"
    )


if __name__ == "__main__":
    main()