"""
EXPERIMENT 1c: Robust root finding for F(z) = sum z^n/n^n
Using continuation method from known roots.
"""
import numpy as np
from scipy.optimize import fsolve
import json

def F_vec(zv, N=150):
    z = complex(zv[0], zv[1])
    result = 0.0 + 0.0j
    for n in range(1, N+1):
        try:
            term = z**n / n**n
        except (OverflowError, ZeroDivisionError):
            break
        if abs(term) > 1e100 or np.isnan(abs(term)) or np.isinf(abs(term)):
            break
        result += term
    return [result.real, result.imag]

# Start from known roots and extend
known_roots = [
    (-4.305, 14.945),
    (-5.386, 32.038),
    (-5.996, 49.129),
    (-6.421, 66.215),
]

roots = []
for kr in known_roots:
    r = complex(kr[0], kr[1])
    val = complex(*F_vec([r.real, r.imag]))
    if abs(val) < 1e-3:
        # Refine with fsolve
        refined = fsolve(F_vec, [r.real, r.imag], full_output=False)
        r = complex(refined[0], refined[1])
        roots.append(r)

# Continue from last known root
if len(roots) > 0:
    last_root = roots[-1]
    # Estimate next root position
    if len(roots) >= 2:
        im_step = roots[-1].imag - roots[-2].imag
        re_step = roots[-1].real - roots[-2].real
    else:
        im_step = 17.08
        re_step = -0.4
    
    # Extend forward
    for k in range(30):
        re_guess = last_root.real + re_step
        im_guess = last_root.imag + im_step
        try:
            refined = fsolve(F_vec, [re_guess, im_guess], full_output=False)
            r = complex(refined[0], refined[1])
            val = complex(*F_vec([r.real, r.imag]))
            if abs(val) < 1e-3 and r.imag > 0 and r.real < 0:
                # Check not duplicate
                is_dup = any(abs(r - ex) < 1.0 for ex in roots)
                if not is_dup:
                    roots.append(r)
                    last_root = r
                    # Update step estimates from last two roots
                    if len(roots) >= 2:
                        im_step = roots[-1].imag - roots[-2].imag
                        re_step = roots[-1].real - roots[-2].real
        except:
            pass

# Also search backward from first root
if len(roots) > 0:
    first_root = roots[0]
    if len(roots) >= 2:
        im_step = roots[1].imag - roots[0].imag
        re_step = roots[1].real - roots[0].real
    else:
        im_step = 17.08
        re_step = 0.4  # going backward means Re increases
    
    for k in range(5):
        re_guess = first_root.real - re_step  # going back
        im_guess = first_root.imag - im_step
        if im_guess < 5:
            break
        try:
            refined = fsolve(F_vec, [re_guess, im_guess], full_output=False)
            r = complex(refined[0], refined[1])
            val = complex(*F_vec([r.real, r.imag]))
            if abs(val) < 1e-3 and r.imag > 0 and r.real < 0:
                is_dup = any(abs(r - ex) < 1.0 for ex in roots)
                if not is_dup:
                    roots.insert(0, r)
                    first_root = r
        except:
            pass

roots.sort(key=lambda r: r.imag)
print(f"Found {len(roots)} roots\n")

if len(roots) >= 2:
    im_spacings = [roots[i].imag - roots[i-1].imag for i in range(1, len(roots))]
    re_shifts = [roots[i].real - roots[i-1].real for i in range(1, len(roots))]
    
    print(f"{'#':>3s}  {'Re':>10s}  {'Im':>10s}  {'Im-gap':>10s}  {'Re-shift':>10s}")
    for i, r in enumerate(roots):
        im_gap = f"{im_spacings[i-1]:10.4f}" if i > 0 else f"{'---':>10s}"
        re_shift = f"{re_shifts[i-1]:10.4f}" if i > 0 else f"{'---':>10s}"
        print(f"{i:3d}  {r.real:10.4f}  {r.imag:10.4f}  {im_gap}  {re_shift}")
    
    mean_im = np.mean(im_spacings)
    std_im = np.std(im_spacings)
    
    print(f"\nIm-spacing: mean={mean_im:.6f}, std={std_im:.6f}")
    print(f"2*pi*e:     {2*np.pi*np.e:.6f}")
    print(f"Ratio:      {mean_im/(2*np.pi*np.e):.6f}")
    print(f"Re-shift:   mean={np.mean(re_shifts):.6f}")
    
    # Save
    results = {
        'roots_re': [r.real for r in roots],
        'roots_im': [r.imag for r in roots],
        'im_spacings': im_spacings,
        're_shifts': re_shifts,
        'two_pi_e': float(2*np.pi*np.e),
        'mean_im_spacing': float(mean_im),
        'std_im_spacing': float(std_im),
    }
    with open('/app/exp1_results.json', 'w') as f:
        json.dump(results, f, indent=2)
    print("\nSaved to /app/exp1_results.json")