#!/usr/bin/env python3
"""Generic Painleve VI tau ratios: Fourier projection versus contour Nystrom.

Requires Python 3.10+ and mpmath 1.3.0. Source: Gavrylenko--Lisovyy,
arXiv:1608.00958v3, Theorem A, (1.4), (1.6)--(1.8). Source residue
exponents are +/-theta; the website's theta values are twice these.

The two methods independently assemble the same explicit hypergeometric
Fredholm operator: power-series convolution versus direct kernel quadrature.
The sigma-form PVI residual is a separate differential-equation check, not
reconstructed Schlesinger transport. Results are numerical observations,
not interval enclosures or a scalar spectral condition. Exit 0: all declared
checks pass; 1: failed numerical check; 2: invalid/unsupported input.
"""

import argparse
import json
import sys
from time import perf_counter

import mpmath as mp

SOURCE = 'https://arxiv.org/pdf/1608.00958'


def trace(matrix):
    return sum(matrix[k, k] for k in range(matrix.rows))


def fmt(value):
    """Decimal strings keep output independent of machine float conversion."""
    if isinstance(value, (mp.mpc, complex)):
        return {'real': mp.nstr(mp.re(value), 20), 'imag': mp.nstr(mp.im(value), 20)}
    return mp.nstr(value, 12)


def distance(first, second):
    return abs(first-second)/(1+abs(second))


class Problem:
    def __init__(self, args):
        self.theta0, self.thetat, self.theta1, self.thetainf = [
            mp.mpf(getattr(args, name)) for name in
            ('theta0', 'thetat', 'theta1', 'thetainf')]
        self.sigma, self.eta = mp.mpf(args.sigma), mp.mpf(args.eta)
        self.alpha = self.sigma**2-self.theta0**2-self.thetat**2
        self.beta = -2*self.thetat*self.theta1
        self.tref, self.t = mp.mpf(args.reference), mp.mpf(args.time)
        self.radii = [mp.mpf(args.radius), mp.mpf(args.second_radius)]
        self.tolerance = mp.mpf(args.tolerance)

    def validate(self):
        values = [self.theta0, self.thetat, self.theta1, self.thetainf,
                  self.sigma, self.eta, self.tref, self.t, self.tolerance, *self.radii]
        if not all(mp.isfinite(x) for x in values) or self.tolerance <= 0:
            raise ValueError('all numerical inputs must be finite; tolerance must be positive')
        if not 0 < abs(self.sigma) < mp.mpf('0.5'):
            raise ValueError('this real chart requires 0 < abs(sigma) < 1/2')
        if not (0 < self.tref < 1 and 0 < self.t < 1) or self.tref == self.t:
            raise ValueError('distinct time and reference must lie in (0,1)')
        if any(not max(self.tref, self.t) < r < 1 for r in self.radii):
            raise ValueError('each radius must satisfy max(time,reference) < R < 1')
        if self.radii[0] == self.radii[1]:
            raise ValueError('choose distinct contour radii')
        combinations = [a+sign_b*b+sign_s*self.sigma
                        for a, b in [(self.theta0, self.thetat),
                                     (self.theta1, self.thetainf)]
                        for sign_b in [-1, 1] for sign_s in [-1, 1]]
        combinations += [2*x for x in [self.theta0, self.thetat,
                                        self.theta1, self.thetainf]]
        self.genericity_gap = min(abs(x-mp.nint(x)) for x in combinations)
        if self.genericity_gap <= 100*mp.eps:
            raise ValueError('resonant exponent combination outside this implemented chart')
        composite_trace = 2*mp.cos(2*mp.pi*self.sigma)
        self.noncommutation_gap = min(abs(composite_trace-2*mp.cos(
            2*mp.pi*(self.theta0+sign*self.thetat))) for sign in [-1, 1])
        if self.noncommutation_gap <= 100*mp.eps:
            raise ValueError('chosen composite trace does not certify noncommutation')

    def prefactor(self, t):
        # Real logarithms for positive t and 1-t. Never omit the second power.
        return mp.exp(self.alpha*mp.log(t)+self.beta*mp.log(1-t))


def hypergeom_coefficients(a, b, sigma, degree):
    """Taylor coefficients of the source's K(a,b,sigma;z), via Pochhammer products."""
    result = []
    for n in range(degree+1):
        matrix = mp.zeros(2)
        for row, sign in enumerate([1, -1]):
            x, y, c = a+b+sign*sigma, a-b+sign*sigma, 2*sign*sigma
            matrix[row, row] = mp.rf(x, n)*mp.rf(y, n)/(mp.rf(c, n)*mp.factorial(n))
            if n:
                factor = sign*(b*b-(a+sign*sigma)**2)/(2*sigma*(1+2*sign*sigma))
                matrix[row, 1-row] = factor*mp.rf(1+x, n-1)*mp.rf(1+y, n-1)/(
                    mp.rf(2+c, n-1)*mp.factorial(n-1))
        result.append(matrix)
    return result


def fourier_block(a, b, sigma, modes):
    """Coefficients of (K(z) K(w)^-1-I)/(z-w), without contour sampling."""
    coefficients = hypergeom_coefficients(a, b, sigma, 2*modes-1)
    inverse = [mp.eye(2)]
    for n in range(1, 2*modes):
        inverse.append(-sum((coefficients[k]*inverse[n-k]
                             for k in range(1, n+1)), mp.zeros(2)))
    block = mp.zeros(2*modes)
    for i in range(modes):
        for j in range(modes):
            value = sum((coefficients[i+j+1-k]*inverse[k]
                         for k in range(j+1)), mp.zeros(2))
            for row in range(2):
                for column in range(2):
                    block[2*i+row, 2*j+column] = value[row, column]
    return block


class Fourier:
    def __init__(self, problem, modes):
        self.problem, self.modes = problem, modes
        self.a = fourier_block(problem.theta1, problem.thetainf, problem.sigma, modes)
        self.base_d = fourier_block(problem.thetat, problem.theta0, -problem.sigma, modes)

    def matrices(self, t, derivatives=0):
        p, modes = self.problem, self.modes
        exponents = [p.sigma, -p.sigma]
        ds = [mp.zeros(2*modes) for _ in range(derivatives+1)]
        for i in range(2*modes):
            for j in range(2*modes):
                power = i//2+j//2+1-exponents[i % 2]+exponents[j % 2]
                base = self.base_d[i, j]*mp.exp(1j*p.eta*(i % 2-j % 2))
                for k in range(derivatives+1):
                    ds[k][i, j] = base*mp.exp((power-k)*mp.log(t))*mp.fprod(
                        power-r for r in range(k))
        return [mp.eye(2*modes)-self.a*ds[0]]+[-self.a*d for d in ds[1:]]

    def tau(self, t):
        return self.problem.prefactor(t)*mp.det(self.matrices(t)[0])

    def differential(self, t):
        """Analytic finite-matrix derivatives through third order of log tau."""
        p = self.problem
        f, fp, fpp, fppp = self.matrices(t, 3)
        inverse = f**-1
        b, c, e = inverse*fp, inverse*fpp, inverse*fppp
        h = trace(b)+p.alpha/t-p.beta/(1-t)
        hp = trace(c-b*b)-p.alpha/t**2-p.beta/(1-t)**2
        hpp = trace(e-3*b*c+2*b*b*b)+2*p.alpha/t**3-2*p.beta/(1-t)**3
        zeta = t*(t-1)*h
        zetap = (2*t-1)*h+t*(t-1)*hp
        zetapp = 2*h+2*(2*t-1)*hp+t*(t-1)*hpp
        v = zetap+p.theta0**2+p.thetat**2+p.theta1**2-p.thetainf**2
        gram = mp.matrix([[2*p.theta0**2, t*zetap-zeta, v],
                          [t*zetap-zeta, 2*p.thetat**2, (t-1)*zetap-zeta],
                          [v, (t-1)*zetap-zeta, 2*p.theta1**2]])
        lhs, rhs = (t*(t-1)*zetapp)**2, -2*mp.det(gram)
        residual = abs(lhs-rhs)/(1+abs(lhs)+abs(rhs))
        return {'h': h, 'hp': hp, 'hpp': hpp, 'zeta': zeta,
                'sigma_residual': residual}


def direct_parametrix(a, b, sigma, z):
    """Direct hypergeometric values and derivatives, without Taylor truncation."""
    k, kp = mp.zeros(2), mp.zeros(2)
    for row, sign in enumerate([1, -1]):
        x, y, c = a+b+sign*sigma, a-b+sign*sigma, 2*sign*sigma
        k[row, row] = mp.hyp2f1(x, y, c, z)
        kp[row, row] = x*y/c*mp.hyp2f1(x+1, y+1, c+1, z)
        factor = sign*(b*b-(a+sign*sigma)**2)/(2*sigma*(1+2*sign*sigma))
        value = mp.hyp2f1(1+x, 1+y, 2+c, z)
        k[row, 1-row] = factor*z*value
        kp[row, 1-row] = factor*(value+z*(1+x)*(1+y)/(2+c)*mp.hyp2f1(
            2+x, 2+y, 3+c, z))
    # Source (1.7): det K=(1-z)^(-2a). This avoids sharing the Fourier inverse.
    adjugate = mp.matrix([[k[1, 1], -k[0, 1]], [-k[1, 0], k[0, 0]]])
    inverse = (1-z)**(2*a)*adjugate
    normalization_error = mp.norm(k*inverse-mp.eye(2))
    return k, kp, inverse, normalization_error


def nystrom(problem, t, nodes, radius):
    p = problem
    points = [radius*mp.exp(2j*mp.pi*k/nodes) for k in range(nodes)]
    twist = mp.diag([t**(-p.sigma)*mp.exp(-1j*p.eta/2),
                     t**p.sigma*mp.exp(1j*p.eta/2)])
    twist_inverse = twist**-1
    right, left, normalization_errors = [], [], []
    for z in points:
        k, kp, inverse, error = direct_parametrix(p.theta1, p.thetainf, p.sigma, z)
        right.append((k, kp, inverse))
        normalization_errors.append(error)
        k, kp, inverse, error = direct_parametrix(p.thetat, p.theta0, -p.sigma, t/z)
        left.append((twist*k*twist_inverse,
                     (-t/z**2)*twist*kp*twist_inverse,
                     twist*inverse*twist_inverse))
        normalization_errors.append(error)
    a, d = mp.zeros(2*nodes), mp.zeros(2*nodes)
    for i, z in enumerate(points):
        for j, w in enumerate(points):
            if i == j:
                aa = right[i][1]*right[i][2]
                dd = -left[i][1]*left[i][2]
            else:
                aa = (right[i][0]*right[j][2]-mp.eye(2))/(z-w)
                dd = (mp.eye(2)-left[i][0]*left[j][2])/(z-w)
            # dz/(2*pi*i) gives weight w/nodes on a counterclockwise circle.
            for row in range(2):
                for column in range(2):
                    a[2*i+row, 2*j+column] = aa[row, column]*w/nodes
                    d[2*i+row, 2*j+column] = dd[row, column]*w/nodes
    determinant = mp.det(mp.eye(2*nodes)-a*d)
    return p.prefactor(t)*determinant, max(normalization_errors)


def run(args):
    mp.mp.dps = args.dps
    p = Problem(args)
    p.validate()
    records, checks = [], {}
    mode_levels = sorted({max(2, args.modes-4), max(3, args.modes-2), args.modes})
    node_levels = sorted({max(4, args.nodes-12), max(6, args.nodes-6), args.nodes})
    fourier_ratios, nystrom_ratios = [], []
    for modes in mode_levels:
        method = Fourier(p, modes)
        ratio = method.tau(p.t)/method.tau(p.tref)
        differential = method.differential(p.t)
        fourier_ratios.append(ratio)
        records.append({'method': 'Fourier', 'modes': modes, 'dps': args.dps,
                        'tau_ratio': fmt(ratio),
                        'sigma_scaled_residual': fmt(differential['sigma_residual'])})
    checks['fourier_last_refinement'] = distance(fourier_ratios[-2], fourier_ratios[-1])
    checks['sigma_form'] = differential['sigma_residual']
    for nodes in node_levels:
        tau, error = nystrom(p, p.t, nodes, p.radii[0])
        reference, reference_error = nystrom(p, p.tref, nodes, p.radii[0])
        ratio = tau/reference
        nystrom_ratios.append(ratio)
        records.append({'method': 'Nystrom', 'nodes': nodes, 'radius': fmt(p.radii[0]),
                        'dps': args.dps, 'tau_ratio': fmt(ratio),
                        'parametrix_normalization_error': fmt(max(error, reference_error))})
    checks['nystrom_last_refinement'] = distance(nystrom_ratios[-2], nystrom_ratios[-1])
    checks['method_agreement'] = distance(nystrom_ratios[-1], fourier_ratios[-1])
    checks['parametrix_normalization'] = max(error, reference_error)
    tau, error2 = nystrom(p, p.t, args.nodes, p.radii[1])
    reference, reference_error2 = nystrom(p, p.tref, args.nodes, p.radii[1])
    second_ratio = tau/reference
    records.append({'method': 'Nystrom', 'nodes': args.nodes, 'radius': fmt(p.radii[1]),
                    'dps': args.dps, 'tau_ratio': fmt(second_ratio)})
    checks['contour_invariance'] = distance(second_ratio, nystrom_ratios[-1])
    checks['parametrix_normalization'] = max(checks['parametrix_normalization'],
                                             error2, reference_error2)
    # Same cutoff, different arithmetic precision: isolate rounding effects.
    with mp.workdps(args.dps+20):
        high = Problem(args)
        high_method = Fourier(high, args.modes)
        high_ratio = high_method.tau(high.t)/high_method.tau(high.tref)
        high_tau, _ = nystrom(high, high.t, args.nodes, high.radii[0])
        high_reference, _ = nystrom(high, high.tref, args.nodes, high.radii[0])
        high_nystrom_ratio = high_tau/high_reference
        checks['fourier_precision_change'] = distance(fourier_ratios[-1], high_ratio)
        checks['nystrom_precision_change'] = distance(nystrom_ratios[-1], high_nystrom_ratio)
        records += [{'method': 'Fourier', 'modes': args.modes, 'dps': args.dps+20,
                     'tau_ratio': fmt(high_ratio)},
                    {'method': 'Nystrom', 'nodes': args.nodes, 'radius': fmt(high.radii[0]),
                     'dps': args.dps+20, 'tau_ratio': fmt(high_nystrom_ratio)}]
        # Independent differentiation of the same finite scalar determinant
        # checks the matrix trace derivative implementation; it is not an
        # additional independent tau construction.
        hd = high_method.differential(high.t)
        log_tau = lambda t: mp.log(high_method.tau(t))
        derivative_errors = [distance(mp.diff(log_tau, high.t, order), hd[name])
                             for order, name in [(1, 'h'), (2, 'hp'), (3, 'hpp')]]
        checks['finite_matrix_derivative_consistency'] = max(derivative_errors)
    passed = all(mp.isfinite(value) and value <= p.tolerance for value in checks.values())
    return {'status': 'PASS' if passed else 'FAIL', 'source': SOURCE,
            'source_conventions': 'residue eigenvalues +/-theta_GL; website theta=2*theta_GL',
            'parameters': {name: getattr(args, name) for name in
                           ['theta0', 'thetat', 'theta1', 'thetainf', 'sigma', 'eta']},
            'time': args.time, 'reference_time': args.reference,
            'normalization': 'tau(t)/tau(reference), real Log(t) and Log(1-t)',
            'prefactor_exponents': {'t': fmt(p.alpha), '1-t': fmt(p.beta)},
            'genericity_distance_from_integers': fmt(p.genericity_gap),
            'noncommuting_composite_trace_gap': fmt(p.noncommutation_gap),
            'tolerance': fmt(p.tolerance), 'checks': {key: fmt(value) for key, value in checks.items()},
            'maximum_checked_error': fmt(max(checks.values())), 'refinements': records,
            'scope': 'Numerical convergence and sigma-form check; no rigorous enclosure, '
                     'reconstructed Schlesinger transport, or scalar spectral assertion.'}


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    for name, default in [('theta0', '0.17'), ('thetat', '0.23'), ('theta1', '0.19'),
                          ('thetainf', '0.31'), ('sigma', '0.27'), ('eta', '0.41'),
                          ('time', '0.12'), ('reference', '0.08'), ('radius', '0.36'),
                          ('second-radius', '0.42'), ('tolerance', '1e-9')]:
        parser.add_argument('--'+name, default=default)
    parser.add_argument('--modes', type=int, default=12)
    parser.add_argument('--nodes', type=int, default=28)
    parser.add_argument('--dps', type=int, default=40)
    parser.add_argument('--json', action='store_true')
    args = parser.parse_args()
    if args.modes < 4 or args.nodes < 8 or args.dps < 25:
        parser.error('require modes >= 4, nodes >= 8, dps >= 25')
    try:
        mp.mp.dps = args.dps
        Problem(args).validate()
    except (ValueError, TypeError) as error:
        parser.error(str(error))
    started = perf_counter()
    try:
        result = run(args)
    except (ValueError, ZeroDivisionError, ArithmeticError) as error:
        result = {'status': 'FAIL', 'error': str(error), 'source': SOURCE}
    result['elapsed_seconds'] = round(perf_counter()-started, 3)
    if args.json:
        print(json.dumps(result, indent=2))
    else:
        print(f"{result['status']}: generic Painleve VI Fredholm comparison")
        if 'error' in result:
            print(result['error'])
        else:
            for row in result['refinements']:
                print(row)
            print('Checks:', result['checks'])
            print('Maximum checked error:', result['maximum_checked_error'])
            print('Tolerance:', result['tolerance'])
            print(result['scope'])
        print('Elapsed seconds:', result['elapsed_seconds'])
    return 0 if result['status'] == 'PASS' else 1


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