# SCRIPT: RS02_INTERFACE
# Reflection (Smith coordinate) versus refraction (Clairaut) at a step in the pullback metric
# g = a(y) dx^2 + b(y) dy^2, x a Killing direction. Kill conditions: PREREGISTRATION.md.
import json, numpy as np
from scipy.integrate import solve_ivp

def gamma_pred(a1, a2, C):
    Y = lambda a: np.sqrt(complex(a - C)) if a - C >= 0 else 1j*np.sqrt(C - a)
    Y1, Y2 = Y(a1), Y(a2)
    return (Y1 - Y2)/(Y1 + Y2)

def refl(a1, b1, a2, b2, w, th_deg, om=10.0, Lfac=12.0):
    C = a1*np.sin(np.radians(th_deg))**2           # Clairaut constant, theta measured in g at the left lead
    p = om*np.sqrt(C)
    L = max(Lfac*w, 3.0)
    prof = lambda y, v1, v2: v1 + (v2 - v1)*0.5*(1 + np.tanh(y/w))
    a = lambda y: prof(y, a1, a2); b = lambda y: prof(y, b1, b2)
    def rhs(y, s):
        u, F = s[0] + 1j*s[1], s[2] + 1j*s[3]
        du = F*np.sqrt(b(y)/a(y)); dF = -np.sqrt(a(y)*b(y))*(om*om - p*p/a(y))*u
        return [du.real, du.imag, dF.real, dF.imag]
    q2 = b2*(om*om - p*p/a2)
    if q2 > 0:
        k2 = np.sqrt(q2); u0 = np.exp(1j*k2*L); F0 = np.sqrt(a2/b2)*1j*k2*u0
    else:
        kap = np.sqrt(-q2); u0 = 1.0 + 0j; F0 = -np.sqrt(a2/b2)*kap*u0   # decaying, normalised at +L
    span = 2*L; nstep = max(2000, int(om*span*8))
    sol = solve_ivp(rhs, [L, -L], [u0.real, u0.imag, F0.real, F0.imag], method='DOP853',
                    rtol=1e-11, atol=1e-14, max_step=span/nstep)
    u = sol.y[0, -1] + 1j*sol.y[1, -1]; F = sol.y[2, -1] + 1j*sol.y[3, -1]
    k1 = np.sqrt(b1*(om*om - p*p/a1)); Y1 = np.sqrt(a1/b1)*k1
    # u = A e^{ik1 y} + R e^{-ik1 y};  F = i Y1 (A e^{ik1 y} - R e^{-ik1 y})  at y = -L
    Ae = 0.5*(u + F/(1j*Y1)); Re = 0.5*(u - F/(1j*Y1))
    A = Ae*np.exp(1j*k1*L); R = Re*np.exp(-1j*k1*L)
    return R/A, C

out = {}
a1, a2, b1 = 4.0, 1.0, 1.0
thc = np.degrees(np.arcsin(np.sqrt(a2/a1)))
print(f"critical angle from Clairaut: {thc:.4f} deg")

# K1: sharp limit and linear convergence
print("\nK1  sharp limit (omega=10), |r - Gamma| at omega*w = 0.04, 0.02, 0.01")
k1rows = []; worst = 0.0
for th in [10, 20, 25, 29, 31, 40, 60, 80]:
    errs = []
    for ow in [0.04, 0.02, 0.01]:
        r, C = refl(a1, b1, a2, 0.5, ow/10.0, th)
        errs.append(abs(r - gamma_pred(a1, a2, C)))
    worst = max(worst, errs[-1])
    sl = np.polyfit(np.log([0.04, 0.02, 0.01]), np.log(errs), 1)[0]
    k1rows.append(dict(theta=th, errs=errs, slope=sl))
    print(f"  theta {th:2d}: {errs[0]:.2e} {errs[1]:.2e} {errs[2]:.2e}  slope {sl:+.2f}")
K1 = worst < 1e-3 and all(r_['slope'] > 0.7 for r_ in k1rows)
print("K1:", "PASS" if K1 else "KILL")

# K2: g_perp drops out
print("\nK2  change b2 by a factor 4 at fixed a (sharp, omega*w = 0.01)")
dmax = 0.0
for th in [10, 20, 29, 31, 50, 80]:
    rA, _ = refl(a1, b1, a2, 0.5, 0.001, th); rB, _ = refl(a1, b1, a2, 2.0, 0.001, th)
    rC, _ = refl(a1, 3.0, a2, 0.5, 0.001, th)
    d = max(abs(abs(rA) - abs(rB)), abs(abs(rA) - abs(rC))); dmax = max(dmax, d)
    print(f"  theta {th:2d}: |r| = {abs(rA):.6f}, {abs(rB):.6f} (b2 x4), {abs(rC):.6f} (b1 x3)   max diff {d:.1e}")
K2 = dmax < 1e-3
print("K2:", "PASS" if K2 else "KILL")

# K3: thick layer, onset of total reflection at the Clairaut critical angle
print("\nK3  thick layer omega*w = 30")
ths = np.arange(1.0, 89.01, 0.25); R2 = []
for th in ths:
    r, _ = refl(a1, b1, a2, 0.5, 3.0, th); R2.append(abs(r)**2)
R2 = np.array(R2)
i = np.where(R2 >= 0.5)[0][0]
th_half = ths[i-1] + (0.5 - R2[i-1])*(ths[i] - ths[i-1])/(R2[i] - R2[i-1])
above = ths > thc + 1e-9
dev_above = np.max(np.abs(np.sqrt(R2[above]) - 1))
below_max = np.max(R2[ths < thc - 1.0])
print(f"  |r|^2 = 1/2 at {th_half:.3f} deg (critical {thc:.3f}); max ||r|-1| above critical {dev_above:.1e}; max |r|^2 more than 1 deg below critical {below_max:.2e}")
K3 = abs(th_half - thc) < 1.5 and dev_above < 1e-9
print("K3:", "PASS" if K3 else "KILL")

# Opposite: index rising across the step
print("\nOpposite  a1=1 -> a2=4 (index rises), sharp and thick")
Omax = 0.0
for w in [0.001, 3.0]:
    for th in np.arange(1.0, 89.0, 2.0):
        r, _ = refl(1.0, 1.0, 4.0, 0.5, w, th); Omax = max(Omax, abs(r))
print(f"  largest |r| for theta < 89 deg: {Omax:.6f}")
OPP = Omax < 1 - 1e-6
print("Opposite:", "PASS" if OPP else "KILL")

out.update(theta_c=thc, K1=K1, K1rows=k1rows, K2=K2, K2_maxdiff=dmax, K3=K3, theta_half=th_half,
           dev_above=dev_above, below_max=below_max, opposite=OPP, opposite_max=Omax,
           thick_curve=dict(theta=ths.tolist(), R2=R2.tolist()))
json.dump(out, open('rs02_results.json', 'w'), indent=1, default=float)