"""Continue first-paper SU(2) periods along explicitly based loops.

Requires mpmath 1.3.0. No network, site imports or hidden data.
Only --figure explicitly writes a file. Other modes print JSON.
z=u/Lambda^2, A=a/Lambda, D=a_D/Lambda; primes mean d/dz.
"""
import argparse
import json
from pathlib import Path
import mpmath as mp


def initial_periods(base):
    c = mp.sqrt(2)/mp.pi
    den = lambda t: mp.sqrt(2+(base-1)*mp.sin(t)**2)
    a = c*mp.quad(lambda t: mp.sqrt(base-mp.cos(t)), [0, mp.pi/2, mp.pi])
    ap = c/2*mp.quad(lambda t: 1/mp.sqrt(base-mp.cos(t)), [0, mp.pi/2, mp.pi])
    d = 2j*c*(base-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 taylor_step(z, p, dp, h, order):
    c = [p, dp]
    for n in range(order-1):
        c.append(-(2*z*n*(n+1)*c[n+1]+(n*(n-1)+mp.mpf('.25'))*c[n])
                 /((z*z-1)*(n+2)*(n+1)))
    return [mp.polyval(list(reversed(c)), h),
            mp.polyval(list(reversed([n*c[n] for n in range(1,len(c))])), h)]


def circle(center, radius, segments=64):
    points = [center+radius*mp.exp(2j*mp.pi*j/segments) for j in range(segments)]
    return points+[mp.mpc(center+radius)]


def paths(base, radius):
    plus_stem = [mp.mpc(base), mp.mpc(1+radius)]
    minus_stem = [mp.mpc(base), base-1j, -1+radius-1j, mp.mpc(-1+radius)]
    return {
        'plus': plus_stem+circle(1,radius)[1:]+plus_stem[-2::-1],
        'minus_below_plus': minus_stem+circle(-1,radius)[1:]+minus_stem[-2::-1],
        'outer': circle(0,base)
    }


def continue_path(path, initial, order, fraction):
    data, z, count = list(initial), path[0], 0
    w_error, min_im_tau = mp.mpf(0), mp.inf
    for target in path[1:]:
        while abs(target-z)>mp.power(10,-mp.mp.dps+5):
            delta = target-z
            distance = min(abs(z-1),abs(z+1))
            h = delta*min(1,fraction*distance/abs(delta))
            data[:2] = taylor_step(z,*data[:2],h,order)
            data[2:] = taylor_step(z,*data[2:],h,order)
            z += h
            count += 1
            w_error = max(w_error,abs(data[0]*data[3]-data[2]*data[1]-2j/mp.pi))
            min_im_tau = min(min_im_tau,mp.im(data[3]/data[1]))
            if count>6000:
                raise ValueError('Too many continuation disks.')
        z = target
    return data, count, w_error, min_im_tau


def fundamental(v):
    return mp.matrix([[v[2],v[0]],[v[3],v[1]]])


def recover(initial, final):
    return (fundamental(initial)**-1*fundamental(final)).T


def maxabs(matrix):
    return max(abs(x) for x in matrix)


def integer_matrix(matrix):
    return [[int(mp.nint(mp.re(matrix[i,j]))) for j in range(2)] for i in range(2)]


def make_figure(filename, base, radius):
    """Render the same polygon paths; stems are traversed both ways."""
    loops = paths(base,radius)
    panels = [('plus','Around +1'),('minus_below_plus','Around -1; lower stem'),('outer','Around both punctures')]
    width, height = 1080, 400
    pieces = [f'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 {width} {height}" role="img" aria-labelledby="title desc">',
              '<title id="title">Three based continuation loops in the dimensionless u-plane</title>',
              '<desc id="desc">Each path starts at a real base greater than one. The minus loop approaches below plus one. All circles turn counterclockwise. Stems are followed outward and back.</desc>',
              '<rect width="1080" height="400" fill="#fff"/>',
              '<defs><marker id="arrow" markerWidth="8" markerHeight="8" refX="6" refY="3" orient="auto"><path d="M0 0 L6 3 L0 6" fill="none" stroke="#126a8a" stroke-width="1.3"/></marker></defs>']
    scale = 135/(float(base)+.3)
    for i,(key,title) in enumerate(panels):
        cx,cy = 180+360*i,210
        xy = lambda z:(cx+scale*float(mp.re(z)),cy-scale*float(mp.im(z)))
        pieces += [f'<text x="{cx}" y="30" text-anchor="middle" font-family="sans-serif" font-size="18" fill="#172d3c">{title}</text>',
                   f'<path d="M{cx-150} {cy} H{cx+150} M{cx} {cy-150} V{cy+150}" stroke="#a7b6c2" fill="none"/>']
        coords = [xy(z) for z in loops[key]]
        line = ' '.join(('M' if j==0 else 'L')+f'{x:.3f} {y:.3f}' for j,(x,y) in enumerate(coords))
        pieces.append(f'<path d="{line}" stroke="#126a8a" stroke-width="2.2" fill="none"/>')
        center,rad = (0,base) if key=='outer' else ((1,radius) if key=='plus' else (-1,radius))
        arc = [xy(center+rad*mp.exp(1j*mp.pi*j/16)) for j in [5,6,7]]
        d = ' '.join(('M' if j==0 else 'L')+f'{x:.3f} {y:.3f}' for j,(x,y) in enumerate(arc))
        pieces.append(f'<path d="{d}" stroke="#126a8a" stroke-width="2.2" fill="none" marker-end="url(#arrow)"/>')
        for val,label in [(-1,'-1'),(1,'+1')]:
            x,y = xy(val)
            pieces += [f'<circle cx="{x:.3f}" cy="{y:.3f}" r="4" fill="#b63d28"/>',f'<text x="{x:.3f}" y="{y+20:.3f}" text-anchor="middle" font-family="sans-serif" font-size="14" fill="#8c2d1c">{label}</text>']
        x,y = xy(base)
        pieces += [f'<circle cx="{x:.3f}" cy="{y:.3f}" r="4" fill="#172d3c"/>',f'<text x="{x:.3f}" y="{y-12:.3f}" text-anchor="middle" font-family="sans-serif" font-size="14" fill="#172d3c">P</text>']
    pieces.append('<text x="540" y="388" text-anchor="middle" font-family="sans-serif" font-size="15" fill="#172d3c">Horizontal: Re z. Vertical: Im z. The base and both punctures are shared; each panel shows one loop.</text></svg>')
    Path(filename).write_text('\n'.join(pieces)+'\n',encoding='utf-8')


def run(args):
    if not 40<=args.dps<=90:
        raise ValueError('Choose40 to90 decimal digits.')
    mp.mp.dps = args.dps
    base,radius = mp.mpf(args.base),mp.mpf(args.radius)
    if not (mp.isfinite(base) and 2<=base<=6 and mp.isfinite(radius) and mp.mpf('.15')<=radius<=mp.mpf('.7')):
        raise ValueError('Require finite2<=base<=6 and0.15<=radius<=0.7.')
    nstr = lambda x:mp.nstr(x,14)
    pair = lambda x:{'real':nstr(mp.re(x)),'imag':nstr(mp.im(x))}
    tolerance = mp.mpf('1e-24')
    initial = initial_periods(base)
    loops = paths(base,radius)
    expected = {'plus':mp.matrix([[1,0],[-2,1]]),
                'minus_below_plus':mp.matrix([[-1,2],[-2,3]]),
                'outer':mp.matrix([[-1,2],[0,-1]])}
    rows,finals = [],{}
    for name,path in loops.items():
        levels=[]
        for order,fraction in [(40,mp.mpf('.2')),(60,mp.mpf('.1'))]:
            final,steps,werr,minimum = continue_path(path,initial,order,fraction)
            M = recover(initial,final)
            residual = maxabs(M-expected[name])
            assert residual<tolerance and werr<tolerance and minimum>0
            levels.append({'order':order,'disk_fraction':nstr(fraction),'steps':steps,
                           'matrix_residual':nstr(residual),'max_W_error':nstr(werr),
                           'minimum_sampled_Im_tau':nstr(minimum)})
        finals[name] = final
        candidate = recover(initial,initial if args.reset_branches else final)
        integers = integer_matrix(candidate)
        int_residual = maxabs(candidate-mp.matrix(integers))
        rows.append({'loop':name,'matrix_rounded_after_check':integer_matrix(M),
                     'refinements':levels,'determinant_error':nstr(abs(mp.det(M)-1)),
                     'reported_matrix':integers,'reported_integer_residual':nstr(int_residual),
                     'reported_benchmark_passes':bool(maxabs(candidate-expected[name])<tolerance)})
    # Continue the already continued analytic germs, without resetting values.
    plus_minus,_,_,_ = continue_path(loops['minus_below_plus'],finals['plus'],60,mp.mpf('.1'))
    minus_plus,_,_,_ = continue_path(loops['plus'],finals['minus_below_plus'],60,mp.mpf('.1'))
    reverse,_,_,_ = continue_path(list(reversed(loops['plus'])),initial,60,mp.mpf('.1'))
    product_errors = {
        'plus_then_minus_equals_Mplus_Mminus':maxabs(recover(initial,plus_minus)-expected['plus']*expected['minus_below_plus']),
        'minus_then_plus_equals_Mminus_Mplus':maxabs(recover(initial,minus_plus)-expected['minus_below_plus']*expected['plus']),
        'plus_reverse_equals_inverse':maxabs(recover(initial,reverse)-expected['plus']**-1)}
    assert max(product_errors.values())<tolerance
    with mp.workdps(args.dps+20):
        initial_high = initial_periods(mp.mpf(args.base))
        high,_,_,_ = continue_path(paths(mp.mpf(args.base),mp.mpf(args.radius))['plus'],initial_high,60,mp.mpf('.1'))
    repeat = max(abs(x-y)/max(1,abs(x),abs(y)) for x,y in zip(high,finals['plus']))
    assert repeat<tolerance
    before = mp.matrix([initial[2],initial[0]])
    after = mp.matrix([finals['plus'][2],finals['plus'][0]])
    gamma = mp.matrix([[0,1]])
    passive = gamma*expected['plus']**-1
    Z0,Zpass,Zfixed = (gamma*before)[0],(passive*after)[0],(gamma*after)[0]
    tau0,tau1 = initial[3]/initial[1],finals['plus'][3]/finals['plus'][1]
    tau_exact = tau0/(1-2*tau0)
    metric0 = mp.im(tau0)*abs(initial[1])**2
    metric1 = mp.im(tau1)*abs(finals['plus'][1])**2
    assert abs(Zpass-Z0)<tolerance and abs(tau1-tau_exact)<tolerance and abs(metric1-metric0)<tolerance
    if args.figure:
        make_figure(args.figure,base,radius)
    return {
        'theory':'flat4d N=2 pure SU(2), first-paper period/charge basis',
        'base_z':nstr(base),'circle_radius':nstr(radius),'circle_segments':64,
        'digits':args.dps,'comparison_tolerance':nstr(tolerance),
        'initial_A':pair(initial[0]),'initial_D':pair(initial[2]),'loops':rows,
        'chronology_and_reversal_residuals':{k:nstr(v) for k,v in product_errors.items()},
        'plus_precision_repeat_scaled_error':nstr(repeat),
        'plus_charge_example':{'initial_row':[0,1],'passive_row':[2,1],
                               'initial_Z_over_Lambda':pair(Z0),'passive_Z_over_Lambda':pair(Zpass),
                               'continued_fixed_row_Z_over_Lambda':pair(Zfixed),
                               'initial_tau':pair(tau0),'continued_tau':pair(tau1),
                               'metric_residual':nstr(abs(metric1-metric0))},
        'invariant_primitive_rows':{'plus':[1,0],'minus_below_plus':[1,-1]},
        'intentional_branch_reset_detected':bool(args.reset_branches and all(not row['reported_benchmark_passes'] for row in rows)),
        'scope':'Finite continuation diagnostics, not interval bounds or a stable-particle inventory.'}


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--base',default='3')
    parser.add_argument('--radius',default='0.5')
    parser.add_argument('--dps',type=int,default=50)
    parser.add_argument('--reset-branches',action='store_true')
    parser.add_argument('--figure',help='Explicit output path for the loop SVG')
    args=parser.parse_args()
    try:
        result=run(args)
    except ValueError as error:
        parser.error(str(error))
    print(json.dumps(result,indent=2))


if __name__=='__main__':
    main()