r"""
Frame dragging and the ergosphere of a Kerr black hole
==============================================================

A rotating (Kerr) black hole of mass :math:`M` and spin parameter
:math:`a = J/M \in [0, M)`, found by Roy Kerr in 1963, drags spacetime
itself around with it. Outside the horizon :math:`r_+ = M +
\sqrt{M^2-a^2}`, in the ergosphere

.. math::

    r_E(\theta) = M + \sqrt{M^2 - a^2\cos^2\theta}

this dragging is so severe that no observer can stay at fixed spatial
coordinates -- everyone is swept around in the direction of rotation, no
matter how they fire their rockets. This example traces out the ergosphere
and horizon as spin increases, the Lense-Thirring frame-dragging angular
velocity of a zero-angular-momentum observer, and how the innermost stable
circular orbit (ISCO) shrinks (prograde) or grows (retrograde) with spin.
The ergosphere is also what makes the Penrose energy-extraction process
possible -- see :doc:`plot_penrose_energy_extraction`.
"""

import matplotlib.pyplot as plt
import numpy as np

from physicskit.relativity.chapters.kerr import KerrBlackHole

# %%
# The ergosphere and horizons, as spin increases
# ------------------------------------------------------
# The event horizon :math:`r_+ = M + \sqrt{M^2-a^2}` shrinks as spin
# increases, while the ergosphere :math:`r_E(\theta)` (widest at the
# equator, where :math:`r_E = 2M`, and touching the horizon at the poles)
# bulges out further beyond it -- the region between the two curves is
# where no observer can remain at fixed :math:`(r,\theta,\phi)`.
theta = np.linspace(0, np.pi, 200)
fig, axes = plt.subplots(1, 3, figsize=(13, 4.5), subplot_kw={"projection": "polar"})
for ax, a in zip(axes, [0.5, 0.9, 0.998]):
    bh = KerrBlackHole(M=1.0, a=a)
    r_ergo = bh.ergosphere_radius(theta)
    ax.plot(theta, r_ergo, color="crimson", label="ergosphere $r_E(\\theta)$")
    ax.plot(theta, np.full_like(theta, bh.outer_horizon_radius), color="black", label="horizon $r_+$")
    ax.plot(-theta, r_ergo, color="crimson")
    ax.plot(-theta, np.full_like(theta, bh.outer_horizon_radius), color="black")
    ax.set_title(f"a/M = {a}")
    ax.set_ylim(0, 2.2)
axes[0].legend(loc="upper right", fontsize=7, bbox_to_anchor=(1.3, 1.2))
plt.tight_layout()

# %%
# Frame dragging: the ZAMO angular velocity rises sharply near the horizon
# --------------------------------------------------------------------------
# A zero-angular-momentum observer (ZAMO) -- one who carries no orbital
# angular momentum of their own, yet is still swept around in
# :math:`\phi` purely by the geometry -- rotates at
#
# .. math::
#
#     \omega(r) = -\frac{g_{t\phi}}{g_{\phi\phi}}
#               = \frac{2Ma}{r^3 + a^2 r + 2Ma^2}
#
# which falls off as the classic Lense-Thirring rate :math:`2Ma/r^3` at
# large :math:`r`, but rises sharply approaching the ergosphere.
fig, ax = plt.subplots(figsize=(6, 4.5))
r = np.linspace(1.01, 15.0, 300)
for a in [0.3, 0.7, 0.998]:
    bh = KerrBlackHole(M=1.0, a=a)
    r_valid = r[r > bh.outer_horizon_radius]
    omega = bh.frame_dragging_angular_velocity(r_valid)
    ax.plot(r_valid, omega, label=f"a/M={a}")
ax.set_xlabel("r [M]")
ax.set_ylabel(r"ZAMO angular velocity $\omega(r)$")
ax.set_title("Lense-Thirring frame dragging")
ax.legend()
plt.tight_layout()

# %%
# ISCO shrinks (prograde) or grows (retrograde) with spin
# --------------------------------------------------------------------------
# The Bardeen-Press-Teukolsky (1972) innermost stable circular orbit radius,
#
# .. math::
#
#     \frac{r_{\text{ISCO}}}{M} = 3 + Z_2 \mp \sqrt{(3-Z_1)(3+Z_1+2Z_2)}
#
# with :math:`Z_1 = 1+(1-a_*^2)^{1/3}[(1+a_*)^{1/3}+(1-a_*)^{1/3}]`,
# :math:`Z_2=\sqrt{3a_*^2+Z_1^2}`, and :math:`a_*=a/M`, uses the
# :math:`-` sign for a prograde (co-rotating) orbit and :math:`+` for
# retrograde: prograde ISCOs shrink toward the horizon as spin increases,
# while retrograde ISCOs grow toward :math:`9M`.
a_values = np.linspace(0.0, 0.999, 60)
isco_pro = [KerrBlackHole(M=1.0, a=a).isco_radius(prograde=True) for a in a_values]
isco_retro = [KerrBlackHole(M=1.0, a=a).isco_radius(prograde=False) for a in a_values]

fig, ax = plt.subplots(figsize=(6, 4.5))
ax.plot(a_values, isco_pro, label="prograde")
ax.plot(a_values, isco_retro, label="retrograde")
ax.axhline(1.0, color="k", linestyle="--", linewidth=1, alpha=0.5, label="horizon, extremal limit")
ax.set_xlabel("spin a/M")
ax.set_ylabel("$r_{ISCO}$ [M]")
ax.set_title("Innermost stable circular orbit vs. spin")
ax.legend()
plt.tight_layout()
plt.show()