# SCRIPT: RS01_POLAR_MIRROR
# Wave version of IQGT v7.3.1 section 8.4: is the polar stratum a mirror for waves,
# and do the Smith-chart phase and lateral shift of the reflected wave reproduce the
# Clairaut geodesic? Kill conditions are in PREREGISTRATION.md (written before this run).
import json, numpy as np
from scipy.integrate import solve_ivp, quad
from scipy.optimize import brentq

K = 3.0
YL = 1.2
def a(y): return 0.25*K*K*np.sin(y)**2
B = 0.25
aL = a(YL)

def Q(y, w, p): return 0.25*(w*w - p*p/a(y))          # b (w^2 - p^2/a)
def rhs(y, s, w, p):
    u, du = s
    return [du, -du/np.tan(y) - Q(y, w, p)*u]

def start_point(w, p, yt, depth):
    f = lambda y0: quad(lambda y: np.sqrt(max(-Q(y, w, p), 0.0)), y0, yt, limit=200)[0] - depth
    lo = 1e-6
    if f(lo) < 0: return lo
    return brentq(f, lo, yt*(1-1e-9))

def reflection(w, p, depth=30.0):
    yt = np.arcsin(np.sqrt(4*p*p/(w*w)/(K*K)))
    y0 = start_point(w, p, yt, depth)
    kap = np.sqrt(-Q(y0, w, p))
    sol = solve_ivp(rhs, [y0, YL], [1.0, kap], args=(w, p), method='DOP853',
                    rtol=1e-12, atol=1e-30, max_step=0.02)
    u, du = sol.y[0, -1], sol.y[1, -1]
    kL = np.sqrt(0.25*(w*w - p*p/aL))
    Bc = 0.5*(u + du/(1j*kL)); Ac = 0.5*(u - du/(1j*kL))
    return Bc/Ac, yt

def geo_dx(C):
    yt = np.arcsin(np.sqrt(4*C/(K*K)))
    # substitution y = yt + s^2 removes the 1/sqrt endpoint singularity
    g = lambda s: 2*s*np.sqrt(B*C)/(np.sqrt(a(yt+s*s))*np.sqrt(a(yt+s*s)-C))
    return 2*quad(g, 0, np.sqrt(YL-yt), limit=400, epsabs=1e-13, epsrel=1e-12)[0]

def wkb_phase(w, C):
    p = w*np.sqrt(C); yt = np.arcsin(np.sqrt(4*C/(K*K)))
    g = lambda s: 2*s*np.sqrt(max(Q(yt+s*s, w, p), 0.0))
    S = quad(g, 0, np.sqrt(YL-yt), limit=400, epsabs=1e-12, epsrel=1e-12)[0]
    return 2*S - np.pi/2

out = {"aL": aL, "rows": []}
omegas = [20, 40, 80, 160, 320]
thetas = [20, 35, 50, 65]
maxdev = 0.0; ctrl = 0.0
for th in thetas:
    C = aL*np.sin(np.radians(th))**2
    dxg = geo_dx(C)
    for w in omegas:
        p = w*np.sqrt(C)
        r, yt = reflection(w, p)
        r2, _ = reflection(w, p, depth=40.0)
        ctrl = max(ctrl, abs(r - r2))
        maxdev = max(maxdev, abs(abs(r) - 1))
        dp = p*1e-5
        rp, _ = reflection(w, p+dp); rm, _ = reflection(w, p-dp)
        D = -np.angle(rp/rm)/(2*dp)
        dphi = np.angle(r*np.exp(-1j*wkb_phase(w, C)))
        out["rows"].append(dict(theta=th, omega=w, C=C, y_turn=yt, abs_r_minus_1=abs(r)-1,
                                D_wave=D, dx_geo=dxg, rel_err=(D-dxg)/dxg, phase_minus_wkb=dphi))
print(f"a_L = {aL:.6f}; max | |r| - 1 | = {maxdev:.2e}; control (start depth 30 vs 40) max |dr| = {ctrl:.2e}")
print(" theta omega   y_turn    D_wave     dx_geo    rel_err    phase-WKB")
for R in out["rows"]:
    print(f"{R['theta']:5d} {R['omega']:5d} {R['y_turn']:8.5f} {R['D_wave']:10.6f} {R['dx_geo']:10.6f} {R['rel_err']:+.3e} {R['phase_minus_wkb']:+.3e}")
slopes = {}
for th in thetas:
    rows = [R for R in out["rows"] if R["theta"] == th]
    x = np.log([R["omega"] for R in rows]); y = np.log([abs(R["rel_err"]) for R in rows])
    slopes[th] = np.polyfit(x, y, 1)[0]
print("log-log slope of |rel_err| vs omega:", {k: round(v, 3) for k, v in slopes.items()})
last = [R for R in out["rows"] if R["omega"] == 320]
K1 = maxdev < 1e-9
K2 = all(s <= -0.8 for s in slopes.values()) and all(abs(R["rel_err"]) < 0.02 for R in last)
K3 = all(abs(R["phase_minus_wkb"]) < 0.02 for R in last)
CT = ctrl < 1e-8
print(f"K1 |r|=1: {'PASS' if K1 else 'KILL'}   K2 shift: {'PASS' if K2 else 'KILL'}   K3 phase: {'PASS' if K3 else 'KILL'}   control: {'PASS' if CT else 'FAIL'}")
out.update(slopes={str(k): v for k, v in slopes.items()}, K1=K1, K2=K2, K3=K3, control=CT, control_max=ctrl, max_abs_r_dev=maxdev)
json.dump(out, open("rs01_results.json", "w"), indent=1, default=float)