# SCRIPT: RS03_PUSH
# The refraction law pushed several ways. Kill conditions: PREREGISTRATION_RS03.md (written before the run).
import json, numpy as np
from scipy.integrate import solve_ivp
from scipy.special import erf

def Yc(a, C):
    return np.sqrt(complex(a - C)) if a - C >= 0 else 1j*np.sqrt(C - a)

def gamma_pred(a1, a2, C):
    Y1, Y2 = Yc(a1, C), Yc(a2, C)
    return (Y1 - Y2)/(Y1 + Y2)

def solve(afun, bfun, a1, b1, a2, b2, C, om, L, breaks=()):
    # integrate (u, F), F = sqrt(a/b) u', from +L (outgoing or decaying) to -L; split at breakpoints
    p2 = om*om*C
    def rhs(y, s):
        u, F = s[0] + 1j*s[1], s[2] + 1j*s[3]
        a, b = afun(y), bfun(y)
        du = F*np.sqrt(b/a); dF = -np.sqrt(a*b)*(om*om - p2/a)*u
        return [du.real, du.imag, dF.real, dF.imag]
    q2 = b2*(om*om - p2/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
    pts = [L] + sorted([x for x in breaks if -L < x < L], reverse=True) + [-L]
    s = [u0.real, u0.imag, F0.real, F0.imag]
    for y0, y1 in zip(pts[:-1], pts[1:]):
        nstep = max(400, int(om*(y0 - y1)*10))
        sol = solve_ivp(rhs, [y0, y1], s, method='DOP853', rtol=1e-12, atol=1e-14, max_step=(y0 - y1)/nstep)
        s = sol.y[:, -1]
    u = s[0] + 1j*s[1]; F = s[2] + 1j*s[3]
    k1 = np.sqrt(b1*(om*om - p2/a1)); Y1 = np.sqrt(a1/b1)*k1
    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, 1/A          # reflection, transmission amplitude (at +L, unit outgoing)

def profile(kind, v1, v2, w):
    if kind == 'step': return lambda y: v1 if y < 0 else v2
    if kind == 'tanh': return lambda y: v1 + (v2 - v1)*0.5*(1 + np.tanh(y/w))
    if kind == 'erf':  return lambda y: v1 + (v2 - v1)*0.5*(1 + erf(y/w))
    if kind == 'ramp': return lambda y: v1 + (v2 - v1)*min(1.0, max(0.0, 0.5 + y/(2*w)))

def refl(kind, a1, b1, a2, b2, w, th, om=10.0, side=1):
    C = a1*np.sin(np.radians(th))**2
    if side == 2:   # wave sent from the right: mirror the layout, same C
        a1, b1, a2, b2 = a2, b2, a1, b1
    L = max(12*w, 3.0) if kind != 'step' else 3.0
    af, bf = profile(kind, a1, a2, w), profile(kind, b1, b2, w)
    br = (0.0,) if kind == 'step' else ((-w, w) if kind == 'ramp' else ())
    r, t = solve(af, bf, a1, b1, a2, b2, C, om, L, br)
    return r, t, C, a1, a2

rng = np.random.default_rng(2026)
res = {}

# W1: sharp step, random parameters
w1 = 0.0; n_ev = 0
for _ in range(400):
    a1, a2 = rng.uniform(0.3, 5, 2); b1, b2 = rng.uniform(0.3, 5, 2); th = rng.uniform(1, 89)
    r, t, C, *_ = refl('step', a1, b1, a2, b2, 0, th)
    w1 = max(w1, abs(r - gamma_pred(a1, a2, C))); n_ev += C > a2
res['W1'] = dict(max_err=w1, evanescent_cases=int(n_ev), passed=w1 < 1e-9)
print(flush=True); print(f"W1 sharp step, 400 random cases ({n_ev} beyond the turning point): max |r - Gamma| = {w1:.1e}  {'PASS' if w1 < 1e-9 else 'KILL'}")

# W2: three smooth profiles converge to the step
a1, a2, b1, b2 = 4.0, 1.0, 1.0, 0.5
W2 = {}; ok2 = True
for kind in ['tanh', 'erf', 'ramp']:
    rows = []
    for th in [10, 25, 29, 31, 60]:
        errs = []
        for ow in [0.08, 0.04, 0.02, 0.01]:
            r, t, C, *_ = refl(kind, a1, b1, a2, b2, ow/10, th)
            errs.append(abs(r - gamma_pred(a1, a2, C)))
        mono = all(e2 < e1 for e1, e2 in zip(errs, errs[1:]))
        q = np.polyfit(np.log([0.04, 0.02, 0.01]), np.log(errs[1:]), 1)[0]
        ok2 &= mono
        rows.append(dict(theta=th, errs=errs, order=q, monotone=mono))
    W2[kind] = rows
    print(f"W2 {kind:4s}: " + "  ".join(f"th{r_['theta']}: {r_['errs'][-1]:.1e} (order {r_['order']:.2f}{'' if r_['monotone'] else ', NOT monotone'})" for r_ in rows))
# Richardson: with order q fitted, the w -> 0 value of r from the last two widths
rich = 0.0
for kind in ['tanh', 'erf', 'ramp']:
    for th in [10, 25, 29, 31, 60]:
        rs = [refl(kind, a1, b1, a2, b2, ow/10, th)[0] for ow in [0.04, 0.02, 0.01]]
        C = a1*np.sin(np.radians(th))**2
        q = np.log2(abs(rs[0] - rs[1])/abs(rs[1] - rs[2]))
        r0 = rs[2] + (rs[2] - rs[1])/(2**q - 1)
        rich = max(rich, abs(r0 - gamma_pred(a1, a2, C)))
ok2 = ok2 and rich < 1e-5
res['W2'] = dict(rows=W2, richardson_max=rich, passed=bool(ok2))
print(f"W2 Richardson to w = 0, all profiles and angles: max |r0 - Gamma| = {rich:.1e}  {'PASS' if ok2 else 'KILL'}")

# W3: reciprocity
d_step = 0.0; d_smooth = 0.0
for th in [5, 15, 25, 29]:      # both sides propagating needs C < a2 = 1 -> th < 30
    rL = refl('step', a1, b1, a2, b2, 0, th)[0]; rR = refl('step', a1, b1, a2, b2, 0, th, side=2)[0]
    d_step = max(d_step, abs(rL + rR))
    for kind in ['tanh', 'erf', 'ramp']:
        for w in [0.003, 0.3]:
            rL = refl(kind, a1, b1, a2, b2, w, th)[0]; rR = refl(kind, a1, b1, a2, b2, w, th, side=2)[0]
            d_smooth = max(d_smooth, abs(abs(rL) - abs(rR)))
ok3 = d_step < 1e-8 and d_smooth < 1e-8
res['W3'] = dict(step=d_step, smooth=d_smooth, passed=ok3)
print(f"W3 reciprocity: |r21 + r12| step {d_step:.1e}; ||r21| - |r12|| smooth {d_smooth:.1e}  {'PASS' if ok3 else 'KILL'}")

# W4: flux conservation
f4 = 0.0
for kind in ['step', 'tanh', 'erf', 'ramp']:
    for w in [0.003, 0.3, 3.0]:
        for th in [5, 20, 29]:
            r, t, C, A1, A2 = refl(kind, a1, b1, a2, b2, w, th)
            f4 = max(f4, abs(abs(r)**2 + (Yc(A2, C).real/Yc(A1, C).real)*abs(t)**2 - 1))
res['W4'] = dict(max_err=f4, passed=f4 < 1e-8)
print(f"W4 flux |r|^2 + (Y2/Y1)|t|^2 = 1: max error {f4:.1e}  {'PASS' if f4 < 1e-8 else 'KILL'}")

# W5: opposite - g_perp matters in a thick layer
d5 = 0.0
for th in np.arange(2, 30, 3.0):
    rA = refl('tanh', a1, b1, a2, 0.5, 0.3, th)[0]; rB = refl('tanh', a1, b1, a2, 2.0, 0.3, th)[0]
    d5 = max(d5, abs(abs(rA) - abs(rB)))
res['W5'] = dict(max_change=d5, passed=d5 > 1e-3)
print(f"W5 opposite: thick layer (omega w = 3), b2 x4 changes |r| by up to {d5:.3f}  {'PASS (g_perp matters off the step)' if d5 > 1e-3 else 'KILL'}")

# W6: total reflection beyond the turning point, and the adiabatic onset
thc = np.degrees(np.arcsin(np.sqrt(a2/a1)))
d6 = 0.0
for kind in ['step', 'tanh', 'erf', 'ramp']:
    for w in [0.003, 0.3, 3.0]:
        for th in [31, 45, 70, 88]:
            d6 = max(d6, abs(abs(refl(kind, a1, b1, a2, b2, w, th)[0]) - 1))
halves = []
from scipy.optimize import brentq
for ow in [10, 30, 100]:
    f = lambda th: abs(refl('tanh', a1, b1, a2, b2, ow/10, th)[0])**2 - 0.5
    halves.append(brentq(f, max(1.0, thc - 15), thc - 1e-6, xtol=1e-4))
    print(f"   omega w = {ow}: half-reflection at {halves[-1]:.4f} deg", flush=True)
gaps = [abs(h - thc) for h in halves]
ok6 = d6 < 1e-9 and gaps[0] > gaps[1] > gaps[2]
res['W6'] = dict(max_dev=d6, theta_c=thc, halves=halves, passed=ok6)
print(f"W6 |r| = 1 beyond theta_c = {thc:.4f}: max dev {d6:.1e}; half-reflection at omega w = 10, 30, 100: " + ", ".join(f"{h:.3f}" for h in halves) + f"  {'PASS' if ok6 else 'KILL'}")

# W7: gap family at normal incidence
w7 = 0.0
for _ in range(200):
    D1, D2 = rng.uniform(0.2, 5, 2); bb1, bb2 = rng.uniform(0.3, 3, 2)
    r = refl('step', 1/D1**2, bb1, 1/D2**2, bb2, 0, 0.0)[0]
    w7 = max(w7, abs(r - np.tanh(0.5*np.log(D2/D1))))
res['W7'] = dict(max_err=w7, passed=w7 < 1e-9)
print(f"W7 gap family Gamma = tanh(ln(D2/D1)/2), 200 pairs: max error {w7:.1e}  {'PASS' if w7 < 1e-9 else 'KILL'}")

json.dump(res, open('rs03_results.json', 'w'), indent=1, default=float)