"""
three-body-animation.py
=======================
Side-by-side animation of the three-body problem in two regimes:

  LEFT  — Chaotic orbit (generic, non-symmetric initial conditions)
  RIGHT — Figure-eight periodic orbit (Chenciner & Montgomery 2000)

The figure-eight is one of the very few exact periodic solutions known for
the three-body problem. Both cases use equal masses and G = 1 (non-dimensional
units). The contrast illustrates Bruns' theorem: no closed-form solution exists
in general, yet special solutions with rich symmetry can be periodic.

Dependencies
------------
  numpy, scipy, matplotlib

Usage
-----
  python three-body-animation.py

To save (uncomment the ani.save() line at the bottom):
  - GIF requires pillow  (pip install pillow)
  - MP4 requires ffmpeg  (system package)

Reference
---------
  Chenciner, A. & Montgomery, R. (2000). A remarkable periodic solution of
  the three-body problem in the case of equal masses. Annals of Mathematics,
  152(3), 881–901.
"""

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

# ── ODE right-hand side ───────────────────────────────────────────────────────

G = 1.0   # gravitational constant (non-dimensional)


def nbody_rhs(t, y, masses):
    """First-order form of the N-body equations of motion (Battin Eq. 2.39)."""
    N = len(masses)
    pos = y[:3 * N].reshape(N, 3)
    vel = y[3 * N:].reshape(N, 3)
    acc = np.zeros_like(pos)
    for i in range(N):
        for j in range(N):
            if i != j:
                r_vec = pos[j] - pos[i]          # r_j - r_i  (attractive)
                r_mag = np.linalg.norm(r_vec)
                acc[i] += G * masses[j] / r_mag ** 3 * r_vec
    return np.concatenate([vel.ravel(), acc.ravel()])


def integrate(pos0, vel0, masses, t_end, n_pts=5000):
    """Integrate with DOP853 (8th-order Runge-Kutta), tight tolerances."""
    y0 = np.concatenate([pos0.ravel(), vel0.ravel()])
    t_eval = np.linspace(0, t_end, n_pts)
    sol = solve_ivp(
        nbody_rhs, [0, t_end], y0, args=(masses,),
        t_eval=t_eval, method='DOP853',
        rtol=1e-11, atol=1e-12,
    )
    N = len(masses)
    pos_out = sol.y[:3 * N].reshape(N, 3, -1)   # shape: (N, xyz, steps)
    return sol.t, pos_out


# ── Initial conditions: figure-eight ─────────────────────────────────────────
# Chenciner & Montgomery (2000) / Moore (1993).
# Three unit masses, G = 1.  Period T ≈ 6.3259.
# CoM is at the origin; total linear and angular momentum are zero.

masses_f8 = np.ones(3)

pos0_f8 = np.array([
    [-0.97000436,  0.24308753, 0.0],
    [ 0.97000436, -0.24308753, 0.0],
    [ 0.0,         0.0,        0.0],
])
vel0_f8 = np.array([
    [ 0.46620368,  0.43236573, 0.0],
    [ 0.46620368,  0.43236573, 0.0],
    [-0.93240737, -0.86473146, 0.0],
])
T_F8      = 6.3259          # one period (non-dimensional time units)
t_end_f8  = 2.05 * T_F8    # just over two full periods


# ── Initial conditions: chaotic ───────────────────────────────────────────────
# Three equal masses at the vertices of a scalene triangle with small
# transverse velocities.  CoM is fixed at the origin.  The asymmetry of
# the configuration guarantees no hidden periodic symmetry.

masses_ch = np.ones(3)

_pos_raw = np.array([
    [-1.0,  0.2, 0.0],
    [ 1.1, -0.1, 0.0],
    [ 0.1, -1.3, 0.0],
], dtype=float)
_vel_raw = np.array([
    [ 0.15,  0.45, 0.0],
    [-0.35,  0.10, 0.0],
    [ 0.20, -0.55, 0.0],
], dtype=float)

# Shift so that CoM position and CoM velocity are exactly zero
pos0_ch = _pos_raw - _pos_raw.mean(axis=0)
vel0_ch = _vel_raw - _vel_raw.mean(axis=0)

t_end_ch = 18.0


# ── Integrate ─────────────────────────────────────────────────────────────────

print("Integrating figure-eight orbit …")
t_f8, pos_f8 = integrate(pos0_f8, vel0_f8, masses_f8, t_end_f8)
print(f"  {len(t_f8)} steps,  final t = {t_f8[-1]:.4f}  (2 × T = {2*T_F8:.4f})")

print("Integrating chaotic orbit …")
t_ch, pos_ch = integrate(pos0_ch, vel0_ch, masses_ch, t_end_ch)
print(f"  {len(t_ch)} steps,  final t = {t_ch[-1]:.4f}")


# ── Animation ─────────────────────────────────────────────────────────────────

BODY_COLORS = ['#e6194b', '#3cb44b', '#4363d8']  # red, green, blue
TRAIL_LEN   = 450    # number of past steps shown as a fading trail
SKIP        = 6      # render every SKIP-th timestep (controls speed)

BG  = '#0d1117'
FG  = '#c9d1d9'
DIM = '#8b949e'

fig, (ax_ch, ax_f8) = plt.subplots(1, 2, figsize=(14, 6))
fig.patch.set_facecolor(BG)
fig.suptitle(
    'Three-Body Problem  —  Chaos vs. Periodic Motion',
    color=FG, fontsize=13, fontweight='bold', y=1.01,
)

for ax in (ax_ch, ax_f8):
    ax.set_facecolor(BG)
    ax.set_aspect('equal')
    ax.tick_params(colors=DIM, labelsize=8)
    for sp in ax.spines.values():
        sp.set_edgecolor('#30363d')

ax_ch.set_title('Chaotic orbit\n(generic initial conditions)',
                color=FG, fontsize=10, pad=8)
ax_ch.set_xlim(-2.8, 2.8)
ax_ch.set_ylim(-2.8, 2.8)

ax_f8.set_title('Figure-eight periodic orbit\n(Chenciner & Montgomery 2000)',
                color=FG, fontsize=10, pad=8)
ax_f8.set_xlim(-1.5, 1.5)
ax_f8.set_ylim(-0.75, 0.75)

# Artist objects: one dot + one trail per body per panel
dots_ch,  trails_ch  = [], []
dots_f8,  trails_f8  = [], []

for col in BODY_COLORS:
    d, = ax_ch.plot([], [], 'o', color=col, ms=10, zorder=3)
    tr, = ax_ch.plot([], [], '-', color=col, alpha=0.5, lw=1.3)
    dots_ch.append(d);  trails_ch.append(tr)

    d, = ax_f8.plot([], [], 'o', color=col, ms=10, zorder=3)
    tr, = ax_f8.plot([], [], '-', color=col, alpha=0.5, lw=1.3)
    dots_f8.append(d);  trails_f8.append(tr)

txt_ch = ax_ch.text(
    0.03, 0.97, '', transform=ax_ch.transAxes,
    color=DIM, fontsize=9, va='top', fontfamily='monospace',
)
txt_f8 = ax_f8.text(
    0.03, 0.97, '', transform=ax_f8.transAxes,
    color=DIM, fontsize=9, va='top', fontfamily='monospace',
)

# Annotate period on the figure-eight panel
ax_f8.text(
    0.03, 0.06,
    f'Period T ≈ {T_F8:.4f}',
    transform=ax_f8.transAxes, color=DIM, fontsize=8, fontfamily='monospace',
)

n_ch = len(t_ch[::SKIP])
n_f8 = len(t_f8[::SKIP])
n_frames = max(n_ch, n_f8)


def update(frame):
    ic = min(frame, n_ch - 1) * SKIP
    i8 = min(frame, n_f8 - 1) * SKIP

    for k in range(3):
        # chaotic panel
        s = max(0, ic - TRAIL_LEN)
        dots_ch[k].set_data([pos_ch[k, 0, ic]], [pos_ch[k, 1, ic]])
        trails_ch[k].set_data(pos_ch[k, 0, s:ic + 1], pos_ch[k, 1, s:ic + 1])

        # figure-eight panel
        s = max(0, i8 - TRAIL_LEN)
        dots_f8[k].set_data([pos_f8[k, 0, i8]], [pos_f8[k, 1, i8]])
        trails_f8[k].set_data(pos_f8[k, 0, s:i8 + 1], pos_f8[k, 1, s:i8 + 1])

    txt_ch.set_text(f't = {t_ch[ic]:.2f}')
    txt_f8.set_text(f't = {t_f8[i8]:.3f}')

    return (*dots_ch, *trails_ch, *dots_f8, *trails_f8, txt_ch, txt_f8)


ani = animation.FuncAnimation(
    fig, update, frames=n_frames, interval=20, blit=True,
)

plt.tight_layout()

# Uncomment to save:
ani.save('img/three-body-animation.gif', writer='pillow', fps=40)
# ani.save('img/three-body-animation.mp4', writer='ffmpeg', fps=40, dpi=150)