"""Oriented SW periods on Re(u/Lambda^2)>1; mpmath 1.3.0.

All integrations/ODE steps use z=u/Lambda^2. A=a/Lambda and
D=a_D/Lambda. Their derivatives below are with respect to z.
No network, site imports, hidden data or automatic file writes.
"""
import argparse
import json
import mpmath as mp


def contours(z):
    """Both-sheet contour integrals after removing square-root endpoints."""
    c = mp.sqrt(2)/mp.pi
    den = lambda t: mp.sqrt(2+(z-1)*mp.sin(t)**2)
    a = c*mp.quad(lambda t: mp.sqrt(z-mp.cos(t)), [0, mp.pi/2, mp.pi])
    ap = c/2*mp.quad(lambda t: 1/mp.sqrt(z-mp.cos(t)), [0, mp.pi/2, mp.pi])
    d = 2j*c*(z-1)*mp.quad(lambda t: mp.cos(t)**2/den(t), [0, mp.pi/2])
    dp = 1j*c*mp.quad(lambda t: 1/den(t), [0, mp.pi/2])
    return [a, ap, d, dp]


def elliptic(z):
    m = 2/(z+1)  # parameter, not modulus
    c = mp.sqrt(2)/(mp.pi*mp.sqrt(z+1))
    return [2*mp.sqrt(2*(z+1))*mp.ellipe(m)/mp.pi, c*mp.ellipk(m),
            2j*mp.sqrt(2*(z+1))*(mp.ellipk(1-m)-mp.ellipe(1-m))/mp.pi,
            1j*c*mp.ellipk(1-m)]


def step(z, p, dp, h, order):
    coefficients = [p, dp]
    for n in range(order-1):
        coefficients.append(-(2*z*n*(n+1)*coefficients[n+1]
                              +(n*(n-1)+mp.mpf('0.25'))*coefficients[n])
                            /((z*z-1)*(n+2)*(n+1)))
    return [mp.polyval(list(reversed(coefficients)), h),
            mp.polyval(list(reversed([n*coefficients[n]
                                     for n in range(1, len(coefficients))])), h)]


def continuation(target, initial, order, fraction):
    z, data, steps = mp.mpc(3), list(initial), 0
    threshold = mp.power(10, -mp.mp.dps+5)
    while abs(target-z) > threshold:
        delta = target-z
        radius = min(abs(z-1), abs(z+1))
        h = delta*min(1, fraction*radius/abs(delta))
        data[:2] = step(z, data[0], data[1], h, order)
        data[2:] = step(z, data[2], data[3], h, order)
        z += h
        steps += 1
        if steps > 3000:
            raise ValueError('Too many disks; use a point farther from z=1.')
    return data, steps


def scaled_error(a, b):
    return abs(a-b)/max(1, abs(a), abs(b))


def run(args):
    mp.mp.dps = args.dps
    target = mp.mpc(args.real, args.imag)
    scale = mp.mpf(args.scale)
    if not (mp.isfinite(target) and target.real >= mp.mpf('1.0001')
            and abs(target) <= 100 and mp.isfinite(scale) and scale > 0):
        raise ValueError('Require Re(z)>=1.0001, |z|<=100, finite positive Lambda/m0.')
    if not (40 <= args.dps <= 100):
        raise ValueError('Choose 40 to100 decimal digits.')
    nstr = lambda x: mp.nstr(x, 14)
    pair = lambda x: {'real': nstr(mp.re(x)), 'imag': nstr(mp.im(x))}
    base, direct = contours(mp.mpf(3)), contours(target)
    tol = mp.mpf('1e-25')
    rows = []
    for order, fraction in [(40, mp.mpf('.2')), (60, mp.mpf('.1'))]:
        values, disks = continuation(target, base, order, fraction)
        err = max(scaled_error(a, b) for a, b in zip(values, direct))
        werr = abs(values[0]*values[3]-values[2]*values[1]-2j/mp.pi)
        assert err < tol and werr < tol
        rows.append({'order': order, 'disk_fraction': nstr(fraction),
                     'disks': disks, 'max_scaled_error': nstr(err),
                     'absolute_W_error': nstr(werr)})
    closed = elliptic(target)
    ellerr = max(scaled_error(a, b) for a, b in zip(closed, direct))
    assert ellerr < tol
    with mp.workdps(args.dps+20):
        stable = contours(mp.mpc(args.real, args.imag))
    precision_error = max(scaled_error(a, b) for a, b in zip(stable, direct))
    assert precision_error < tol
    tau = direct[3]/direct[1]
    assert mp.im(tau) > 0 and abs(base[3]/base[1]-1j) < tol
    chosen = direct[2]/direct[0] if args.swap_period_ratio else tau
    base_chosen = base[2]/base[0] if args.swap_period_ratio else base[3]/base[1]
    report = {
        'theory': 'flat 4d N=2 pure SU(2), first-paper period basis',
        'precision_digits': args.dps, 'comparison_tolerance': nstr(tol),
        'z_u_over_Lambda_squared': pair(target), 'Lambda_over_m0': nstr(scale),
        'a_over_m0': pair(scale*direct[0]), 'aD_over_m0': pair(scale*direct[2]),
        'm0_times_da_du': pair(direct[1]/scale),
        'm0_times_daD_du': pair(direct[3]/scale),
        'tau_from_derivatives': pair(tau), 'period_ratio': pair(direct[2]/direct[0]),
        'ODE_comparisons': rows, 'elliptic_scaled_error': nstr(ellerr),
        'precision_repeat_scaled_error': nstr(precision_error),
        'reported_coupling': pair(chosen),
        'basepoint_tau_equals_i_test_passes': bool(abs(base_chosen-1j) < tol),
        'intentional_ratio_error_detected': bool(args.swap_period_ratio and abs(base_chosen-1j)>mp.mpf('.1')),
        'scope': 'Finite diagnostics, not certified interval bounds or a BPS spectrum.'
    }
    weak = []
    for z in [20, 40, 80]:
        a, ap, d, dp = contours(mp.mpf(z))
        deviation = a/mp.sqrt(2*z)-1
        weak.append({'z': z, 'z_squared_times_electric_relative_shift': nstr(z*z*deviation),
                     'shift_after_first_correction': nstr(deviation+mp.mpf(1)/(16*z*z))})
    report['weak_coupling'] = weak
    if args.instanton:
        # The supplied one-box Euler factors, in the NY coordinate b=-i alpha.
        alpha = mp.mpf(1)
        b = -1j*alpha
        regulator = []
        for h in [mp.mpf('.2'), mp.mpf('.1'), mp.mpf('.05')]:
            e1, e2 = h, 2*h
            E = e1+e2
            terms = [1/(e1*e2*2*b*(E-2*b)),
                     1/(e1*e2*(-2*b)*(E+2*b))]
            coefficient = e1*e2*sum(terms)
            assert abs(coefficient-2/(E*E+4*alpha*alpha)) < tol
            regulator.append({'h_over_alpha': nstr(h),
                              'alpha_squared_Fscript_q_coefficient': nstr(mp.re(alpha*alpha*coefficient))})
        extracted = []
        for z in [20, 40, 80]:
            a, ap, d, dp = contours(mp.mpf(z))
            perturbative_d = 1j*a/mp.pi*(mp.log(4*a*a)-2)
            value = a**3*(d-perturbative_d)/1j
            extracted.append({'z': z, 'C_a': nstr(mp.re(value)),
                              'difference_from_limit': nstr(mp.re(value)-1/(4*mp.pi))})
        report['instanton_extension_requires_SW16'] = {
            'dictionary': 'alpha=a/sqrt(2), b=-i alpha, q=Lambda^4/4, Fsite=Fscript/(2 pi i)',
            'coefficientwise_regulator': regulator, 'period_extraction': extracted,
            'C_a_limit': nstr(1/(4*mp.pi)),
            'limit_order': 'Extract q coefficient, then remove epsilon1 and epsilon2.'}
    return report


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--real', default='2')
    parser.add_argument('--imag', default='0.5')
    parser.add_argument('--scale', default='1', help='Lambda/m0 in a fixed mass unit')
    parser.add_argument('--dps', type=int, default=50)
    parser.add_argument('--swap-period-ratio', action='store_true')
    parser.add_argument('--instanton', action='store_true')
    args = parser.parse_args()
    try:
        output = run(args)
    except ValueError as error:
        parser.error(str(error))
    print(json.dumps(output, indent=2))


if __name__ == '__main__':
    main()
