# 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


def refl_resolved(kind, a1, b1, a2, b2, w, th, om=10.0):
    C = a1*np.sin(np.radians(th))**2
    Lh = max(12*w, 3.0)
    af, bf = profile(kind, a1, a2, w), profile(kind, b1, b2, w)
    br = tuple(sorted({0.0, w, -w, 3*w, -3*w, 8*w, -8*w}))
    r, t = solve(af, bf, a1, b1, a2, b2, C, om, Lh, br)
    return r, C