"""Lambert: the one-parameter family of conics through two fixed points."""
import numpy as np
import matplotlib.pyplot as plt
from _manstyle import CREAM, INK, MAROON, BLUE, GREEN, AMBER, PURPLE, despine
from _lambertp import conic_through, p_window, tof_of_p

r1, r2, dth = 8000.0, 14000.0, np.radians(110.0)
P1, P2 = np.array([r1, 0.0]), r2*np.array([np.cos(dth), np.sin(dth)])
pi_, pii = p_window(r1, r2, dth)

fig, ax = plt.subplots(figsize=(8.4, 6.6))
cases = [(pi_*1.25, PURPLE, "ellipse, $p$ near $p_{\\rm i}$"),
         (0.5*(pi_ + pii), BLUE, "ellipse"),
         (pii*0.985, GREEN, "ellipse, near-parabolic"),
         (pii, AMBER, "parabola, $p=p_{\\rm ii}$"),
         (1.35*pii, MAROON, "hyperbola")]
for p, col, lab in cases:
    e, f1 = conic_through(p, r1, r2, dth)
    lim = np.arccos(-1/e) - 1e-3 if e >= 1 else np.pi
    ff = np.linspace(-lim, lim, 1500)
    rr = p/(1 + e*np.cos(ff))
    ok = rr < 4.0*r2
    ang = ff - f1
    ax.plot(rr[ok]*np.cos(ang[ok]), rr[ok]*np.sin(ang[ok]), color=col, lw=0.9, alpha=0.35)
    fa = np.linspace(f1, f1 + dth, 300)
    ra = p/(1 + e*np.cos(fa))
    t = tof_of_p(p*(1 - 1e-9) if e == 1 else p, r1, r2, dth)/3600
    xa, ya = ra*np.cos(fa - f1), ra*np.sin(fa - f1)
    ax.plot(xa, ya, color=col, lw=2.6,
            label=f"{lab}:  $p={p:.0f}$ km,  $e={e:.2f}$,  $\\Delta t={t:.2f}$ h")
    for frac in (0.35, 0.72):                       # direction of motion: f grows from f1 to f1 + dtheta
        i = int(frac*(len(fa) - 1))
        ax.annotate("", xy=(xa[i + 1], ya[i + 1]), xytext=(xa[i], ya[i]), zorder=7,
                    arrowprops=dict(arrowstyle="-|>", color=col, lw=2.6, mutation_scale=20))

ax.plot(*zip(P1, P2), color=INK, lw=1.2, ls=(0, (4, 3)))
mid = 0.5*(P1 + P2)
ax.text(mid[0] - 2700, mid[1] - 700, "chord $c$", fontsize=11)
for P, lab, dx in ((P1, r"$\mathbf{r}_1$", 500), (P2, r"$\mathbf{r}_2$", -2600)):
    ax.plot([0, P[0]], [0, P[1]], color=INK, lw=1.4)
    ax.plot(*P, "o", color=INK, ms=8, zorder=6)
    ax.text(P[0] + dx, P[1] + 500, lab, fontsize=13)
th = np.linspace(0, dth, 60)
ax.plot(2500*np.cos(th), 2500*np.sin(th), color=INK, lw=1)
ax.text(2700*np.cos(dth/2), 2700*np.sin(dth/2), r"$\Delta\theta$", fontsize=13)
ax.plot(0, 0, "o", color=INK, ms=9)
ax.text(-1500, -1500, "focus", fontsize=10)
ax.set_xlim(-26000, 22000); ax.set_ylim(-17000, 30000)
ax.set_aspect("equal"); ax.axis("off")
ax.legend(frameon=False, fontsize=9.5, loc="lower left")
ax.set_title(r"Two points, one focus: every $p$ in the window gives a conic --- and a different clock",
             fontsize=11.5)
fig.tight_layout()
fig.savefig("img/fig-lam-family.png", dpi=160, bbox_inches="tight")
print("wrote img/fig-lam-family.png")