"""Figures for ballistic-trajectories-bates.qmd (img/fig-bal-land-lt.png, -eq.png, -gt.png): does the flight come down?

For Q < 1 (0.97), Q = 1 and Q > 1 (1.2): the paths from one burnout point 1000 km above the surface (top) for several flight-path angles.
Solid: the path of the vehicle.  If the ellipse dips below the surface, the path runs from burnout to the impact point (a cross) and
the rest of the conic (the part under the ground and the part before burnout) is dotted.  If the ellipse stays above the surface it is
a satellite orbit: the whole ellipse is flown and there is no 'rest'.
The dot on a landing path is where the note's free flight ends (re-entry at r_re = r_bo); psi in the legend is measured to it.
1000 km is higher than a missile's burnout; it is chosen so that the gap between the surface and a non-landing ellipse can be seen.
With r_bo = R_E every gamma > 0 lands (fig-bal-paths.py).

Run from the sfm-slides directory:   python3 codes/fig-bal-landing.py
"""
import numpy as np
import _figfont                                # all figure text at one scale
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from matplotlib.lines import Line2D
from _manstyle import INK, MAROON, BLUE, GREEN, AMBER, PURPLE
from _ballistic import RE, psi_of, ecc, lands, perigee_radius, psi_impact, gamma_land

D = np.pi / 180
H = 1000.0
R_BO = RE + H
GREY = "#8a8478"
CASES = (
    ("lt", 0.97, [1.0, 3.0, 10.0, 25.0, 45.0, 65.0]),
    ("eq", 1.00, [2.0, 5.0, 12.0, 25.0, 45.0, 70.0]),
    ("gt", 1.20, [3.0, 9.0, 20.0, 35.0, 55.0]),
)
NOLAND = [BLUE, GREEN]
LAND = [AMBER, MAROON, PURPLE, "#a3315f", "#5b3a1e"]
DOT = (0, (1, 2.4))


def conic(Q, g):
    """The ellipse through the burnout point as a function of true anomaly; burnout at the top, motion clockwise.

    Returns f (from burnout round to burnout again), the points (km) and the burnout true anomaly.
    """
    e = ecc(Q, g); a = R_BO / (2 - Q); p = a * (1 - e * e)
    fbo = np.pi - psi_of(Q, g) / 2
    f = np.linspace(fbo, fbo + 2 * np.pi, 3000)
    r = p / (1 + e * np.cos(f)); ang = np.pi / 2 - (f - fbo)
    return f, r * np.cos(ang), r * np.sin(ang), fbo


if __name__ == "__main__":
    for tag, Q, gl in CASES:
        fig, A = plt.subplots(figsize=(10.4, 6.4))
        A.add_patch(Circle((0, 0), RE, color="#dfe9f3", ec=BLUE, lw=1.4, zorder=1))
        A.add_patch(Circle((0, 0), R_BO, fill=False, ec=GREY, lw=0.8, ls=(0, (5, 3)), alpha=0.8, zorder=2))
        nl = ll = 0
        xs, ys = [0], [R_BO]
        handles = []
        for gd in gl:
            g = gd * D
            f, x, y, fbo = conic(Q, g)
            psi = psi_of(Q, g)
            rp = perigee_radius(Q, g, R_BO)
            if lands(Q, g, R_BO):
                col = LAND[ll]; ll += 1
                pim = psi_impact(Q, g, R_BO)
                flown = f <= fbo + pim
                A.plot(x[~flown], y[~flown], color=col, lw=1.7, ls=DOT, zorder=3)
                A.plot(x[flown], y[flown], color=col, lw=2.4, zorder=5)
                ang = np.pi / 2 - pim
                A.scatter([RE * np.cos(ang)], [RE * np.sin(ang)], marker="x", color=col, s=70, linewidths=2.4, zorder=8)
                k = np.argmin(abs(f - (2 * np.pi - fbo)))               # the point of re-entry, r = r_bo, f = 2 pi - f_bo
                A.scatter([x[k]], [y[k]], color=INK, s=16, zorder=7)
                lab = f"$\\gamma={gd:g}^\\circ$: lands, $\\psi={np.degrees(psi):.0f}^\\circ$, $r_p={rp:.0f}$ km"
            else:
                col = NOLAND[nl]; nl += 1
                A.plot(x, y, color=col, lw=2.4, zorder=5)               # the whole ellipse is flown: a satellite orbit
                lab = f"$\\gamma={gd:g}^\\circ$: stays up, $r_p={rp:.0f}$ km"
            xs += list(x); ys += list(y)
            handles.append(Line2D([0], [0], color=col, lw=2.4, label=lab))
            print(f"Q={Q} gamma={gd:5.1f}: psi={np.degrees(psi):7.2f}  r_p={rp:8.1f}  {'lands' if lands(Q, g, R_BO) else 'stays up'}")
        A.scatter([0], [R_BO], color=AMBER, s=60, zorder=9, edgecolor=INK, linewidths=0.6)
        A.text(-0.05 * R_BO, R_BO * 1.03, "burnout", color=AMBER, fontsize=10.5, ha="right", va="bottom")
        A.scatter([0], [0], color=INK, s=14, zorder=6); A.text(0.03 * RE, -0.12 * RE, "$F$", fontsize=11)
        handles += [Line2D([0], [0], color=INK, lw=2.4, label="path of the vehicle (solid)"),
                    Line2D([0], [0], color=INK, lw=1.7, ls=DOT, label="rest of the conic (dotted)"),
                    Line2D([0], [0], color=INK, marker="x", ms=8, mew=2.2, lw=0, label="impact point"),
                    Line2D([0], [0], color=INK, marker="o", ms=4, lw=0, label="end of free flight ($r_{re}=r_{bo}$)"),
                    Line2D([0], [0], color=GREY, lw=0.9, ls=(0, (5, 3)), label="burnout radius $r_{bo}$")]
        xs, ys = np.array(xs), np.array(ys)
        A.set_aspect("equal"); A.axis("off")
        A.set_xlim(1.04 * min(xs.min(), -R_BO), 1.04 * xs.max()); A.set_ylim(1.04 * min(ys.min(), -R_BO), 1.04 * ys.max())
        A.legend(handles=handles, loc="center left", bbox_to_anchor=(1.0, 0.5), fontsize=9.6, frameon=False)
        qs = {"lt": "$Q=0.97<1$", "eq": "$Q=1$", "gt": "$Q=1.2>1$"}[tag]
        A.set_title(f"{qs}, burnout 1000 km above the surface", fontsize=12, loc="left")
        fig.savefig(f"img/fig-bal-land-{tag}.png", dpi=150, bbox_inches="tight")
        plt.close(fig)
    print("threshold gamma [deg]:", {Q: round(gamma_land(Q, R_BO) / D, 2) for Q in (0.97, 1.0, 1.2)})