r"""
Chandrasekhar's dynamical friction: a satellite spiraling inward
=====================================================================

Chandrasekhar (1943) showed that a massive body moving through a sea of
much lighter field stars feels a systematic drag, built up from the
cumulative small-angle gravitational deflections of every star it
passes:

.. math::

    \frac{d\vec v_M}{dt} = -4\pi G^2 M \rho\,\ln\Lambda\,
    \frac{f(v_M)}{v_M^3}\,\vec v_M,

with :math:`\rho` the local background density. This example puts a
test satellite of mass :math:`M` on a circular orbit inside an NFW halo
and evolves its orbital radius under this drag, using the *fast*
(:math:`v\gg\sigma`) limit :math:`f(v)/v^3\to1/v^2` appropriate for a
satellite falling through the dilute outer halo, with
:func:`~physicskit.astro.galactic_dynamics.nfw_density` supplying
exactly the local density that sets the drag's strength, and
:func:`~physicskit.astro.galactic_dynamics.nfw_potential` giving the
orbital energy whose loss drives the satellite's inward spiral.
"""

# %%
import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import solve_ivp

from physicskit.astro.galactic_dynamics import (
    circular_velocity,
    nfw_density,
    nfw_enclosed_mass,
    nfw_potential,
)

# %%
# The host halo and the satellite
# -------------------------------------
rho_s, r_s = 1.0, 1.0
ln_Lambda = 5.0
M_sat = 0.02  # satellite mass, in the same G=1 units as the halo


def v_circ(r):
    return circular_velocity(r, lambda rr: nfw_enclosed_mass(rr, rho_s, r_s))


def specific_energy(r):
    r"""Specific orbital energy of a circular orbit at radius r, per unit satellite mass."""
    return 0.5 * v_circ(r) ** 2 + nfw_potential(r, rho_s, r_s)


def drag_deceleration(r):
    r"""Fast-satellite (v >> sigma) limit of Chandrasekhar's dynamical friction, |dv/dt|."""
    return 4.0 * np.pi * M_sat * nfw_density(r, rho_s, r_s) * ln_Lambda / v_circ(r) ** 2


# %%
# Orbital decay as an energy-loss ODE
# -----------------------------------------
# For a satellite spiraling slowly inward through a near-circular
# sequence of orbits, the drag dissipates specific orbital energy at
# rate :math:`dE/dt=-a_{\rm drag}(r)\,v_{\rm circ}(r)`; converting that
# to a radius-decay rate needs :math:`dE/dr`, evaluated here by a simple
# centered finite difference on :func:`specific_energy`.
def dr_dt(t, r):
    r = r[0]
    h = 1e-4 * r
    dE_dr = (specific_energy(r + h) - specific_energy(r - h)) / (2 * h)
    dE_dt = -drag_deceleration(r) * v_circ(r)
    return [dE_dt / dE_dr]


r0 = 8.0  # several scale radii out, in the dilute outer halo
r_floor = 0.05 * r_s  # stop before the fast-satellite/point-mass approximation breaks down


def reached_floor(t, r):
    return r[0] - r_floor


reached_floor.terminal = True
reached_floor.direction = -1

t_span = (0.0, 4000.0)
sol = solve_ivp(dr_dt, t_span, [r0], events=reached_floor, dense_output=True, rtol=1e-9, atol=1e-12, max_step=5.0)

t_end = sol.t[-1]
t_plot = np.linspace(0.0, t_end, 400)
r_plot = sol.sol(t_plot)[0]

print(f"initial orbital radius: r0 = {r0}")
print(f"time to spiral down to r={r_floor} (a small fraction of r_s): t = {t_end:.1f}")
print(f"final orbital radius: r = {r_plot[-1]:.4f}")
print(f"drag deceleration at r0: {drag_deceleration(r0):.3e}  (grows as the satellite sinks into denser halo gas)")
print(f"drag deceleration at r={r_plot[-1]:.2f}: {drag_deceleration(r_plot[-1]):.3e}")

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 4.5))
ax1.plot(t_plot, r_plot, color="steelblue")
ax1.set_xlabel("t")
ax1.set_ylabel("orbital radius r")
ax1.set_title("Dynamical friction: a satellite's orbit decaying inward")

r_grid = np.linspace(0.3, r0, 200)
ax2.semilogy(r_grid, drag_deceleration(r_grid), color="firebrick")
ax2.set_xlabel("r")
ax2.set_ylabel(r"drag deceleration $|dv/dt|$")
ax2.set_title("Drag strengthens as the satellite sinks inward")
fig.tight_layout()

plt.show()