"""Rank-one ABJM mass responses on a specified toric sphere.

Python 3.13; mpmath 1.3.0. No repository imports or network access.
Conventions and sources: KNP 2411.07334v3 (2.4), (4.11), (A.8), (A.10);
BCP 1808.10554v2 (3.7), (3.9), (3.43). Z_b(x)=D_b(x/2)**2.
The source x is real and dimensionless. Derivatives are at x=0, fixed b.
"""

import argparse
import json
import mpmath as mp


def determinant(w, b_squared):
    """KNP A.8, including its absolute normalization, for odd b_squared."""
    b = mp.sqrt(b_squared)
    n = (b_squared + 1) // 2
    terms = [2 * mp.cosh(mp.pi * w / b + 1j * mp.pi / b_squared
                       * (mp.mpf(n + 1) / 2 - j)) for j in range(1, n + 1)]
    value = 1 / mp.fprod(terms)
    if abs(mp.im(value)) > mp.mpf(10) ** (-mp.mp.dps + 8):
        raise ArithmeticError("Conjugate determinant factors did not pair.")
    return mp.re(value)


def product_responses(b_squared):
    b = mp.sqrt(b_squared)
    n = (b_squared + 1) // 2
    a = mp.pi / (2 * b)  # x/2 is the argument of D_b
    sec2 = [1 / mp.cos(mp.pi / b_squared * (mp.mpf(n + 1) / 2 - j)) ** 2
            for j in range(1, n + 1)]
    return (2 * a ** 2 * mp.fsum(sec2),
            2 * a ** 4 * mp.fsum(4 * s - 6 * s ** 2 for s in sec2))


def kernel(t, b):
    """Exponential form avoids overflow; call integrands for t=0 limits."""
    q = b + 1 / b
    return (2 * mp.exp(-q * t / 2) * (-mp.expm1(-q * t))
            / ((-mp.expm1(-2 * b * t)) * (-mp.expm1(-2 * t / b))))


def integrated_responses(b_squared, mass, digits):
    """Finite quadrature and rigorous omitted-tail inequalities, separately."""
    with mp.workdps(digits):
        b = mp.sqrt(b_squared)
        x = mp.mpf(mass)
        q = b + 1 / b
        c = q / 2
        cutoff = (digits + 15) * mp.log(10) / c
        points = [mp.mpf(0), mp.mpf(1), mp.mpf(4), mp.mpf(10), cutoff]
        h2 = mp.quad(lambda t: q if t == 0 else 2 * t * kernel(t, b), points)
        h4 = mp.quad(lambda t: 0 if t == 0 else -2 * t ** 3 * kernel(t, b), points)
        delta = mp.quad(lambda t: q * x ** 2 / 2 if t == 0
                        else 4 * kernel(t, b) / t * mp.sin(x * t / 2) ** 2, points)
        d = (-mp.expm1(-2 * b * cutoff)) * (-mp.expm1(-2 * cutoff / b))
        bounds = {
            "Fxx": 4 * mp.gammainc(2, c * cutoff, mp.inf) / (d * c ** 2),
            "Fxxxx": 4 * mp.gammainc(4, c * cutoff, mp.inf) / (d * c ** 4),
            "delta_F": 8 * mp.exp(-c * cutoff) / (d * c * cutoff),
        }
        return {"Fxx": +h2, "Fxxxx": +h4, "delta_F": +delta,
                "cutoff": +cutoff, "tail_bounds": bounds}


def scaled_difference(a, b):
    return abs(a - b) / max(mp.mpf(1), abs(a), abs(b))


def calculate(b_squared=1, mass="0.2", digits=50, wrong=False):
    mp.mp.dps = digits
    x = mp.mpf(mass)
    computed = integrated_responses(b_squared, mass, digits)
    repeated = integrated_responses(b_squared, mass, digits + 20)
    z0 = determinant(mp.mpf(0), b_squared) ** 2
    z_probe = determinant(x / 2, b_squared) ** 2
    delta_product = -mp.log(z_probe / z0)
    h2, h4 = product_responses(b_squared)
    tests = [scaled_difference(computed["Fxx"], h2),
             scaled_difference(computed["Fxxxx"], h4),
             scaled_difference(computed["delta_F"], delta_product)]
    repeats = {key: scaled_difference(computed[key], repeated[key])
               for key in ("Fxx", "Fxxxx", "delta_F")}
    tests.extend(repeats.values())
    tolerance = mp.mpf(10) ** (-(digits - 15))
    if max(tests) >= tolerance:
        raise ArithmeticError("Finite product, quadrature or precision check failed.")

    # Finite differences intentionally retain a finite-step bias. They are not
    # included among the six high-precision comparisons above.
    def delta(z):
        return -2 * mp.log(determinant(z / 2, b_squared)
                           / determinant(mp.mpf(0), b_squared))

    differences = []
    for step in ("0.08", "0.04", "0.02"):
        h = mp.mpf(step)
        second = 2 * delta(h) / h ** 2
        fourth = (2 * delta(2 * h) - 8 * delta(h)) / h ** 4
        differences.append({"step": h, "Fxx": second, "Fxxxx": fourth,
                            "second_bias": second - h2, "fourth_bias": fourth - h4})

    out = {
        "model": "compact U(1)_1 x U(1)_-1; positive-H toric sphere",
        "source": "T=diag(1,1,-1,-1)/2; x=L*sigma_F; D_b argument x/2",
        "b_squared": b_squared, "probe_mass_x": x, "working_digits": digits,
        "derivative_locus": "x=0, fixed b and L; these are mass derivatives",
        "Z_zero_mass": z0, "Z_probe_product": z_probe,
        "Z_probe_integral": z0 * mp.exp(-computed["delta_F"]),
        "Fxx_integral": computed["Fxx"], "Fxxxx_integral": computed["Fxxxx"],
        "Fxx_product": h2, "Fxxxx_product": h4,
        "integration_cutoff": computed["cutoff"],
        "analytic_tail_bounds": computed["tail_bounds"],
        "precision_repeat_scaled_differences": repeats,
        "high_precision_comparisons": 6, "comparison_tolerance": tolerance,
        "max_scaled_difference": max(tests), "finite_step_diagnostics": differences,
    }
    if b_squared == 1:
        out["round_cT"] = 32 * computed["Fxx"] / mp.pi ** 2
    if wrong:
        # Actually differentiate the wrongly sourced determinant, not just a
        # modified printed target. It is F_correct(2*x).
        wrong_f = lambda z: -2 * mp.log(determinant(z, b_squared))
        wrong_h2 = mp.diff(wrong_f, 0, 2)
        wrong_h4 = mp.diff(wrong_f, 0, 4)
        caught = (abs(wrong_h2 - h2) > 1 and abs(wrong_h4 - h4) > 1
                  and scaled_difference(wrong_h2, 4 * h2) < tolerance
                  and scaled_difference(wrong_h4, 16 * h4) < tolerance)
        if not caught:
            raise ArithmeticError("The deliberate source-normalization error was not detected.")
        out["deliberate_failure"] = {"wrong_Fxx": wrong_h2, "wrong_Fxxxx": wrong_h4,
                                     "second_ratio": wrong_h2 / h2,
                                     "fourth_ratio": wrong_h4 / h4, "detected": bool(caught)}
    return out


def printable(value):
    if isinstance(value, dict):
        return {key: printable(item) for key, item in value.items()}
    if isinstance(value, list):
        return [printable(item) for item in value]
    if isinstance(value, mp.mpf):
        return mp.nstr(value, 25)
    return value


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--b-squared", type=int, choices=(1, 3, 5, 7, 9), default=1)
    parser.add_argument("--mass", default="0.2", help="real dimensionless probe source; |x|<=2")
    parser.add_argument("--dps", type=int, default=50, help="working decimal digits, 40 to 100")
    parser.add_argument("--wrong-mass-normalization", action="store_true")
    args = parser.parse_args()
    if not 40 <= args.dps <= 100:
        parser.error("--dps must lie between 40 and 100")
    try:
        with mp.workdps(args.dps):
            x = mp.mpf(args.mass)
            if not mp.isfinite(x) or abs(x) > 2:
                raise ValueError
    except (ValueError, TypeError):
        parser.error("--mass must be finite and real, with absolute value at most 2")
    print(json.dumps(printable(calculate(args.b_squared, args.mass, args.dps,
                                        args.wrong_mass_normalization)), indent=2))


if __name__ == "__main__":
    main()