#!/usr/bin/env python3
"""Verify the Chapter 7 classical-block accessory jet by monodromy.

The script integrates the determinant-one first-order form of

    psi''(z) + T_op(z; t) psi(z) = 0

around a counterclockwise circle enclosing z = 0 and z = t.  It compares
the resulting composite trace with -2 cos(pi theta_0t) after retaining:

    order 0: kappa / t,
    order 1: kappa / t + f1,
    order 2: kappa / t + f1 + 2 f2 t.

Requirements: Python 3 and mpmath.
"""

from __future__ import annotations

import argparse

import mpmath as mp


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--dps",
        type=int,
        default=40,
        help="decimal working precision (default: 40)",
    )
    parser.add_argument(
        "--steps",
        type=int,
        default=5000,
        help="fixed fourth-order Runge–Kutta steps (default: 5000)",
    )
    parser.add_argument(
        "--radius",
        default="0.2",
        help="contour radius as an mpmath decimal string (default: 0.2)",
    )
    return parser.parse_args()


def add_scaled(
    vector: list[mp.mpc],
    scale: mp.mpf | mp.mpc,
    increment: list[mp.mpc],
) -> list[mp.mpc]:
    return [
        vector[index] + scale * increment[index]
        for index in range(4)
    ]


def monodromy(
    t: mp.mpf,
    truncation: int,
    steps: int,
    radius: mp.mpf,
) -> tuple[mp.mpc, mp.mpc]:
    """Return the composite trace and determinant for one truncation."""

    delta_0 = delta_t = mp.mpf(2) / 9
    delta_1 = delta_infinity = mp.mpf(6) / 25
    delta_internal = mp.mpf(4) / 25

    kappa = delta_internal - delta_0 - delta_t
    f1 = delta_internal / 2
    f2 = mp.mpf(587) / 11700
    accessory = kappa / t

    if truncation >= 1:
        accessory += f1
    if truncation >= 2:
        accessory += 2 * f2 * t

    lambda_infinity = (
        delta_infinity - delta_0 - delta_t - delta_1
    )
    center = t / 2
    step = 2 * mp.pi / steps

    def rhs(phi: mp.mpf, state: list[mp.mpc]) -> list[mp.mpc]:
        phase = mp.exp(1j * phi)
        z = center + radius * phase
        dz_dphi = 1j * radius * phase
        potential = (
            delta_0 / z**2
            + delta_t / (z - t) ** 2
            + delta_1 / (z - 1) ** 2
            + lambda_infinity / (z * (z - 1))
            + t
            * (t - 1)
            * accessory
            / (z * (z - 1) * (z - t))
        )

        # Row-major fundamental matrix Y = [[a, b], [c, d]].
        a, b, c, d = state
        return [
            dz_dphi * c,
            dz_dphi * d,
            -dz_dphi * potential * a,
            -dz_dphi * potential * b,
        ]

    state = [mp.mpc(1), mp.mpc(0), mp.mpc(0), mp.mpc(1)]
    phi = mp.mpf(0)

    for _ in range(steps):
        k1 = rhs(phi, state)
        k2 = rhs(phi + step / 2, add_scaled(state, step / 2, k1))
        k3 = rhs(phi + step / 2, add_scaled(state, step / 2, k2))
        k4 = rhs(phi + step, add_scaled(state, step, k3))
        state = [
            state[index]
            + step
            * (
                k1[index]
                + 2 * k2[index]
                + 2 * k3[index]
                + k4[index]
            )
            / 6
            for index in range(4)
        ]
        phi += step

    a, b, c, d = state
    return a + d, a * d - b * c


def main() -> None:
    args = parse_args()
    if args.dps < 20:
        raise SystemExit("--dps must be at least 20")
    if args.steps < 100:
        raise SystemExit("--steps must be at least 100")

    mp.mp.dps = args.dps
    radius = mp.mpf(args.radius)
    target = -2 * mp.cos(3 * mp.pi / 5)
    t_values = [mp.mpf(value) for value in ("0.08", "0.04", "0.02", "0.01")]
    largest_t = max(t_values)
    lower_radius = largest_t / 2
    upper_radius = 1 - largest_t / 2

    if not lower_radius < radius < upper_radius:
        raise SystemExit(
            "--radius must satisfy "
            f"{lower_radius} < radius < {upper_radius}; "
            "the circle must enclose 0 and every t value but exclude 1"
        )

    print("Classical-block composite-monodromy check")
    print(f"mpmath={mp.__version__}")
    print(f"dps={args.dps}, steps={args.steps}, radius={radius}")
    print(f"target trace={mp.nstr(target, 24)}")
    print()
    print(
        "t       error(kappa/t)   error(+f1)      "
        "error(+2 f2 t)  max final |det-1|"
    )

    for t in t_values:
        errors: list[mp.mpf] = []
        determinant_errors: list[mp.mpf] = []

        for truncation in range(3):
            trace, determinant = monodromy(
                t=t,
                truncation=truncation,
                steps=args.steps,
                radius=radius,
            )
            errors.append(abs(trace - target))
            determinant_errors.append(abs(determinant - 1))

        print(
            f"{mp.nstr(t, 4):<7} "
            f"{mp.nstr(errors[0], 12):<16} "
            f"{mp.nstr(errors[1], 12):<16} "
            f"{mp.nstr(errors[2], 12):<16} "
            f"{mp.nstr(max(determinant_errors), 6)}"
        )


if __name__ == "__main__":
    main()