#!/usr/bin/env python3
"""Reproduce the Gauss starter with local series and gamma-function controls.

Python 3.10+ and mpmath 1.3.0 are the supported baseline. No network or data
files are needed. The finite series are evaluated independently at three
interior match points. These are numerical checks, not interval enclosures.
Exit 0 means all declared checks passed; exit 1 means a numerical check failed;
argparse uses exit 2 for invalid input. JSON results contain decimal strings.
"""

import argparse
import json
import sys

import mpmath as mp


def series(a, b, c, z, terms):
    """Return a truncated 2F1 polynomial and its first two derivatives."""
    coefficients = [mp.mpf(1)]
    for k in range(1, terms):
        coefficients.append(coefficients[-1] * (a+k-1) * (b+k-1) / ((c+k-1)*k))
    value = mp.polyval(list(reversed(coefficients)), z)
    first = mp.polyval(list(reversed([k*coefficients[k] for k in range(1, terms)])), z)
    second = mp.polyval(list(reversed([k*(k-1)*coefficients[k] for k in range(2, terms)])), z)
    return value, first, second


def wronskian(f, g):
    return f[0]*g[1] - f[1]*g[0]


def evaluate(terms, dps, tolerance):
    mp.mp.dps = dps
    a, b, c = mp.mpf(1)/4, mp.mpf(1)/3, mp.mpf(5)/4
    delta = c-a-b
    exact_a = mp.gamma(c)*mp.gamma(delta)/(mp.gamma(c-a)*mp.gamma(c-b))
    exact_b = mp.gamma(c)*mp.gamma(-delta)/(mp.gamma(a)*mp.gamma(b))
    rows = []
    errors = []
    for z in [mp.mpf(1)/3, mp.mpf(1)/2, mp.mpf(2)/3]:
        s = 1-z
        f0 = series(a, b, c, z, terms)
        h = series(a, b, 1-delta, s, terms)
        f1 = (h[0], -h[1], h[2])
        h = series(c-a, c-b, 1+delta, s, terms)
        g1 = (s**delta*h[0],
              -delta*s**(delta-1)*h[0]-s**delta*h[1],
              delta*(delta-1)*s**(delta-2)*h[0]
              +2*delta*s**(delta-1)*h[1]+s**delta*h[2])
        denominator = wronskian(f1, g1)
        actual_a = wronskian(f0, g1)/denominator
        actual_b = wronskian(f1, f0)/denominator
        ab_error = max(abs(actual_a/exact_a-1), abs(actual_b/exact_b-1))
        abel_error = abs(denominator/(-delta*z**(-c)*s**(delta-1))-1)
        ode_errors = []
        for f in (f0, f1, g1):
            pieces = [z*s*f[2], (c-(a+b+1)*z)*f[1], -a*b*f[0]]
            ode_errors.append(abs(sum(pieces))/(1+sum(abs(p) for p in pieces)))
        # A basis phase must be compensated in its connection coefficient.
        phase = mp.exp(2j*mp.pi*delta)
        transformed_b = wronskian(f1, f0)/wronskian(f1, [phase*v for v in g1])
        branch_error = abs(transformed_b*phase/actual_b-1)
        maximum = max(ab_error, abel_error, *ode_errors, branch_error)
        errors.append(maximum)
        rows.append({"z": mp.nstr(z, 12), "A": mp.nstr(actual_a, 18),
                     "B": mp.nstr(actual_b, 18),
                     "coefficient_relative_error": mp.nstr(ab_error, 8),
                     "abel_relative_error": mp.nstr(abel_error, 8),
                     "ode_scaled_residual": mp.nstr(max(ode_errors), 8),
                     "basis_phase_error": mp.nstr(branch_error, 8)})
    passed = all(mp.isfinite(error) and error <= tolerance for error in errors)
    return {"status": "PASS" if passed else "FAIL", "terms": terms, "dps": dps,
            "tolerance": mp.nstr(tolerance, 8),
            "reference_A": mp.nstr(exact_a, 18), "reference_B": mp.nstr(exact_b, 18),
            "maximum_checked_error": mp.nstr(max(errors), 8), "match_points": rows}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--terms", type=int, default=240)
    parser.add_argument("--dps", type=int, default=60)
    parser.add_argument("--tolerance", default="1e-30")
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()
    if args.terms < 8 or args.dps < 20:
        parser.error("--terms must be at least 8 and --dps at least 20")
    mp.mp.dps = args.dps
    try:
        tolerance = mp.mpf(args.tolerance)
    except (ValueError, TypeError):
        parser.error("--tolerance must be a finite positive number")
    if not mp.isfinite(tolerance) or tolerance <= 0:
        parser.error("--tolerance must be a finite positive number")
    result = evaluate(args.terms, args.dps, tolerance)
    if args.json:
        print(json.dumps(result, indent=2))
    else:
        print(f"{result['status']}: {args.terms} terms, {args.dps} working digits")
        print(f"A = {result['reference_A']}; B = {result['reference_B']}")
        for row in result['match_points']:
            print(f"z={row['z']}: coefficient error {row['coefficient_relative_error']}, "
                  f"Abel error {row['abel_relative_error']}, ODE residual {row['ode_scaled_residual']}")
        print(f"Maximum checked error: {result['maximum_checked_error']}; tolerance {result['tolerance']}")
    return 0 if result['status'] == 'PASS' else 1


if __name__ == "__main__":
    sys.exit(main())