# SCRIPT: BH-FIGS-V10
# Figures f19 (reflection and refraction on the pullback metric; the Smith-chart map; the sine guard)
# and f20 (the photon ring read by the wobble generator), computed from the formulas checked in
# RS-03, V10_PUSH_OTHERS and PHOTON_RING_WALL.
import os, numpy as np, matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp
_H = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(_H, '..', 'paper', 'fig') if os.path.isdir(os.path.join(_H, '..', 'paper', 'fig')) else os.path.join(_H, '..', 'v10', 'fig')
BG='#0f1115'; INK='#e6e8ec'; MUT='#9aa3ae'; RULE='#2a313c'
C1='#3987e5'; C2='#d95926'; C3='#199e70'; C4='#c98500'
plt.rcParams.update({'figure.facecolor':BG,'axes.facecolor':BG,'savefig.facecolor':BG,'text.color':INK,
  'axes.labelcolor':INK,'xtick.color':MUT,'ytick.color':MUT,'axes.edgecolor':RULE,'font.family':'serif',
  'font.serif':['TeX Gyre Pagella','URW Palladio L','Palatino','DejaVu Serif'],'mathtext.fontset':'cm','font.size':11})
def clean(ax):
    for s in ['top','right']: ax.spines[s].set_visible(False)
def Yc(a, C): return np.sqrt(complex(a - C)) if a >= C else 1j*np.sqrt(C - a)
def gam(a1, a2, C): Y1, Y2 = Yc(a1, C), Yc(a2, C); return (Y1 - Y2)/(Y1 + Y2)

# thick-layer reflection by direct integration (same solver as RS-03)
def refl_tanh(a1, b1, a2, b2, w, th, om=10.0):
    C = a1*np.sin(np.radians(th))**2; p2 = om*om*C; L = max(12*w, 3.0)
    a = lambda y: a1 + (a2 - a1)*0.5*(1 + np.tanh(y/w)); b = lambda y: b1 + (b2 - b1)*0.5*(1 + np.tanh(y/w))
    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 - p2/a(y))*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: u0 = 1.0 + 0j; F0 = -np.sqrt(a2/b2)*np.sqrt(-q2)*u0
    sol = solve_ivp(rhs, [L, -L], [u0.real, u0.imag, F0.real, F0.imag], method='DOP853', rtol=1e-10, atol=1e-13, max_step=2*L/3000)
    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 - p2/a1)); Y1 = np.sqrt(a1/b1)*k1
    Ae = 0.5*(u + F/(1j*Y1)); Re = 0.5*(u - F/(1j*Y1))
    return (Re*np.exp(-1j*k1*L))/(Ae*np.exp(1j*k1*L))

# ---------------- f19 ----------------
a1, a2 = 4.0, 1.0; thc = np.degrees(np.arcsin(np.sqrt(a2/a1)))
th = np.linspace(0, 89.5, 600)
R_step = np.array([abs(gam(a1, a2, a1*np.sin(np.radians(t))**2))**2 for t in th])
R_up = np.array([abs(gam(1.0, 4.0, np.sin(np.radians(t))**2))**2 for t in th])
thk = np.arange(1, 89, 2.0)
R_thickA = np.array([abs(refl_tanh(a1, 1.0, a2, 0.5, 0.06, t))**2 for t in thk])
R_thickB = np.array([abs(refl_tanh(a1, 1.0, a2, 2.0, 0.06, t))**2 for t in thk])
print('f19: thick layer max |R_A - R_B| =', np.abs(R_thickA - R_thickB).max())

fig = plt.figure(figsize=(12.6, 4.5))
ax = fig.add_axes([0.05, 0.14, 0.29, 0.72]); bx = fig.add_axes([0.385, 0.08, 0.25, 0.84]); cx = fig.add_axes([0.70, 0.14, 0.285, 0.72])
for q in (ax, cx): clean(q)
ax.plot(th, R_step, color=C1, lw=2.4, label='sharp step')
ax.plot(th, R_up, color=C3, lw=2.0, label=r'opposite, $g_\parallel$ rises')
ax.plot(thk, R_thickA, 'o-', color=C2, ms=3.4, lw=1, label=r'layer $\omega w$=0.6, $g_\perp$=0.5')
ax.plot(thk, R_thickB, 's-', color=C4, ms=3.2, lw=1, mfc='none', label=r'layer $\omega w$=0.6, $g_\perp$=2')
ax.axvline(thc, color=RULE, lw=1.2, ls=(0, (4, 3)))
ax.text(thc + 1.5, 0.30, 'Clairaut turning\npoint: total\nreflection', color=MUT, fontsize=9)
ax.set_xlabel(r'angle of incidence $\theta$ (deg), $C=g_\parallel\sin^2\theta$'); ax.set_ylabel(r'reflected power $|\Gamma|^2$')
ax.set_ylim(-0.03, 1.05); ax.legend(loc='center right', frameon=True, facecolor=BG, edgecolor=BG, framealpha=1, fontsize=8.6, labelcolor=INK, bbox_to_anchor=(1.0, 0.72))
ax.set_title(r'a  $g_\parallel$: 4 $\to$ 1, $g_\perp$ drops out at the step', color=INK, fontsize=10.5, loc='left')

# (b) the Smith chart map
bx.set_aspect('equal'); bx.axis('off')
tt = np.linspace(0, 2*np.pi, 400); bx.plot(np.cos(tt), np.sin(tt), color=MUT, lw=1.2)
bx.plot([-1, 1], [0, 0], color=RULE, lw=0.9)
for x in [0.5, 0.0, -0.5]:   # constant-resistance circles of the chart, for orientation
    r = (1 - x)/2; c = (1 + x)/2; bx.plot(c + r*np.cos(tt), r*np.sin(tt), color=RULE, lw=0.6)
G = np.array([gam(a1, a2, a1*np.sin(np.radians(t))**2) for t in np.linspace(0, 89.9, 800)])
bx.plot(G.real, G.imag, color=C1, lw=2.6)
bx.annotate('', xy=(G[-1].real, G[-1].imag), xytext=(G[-40].real, G[-40].imag), arrowprops=dict(arrowstyle='->', color=C1, lw=1.6))
bx.plot(G[0].real, G[0].imag, 'o', color=C1, ms=5)
bx.text(G[0].real - 0.05, G[0].imag + 0.09, 'normal\nincidence', color=C1, fontsize=8.6, ha='center')
bx.text(0.5, -0.62, 'past the turning\npoint: on the rim', color=C1, fontsize=8.6, ha='center')
bx.plot(0, 0, 'o', color=C2, ms=7); bx.text(0.02, -0.16, 'matched horizon\n$\\Gamma=0$', color=C2, fontsize=8.6, ha='center', va='top')
wt = np.array([0.5, 1, 2, 4]); gf = np.tanh(wt/2)
bx.plot(gf, 0*gf, 'D', color=C3, ms=4.5)
bx.text(-0.98, 0.35, r'fermion record' '\n' r'$\Gamma=\tanh(\omega/2T_H)$' '\n' r'($\omega/T_H$ = 0.5, 1, 2, 4)', color=C3, fontsize=8.4)
bx.annotate('', xy=(gf[0], 0.03), xytext=(-0.45, 0.33), arrowprops=dict(arrowstyle='->', color=C3, lw=1))
bx.set_xlim(-1.08, 1.08); bx.set_ylim(-1.08, 1.08)
bx.set_title(r'b  one chart: step, horizon, fermion record', color=INK, fontsize=10.5, loc='left')

# (c) the sine guard
rng = np.random.default_rng(5)
def proj_family(n):
    Hs = [rng.normal(size=(n, n)) + 1j*rng.normal(size=(n, n)) for _ in range(3)]; Hs = [(h + h.conj().T)/2 for h in Hs]
    def P(x, y):
        H = Hs[0] + np.cos(2*x)*Hs[1] + np.sin(2*x)*y*Hs[2] + y**2*Hs[2]
        w, v = np.linalg.eigh(H); return np.outer(v[:, 0], v[:, 0].conj())
    return P
pts = {2: [], 3: []}
for n in (2, 3):
    for _ in range(160):
        P = proj_family(n); x, y = rng.uniform(-1, 1, 2); h = 1e-5
        dx = (P(x + h, y) - P(x - h, y))/(2*h); dy = (P(x, y + h) - P(x, y - h))/(2*h); Pi = P(x, y)
        g = lambda A, B: 0.5*np.trace(A@B).real; Om = lambda A, B: (1j*np.trace(Pi@(A@B - B@A))).real
        gxx = g(dx, dx); e1 = dx/np.sqrt(gxx); t = dy - g(dy, e1)*e1; e2 = t/np.sqrt(g(t, t))
        thv = rng.uniform(0, 1.55); V = np.sin(thv)*e1 + np.cos(thv)*e2
        pts[n].append((np.degrees(thv), abs(Om(V, dx))/(2*np.sqrt(gxx))))
tg = np.linspace(0, 90, 300)
cx.fill_between(tg, np.cos(np.radians(tg)), 1.05, color=C2, alpha=0.10, lw=0)
cx.plot(tg, np.cos(np.radians(tg)), color=INK, lw=1.8)
p3 = np.array(pts[3]); p2 = np.array(pts[2])
cx.plot(p3[:, 0], p3[:, 1], 'o', color=C1, ms=3.0, alpha=0.8, label='three-level families')
cx.plot(p2[:, 0], p2[:, 1], 'D', color=C4, ms=3.0, label='two-level families (equality)')
cx.text(50, 0.85, 'forbidden', color=C2, fontsize=9.5)
cx.text(4, 0.30, r'$|\Omega(V,\partial_x)|\leq 2\sqrt{g_\parallel-C}$' '\n' r'$=2\sqrt{g_\parallel}\cos\theta=2Y/\omega$', color=INK, fontsize=9.5, bbox=dict(facecolor=BG, edgecolor=RULE, pad=4))
cx.set_xlabel(r'path angle $\theta$ from the normal (deg)'); cx.set_ylabel(r'$|\Omega(V,\partial_x)|\,/\,2\sqrt{g_\parallel}$')
cx.set_ylim(-0.02, 1.05); cx.legend(loc='lower left', frameon=False, fontsize=8.6, labelcolor=INK)
cx.set_title(r'c  the sine guard: one Gram determinant', color=INK, fontsize=10.5, loc='left')
fig.savefig(os.path.join(OUT, 'f19_refraction.png'), dpi=200); plt.close(fig)
print('max ratio three-level =', p3[:, 1].max(), ' two-level dev from cos =', np.abs(p2[:, 1] - np.cos(np.radians(p2[:, 0]))).max())

# ---------------- f20 ----------------
M = 1.0
def rate(Q):
    u0 = 2/(3*M + np.sqrt(9*M*M - 8*Q*Q)); return np.sqrt(6*M*u0 - 6*Q*Q*u0*u0 - 1), u0
def measured(Q):
    lam, u0 = rate(Q)
    f = lambda ph, s: [s[1], -s[0] + 3*M*s[0]**2 - 2*Q*Q*s[0]**3]
    sol = solve_ivp(f, [0, 12], [u0 + 1e-9, 0.0], method='DOP853', rtol=1e-13, atol=1e-16, dense_output=True)
    ph = np.linspace(4, 11, 50); return np.polyfit(ph, np.log(np.abs(sol.sol(ph)[0] - u0)), 1)[0]
Qs = np.linspace(0, 1.0, 200); lam = np.array([rate(q)[0] for q in Qs])
Qm = np.array([0, 0.3, 0.5, 0.7, 0.9]); lm = np.array([measured(q) for q in Qm])
print('f20: measured growth per radian', dict(zip(Qm, np.round(lm, 6))))
fig = plt.figure(figsize=(11.6, 4.4))
ax = fig.add_axes([0.065, 0.15, 0.40, 0.72]); bx = fig.add_axes([0.575, 0.15, 0.40, 0.72])
for q in (ax, bx): clean(q)
ax.plot(Qs, np.exp(-np.pi*lam), color=C1, lw=2.4, label=r'$e^{-\pi\lambda/\Omega}$ from the generator')
ax.plot(Qm, np.exp(-np.pi*lm), 'o', color=C2, ms=6, label='nonlinear null orbits, fitted')
ax.axhline(np.exp(-np.pi), color=RULE, lw=1, ls=(0, (4, 3)))
ax.text(0.02, np.exp(-np.pi) + 0.003, r'Schwarzschild: $e^{-\pi}=0.0432$', color=MUT, fontsize=9.5)
ax.set_xlabel(r'charge $Q/M$ (Reissner–Nordström)'); ax.set_ylabel('brightness ratio of successive images')
ax.legend(loc='upper left', frameon=False, fontsize=9, labelcolor=INK)
ax.set_title(r'a  dimming per half orbit is $e^{-\pi}$ only when $Q=0$', color=INK, fontsize=10.5, loc='left')
n = np.arange(0, 5); tau = 3*np.sqrt(3)*np.pi
bx.bar(n*tau, np.exp(-np.pi*n), width=3.0, color=C1)
for k in range(5): bx.text(k*tau, np.exp(-np.pi*k)*1.6, f'n = {k}', color=MUT, fontsize=9, ha='center')
bx.set_yscale('log'); bx.set_ylim(2e-6, 4)
bx.set_xlabel(r'arrival time after the direct image ($M$)'); bx.set_ylabel(r'relative brightness $e^{-n\pi}$')
bx.set_xticks(n*tau); bx.set_xticklabels(['0'] + [f'{k}' + r'$\cdot 3\sqrt{3}\pi$' for k in range(1, 5)])
bx.set_title(r'b  Schwarzschild: each half orbit later by $3\sqrt{3}\pi M$', color=INK, fontsize=10.5, loc='left')
fig.savefig(os.path.join(OUT, 'f20_photon_ring.png'), dpi=200); plt.close(fig)