#!/usr/bin/env python3

# pykepler: Numerically solve the restricted 3-body problem and animate the
# movement of the bodies
#
# Copyright (C) 2025 Thomas Guillermo Albers Raviola
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

from functools import partial

import numpy as np
from scipy.integrate import solve_ivp

import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib import rcParams
from matplotlib.animation import FuncAnimation

def unpack_r(rv):
    return rv[0:2], rv[2:4], rv[4:6]

def unpack_v(rv):
    return rv[6:8], rv[8:10], rv[10:12]

def distance(x1, x2):
    return np.sqrt(np.sum((x1 - x2)**2))

def grav_force(r1, r2):
    d = distance(r1, r2)
    return (r2 - r1) / d**3

#
# Initial value problem has form
# x' = F(x, t)
#
# Set x = [r, v], where r is position of body and v = r'
# Thus, x' = [r', v'] = [v, a] = [v, -G*M*m/abs(d)**2 * e]
#

def func(t, rv, mu):
    # Unpack components
    r1, r2, r3 = unpack_r(rv)
    v1, v2, v3 = unpack_v(rv)

    f12 = grav_force(r1, r2)
    f13 = grav_force(r1, r3)
    f23 = grav_force(r2, r3)

    return np.concatenate((v1, v2, v3,
                           mu[1] * f12 + mu[2] * f13,
                           -mu[0] * f12 + mu[2] * f23,
                           -mu[1] * f23 - mu[0] * f13))


def angle(x1, x2):
    ux1 = x1 / np.linalg.norm(x1)
    ux2 = x2 / np.linalg.norm(x2)
    return 180.0 * np.arccos(np.clip(np.dot(ux1, ux2), -1.0, 1.0)) / np.pi


def update(sol, text_artist, tri_artist, pnt_artists, pos_artists, frame):
    '''Update plot of the IVP up to the current timestamp'''

    if frame > 0:
        r1, r2, r3 = unpack_r(sol.y)

        r1 = r1[:, frame-1]
        r2 = r2[:, frame-1]
        r3 = r3[:, frame-1]

        tri_artist.set_data([r1[0], r2[0], r3[0], r1[0]],
                            [r1[1], r2[1], r3[1], r1[1]])

        alpha = angle(r2 - r1, r3 - r1)
        beta = angle(r1 - r2, r3 - r2)
        gamma = angle(r1 - r3, r2 - r3)

        a = distance(r2, r3)
        b = distance(r3, r1)
        c = distance(r2, r1)

        fmt = 'α:{:.2f}, β:{:.2f}, γ:{:.2f}\na: {:.2f}, b:{:.2f}, c:{:.2f}'
        text_artist.set_text(fmt.format(alpha, beta, gamma, a, b, c))


    for artist, r in zip(pos_artists, unpack_r(sol.y)):
        artist.set_data(r[0, 0:frame], r[1, 0:frame])

    for artist, r in zip(pnt_artists, unpack_r(sol.y)):
        artist.set_data([r[0, frame-1]], [r[1, frame-1]])

    return [text_artist, tri_artist] + pos_artists + pnt_artists


def animate_ivp(sol, fps=20):
    def progress(current_frame, total_frames):
        str = '\r  Rendering frame [{}/{}] ...'
        print(str.format(current_frame + 1, total_frames), end='', flush=True)

    fig, ax = plt.subplots(figsize=(10,10), dpi=300)
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_xlim(-1, 1)
    ax.set_ylim(-1, 1)

    pos_artists = []
    pnt_artists = []

    tri_artist, = ax.plot([], [], ls='--', color='black')
    for i in range(3):
        pos_artists.append(ax.plot([], [])[0])
        pnt_artists.append(ax.plot([], [], '.', label='m{}'.format(i + 1))[0])

    text_artist = ax.text(0.05, 0.95, '',
                          backgroundcolor='#7F7F7F33',
                          transform=ax.transAxes,
                          verticalalignment='top') # bbox=props

    ax.legend()

    frames = np.arange(0, len(sol.t))
    animation = FuncAnimation(fig,
                              partial(update, sol,
                                      text_artist, tri_artist,
                                      pnt_artists, pos_artists),
                              frames=frames,
                              interval=sol.t[-1], blit=True, repeat=False)

    animation.save('three-body.mp4', fps=fps, progress_callback=progress)
    # Reserve newline for progress
    print('')


def main():
    mu1 = .7
    mu2 = 1 - mu1
    mu = np.array([mu1, mu2, 1e-12])
    omega = 1

    ir1 = mu[1] * np.array([-1, 0])
    ir2 = mu[0] * np.array([1, 0])
    # Trojan satelite at L4 (ahead)
    ir3 = np.array([.5 - mu2, np.sqrt(3) / 2])

    iv1 = mu[1] * np.array([0, -1])
    iv2 = mu[0] * np.array([0, 1])
    # yz - zy, zx - xz
    iv3 = np.array([- ir3[1], ir3[0]])

    irv = np.concatenate((ir1, ir2, ir3, iv1, iv2, iv3))

    t_span = (0.0, 4)
    t_eval = np.linspace(t_span[0], t_span[1], 1024)

    sol = solve_ivp(func, t_span, irv,
                    t_eval=t_eval, method='DOP853',
                    rtol=1e-10, atol=1e-20,
                    args=(mu,))

    animate_ivp(sol)


if __name__ == '__main__':
    main()