"""Lambert: time of flight as a function of p, short way and long way."""
import numpy as np
import matplotlib.pyplot as plt
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine
from _lambertp import p_window, tof_of_p, lambert_p

r1, r2 = 6678.137, 42164.0
fig, axs = plt.subplots(1, 2, figsize=(12.4, 5.2))

for ax, deg in zip(axs, (175.0, 185.0)):
    dth = np.radians(deg)
    pi_, pii = p_window(r1, r2, dth)
    if deg < 180:
        ps = np.linspace(pi_*(1 + 1e-5), 1.25*pii, 3000)
    else:
        ps = np.linspace(0.02*pi_, pii*(1 - 1e-5), 3000)
    ts = np.array([tof_of_p(p, r1, r2, dth) for p in ps])/3600
    e_ok = (ps > pi_) & (ps < pii)
    ax.semilogy(ps[e_ok], ts[e_ok], color=BLUE, lw=2.6, label="ellipse")
    ax.semilogy(ps[~e_ok], ts[~e_ok], color=MAROON, lw=2.6, label="hyperbola")
    for pp, lab, ha in ((pi_, r"$p_{\rm i}$", "right"), (pii, r"$p_{\rm ii}$", "left")):
        ax.axvline(pp, color=AMBER, lw=1.4, ls=(0, (5, 3)))
        ax.text(pp, 0.065, f" {lab} ", color=AMBER, fontsize=12, ha=ha)
    for dt, col in ((4.0, GREEN), (2.0, PURPLE)):
        p, e, *_ = lambert_p(r1, r2, dth, dt*3600)
        ax.axhline(dt, color=col, lw=1, alpha=0.6)
        ax.plot(p, dt, "o", color=col, ms=9, zorder=6)
        ax.annotate(f"$\\Delta t={dt:.0f}$ h: $p={p:.1f}$, $e={e:.4f}$", xy=(p, dt),
                    xytext=(12, 8 if dt > 3 else -18), textcoords="offset points",
                    color=col, fontsize=10)
    ax.set_ylim(0.05, 500)
    ax.set_xlabel(r"semi-latus rectum $p$  [km]")
    ax.set_ylabel(r"time of flight $\Delta t$  [h]")
    way = "short way, TOF falls with $p$" if deg < 180 else "long way, TOF rises with $p$"
    ax.set_title(f"$\\Delta\\theta={deg:.0f}^\\circ$ --- {way}", fontsize=11.5)
    ax.legend(frameon=False, fontsize=10, loc="upper right" if deg < 180 else "upper left")
    despine(ax)

fig.tight_layout()
fig.savefig("img/fig-lam-tof.png", dpi=160, bbox_inches="tight")
print("wrote img/fig-lam-tof.png")