"""Animation of the phasing-loop trajectory of lunar-phasing-and-low-thrust.qmd, on the engine of codes/_manim.py.

Left: the craft injects towards the Moon at once, from the first perigee; it reaches the apogee A of the transfer ellipse and the Moon is not there.
Right: the craft first flies three phasing loops (burn 1 raises the apogee to 137 300 km, burn 2 at the third return is the rest of the trans-lunar
injection) and reaches the Moon's sphere of influence as the Moon does.  Both panels start together and share one clock.

The Moon moves at n_Moon = sqrt((mu_E + mu_M)/D^3), as in the notes; the sphere of influence is drawn to scale.  The numbers the note quotes are asserted.
Run from the sfm-slides directory:   python3 codes/anim-lpl.py [--check | --debug]     (afterwards:  python3 _filters/gif-frames.py  from the project root)
"""
import os
import sys
import numpy as np
from scipy.optimize import brentq
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER, PURPLE
from _manim import MU, Orbit, Leg, Companion, Panel, animate
from _lunar import N_MOON, D, R_SOI
from _lunarlt import R0, V_C0, DAY, HOUR, hohmann_to_moon, phasing_orbit, required_wait, loop_period

pi = np.pi
MOONC = "#6b655b"


class MoonOrbit(Orbit):
    """The Moon's circle, run at n_Moon (the notes' value, which includes the Moon's own mass) instead of the engine's sqrt(mu_E/D^3)."""

    def time(self, f):
        return np.asarray(f, float)/N_MOON

    def f_of_time(self, t):
        return np.mod(N_MOON*np.asarray(t, float), 2*pi)


def build():
    H = hohmann_to_moon()
    phi0 = np.radians(123.068)                                    # Example 2a of patched-conics.qmd (found again in the verify script)
    phi_a, k, pre = np.radians(35.0), 3, 0.4*HOUR
    W = required_wait(phi0, phi_a)
    P3 = phasing_orbit(loop_period(W, k))
    leo = Orbit.circ(R0, color=BLUE)
    loop = Orbit.rprq(R0, P3["ra"], w=0.0, color=MAROON)
    tli = Orbit.rprq(R0, D, w=0.0, color=GREEN)
    moon_c = MoonOrbit.circ(D, color=MOONC)
    n_leo = np.sqrt(MU/R0**3)
    f_leo0 = 2*pi - n_leo*pre                                     # the craft starts 0.4 h before the first burn, at perigee (angle 0)

    # where the transfer ellipse meets the Moon's sphere of influence, for the Moon at phi_0 at the moment of the trans-lunar burn
    def gap(t):
        f = tli.f_of_time(t)
        return np.linalg.norm(tli.pos(f)[:2] - D*np.array([np.cos(phi0 + N_MOON*t), np.sin(phi0 + N_MOON*t)])) - R_SOI
    ts = np.linspace(0, H["T"]/2, 4000)
    g = np.array([gap(t) for t in ts])
    i = int(np.argmax((g[:-1] > 0) & (g[1:] <= 0)))
    t1 = brentq(gap, ts[i], ts[i + 1], xtol=1e-6)
    f_patch = float(tli.f_of_time(t1))
    assert abs(t1/HOUR - 78.859) < 0.005, t1/HOUR                 # Example 2a: 78.86 h

    # right panel: three loops, then the trans-lunar ellipse to the sphere of influence
    legs_b = [Leg(leo, f_leo0, 2*pi), Leg(loop, 0.0, 2*pi*k), Leg(tli, 0.0, f_patch)]
    moon_b = Companion(moon_c, phi_a - N_MOON*pre, color=MOONC, ghost=True, ms=22, label="Moon", label_off=(14000, 14000), halo=R_SOI)
    lims = ((-440000, 440000), (-440000, 440000))
    labels = [dict(text="Moon's orbit", pos=(260000, -415000), color=MOONC)]
    PB = Panel("(b) after three phasing loops", legs_b, lims, companions=[moon_b], arrow=0.045, labels=labels, outlines=[moon_c],
               expect=dict(dv=[2.9437, 0.1627], tol=6e-5, leg_h=[0.4, W/HOUR, t1/HOUR], kinds=["prograde", "prograde"], tol_h=2e-3))
    t_end_b = PB.info["t_total"]
    # the Moon at the end of panel (b) is where the craft meets its sphere of influence
    mb = moon_b.pos(t_end_b)[:2]
    assert abs(np.linalg.norm(PB.legs[2].end[0][:2] - mb) - R_SOI) < 1.0
    ang = np.degrees(np.arctan2(moon_b.pos(pre + W)[1], moon_b.pos(pre + W)[0]))
    assert abs(ang - 123.068) < 1e-3, ang                          # at the trans-lunar burn the Moon is phi_0 ahead of the perigee

    # left panel: the same transfer ellipse, burnt at once (the whole trans-lunar injection at the first perigee), to the apogee A
    legs_a = [Leg(leo, f_leo0, 2*pi), Leg(tli, 0.0, pi)]
    moon_a = Companion(moon_c, phi_a - N_MOON*pre, color=MOONC, ghost=True, ms=22, label="Moon", label_off=(14000, 14000), halo=R_SOI)
    PA = Panel("(a) injection at once: the craft reaches A, the Moon is elsewhere", legs_a, lims, companions=[moon_a], arrow=0.045, labels=labels[:], outlines=[moon_c],
               expect=dict(dv=[H["dv_tli"]], tol=6e-5, leg_h=[0.4, H["T"]/2/HOUR], kinds=["prograde"], tol_h=2e-3))
    ang_a = np.degrees(np.arctan2(moon_a.pos(PA.info["t_total"])[1], moon_a.pos(PA.info["t_total"])[0]))
    assert abs((180 - ang_a) - 79.29) < 0.02, ang_a               # the Moon is 79.3 deg short of A
    return [PA, PB], dict(K=1.3, dpi=60, colors=128, figsize=(15.8, 7.0),
                          suptitle="the craft (black dot) and the Moon (ring, with its sphere of influence to scale)")


if __name__ == "__main__":
    panels, kw = build()
    for P in panels:
        print(P.table())
    if "--check" in sys.argv:
        sys.exit(0)
    if "--debug" in sys.argv:
        animate(panels, "img/anim-lpl-phasing-loop.gif", debug=os.path.join(os.environ["ANIM_DEBUG"], "lpl"), **kw)
    else:
        animate(panels, "img/anim-lpl-phasing-loop.gif", **kw)