"""Dependency map of the SFM notes for courses/sfm/study-plan.qmd (img/fig-study-plan.png).

An arrow from A to B means "read A before B".  The numbers are the steps of the study plan; the shaded boxes
are the core of AE313 (stages 1 to 4).

Run from the sfm-slides directory:   python3 codes/fig-study-plan.py
"""
import matplotlib.pyplot as plt
from matplotlib.patches import FancyBboxPatch
from _manstyle import INK, BLUE

W, H = 170, 40
c1, c2, c3, c4 = 108, 292, 476, 660
r1, r2, r3, r4, r5, r6, r7, r8, r9 = 80, 150, 220, 290, 360, 430, 500, 570, 640

BOXES = {                                   # key: (centre x, centre y, label, core of AE313?)
    "nbody": (c2, r1, "1  N-body problem", True),
    "twobody": (c2, r2, "2  Two-body problem", True),
    "grav": (c4, r2, "13  Earth's gravity", False),
    "tut": (c1, r3, "3  Tutorial 1", True),
    "peri": (c2, r3, "4  Perifocal frame", True),
    "elem": (c2, r4, "5  Orbital elements", True),
    "time": (c1, r5, "6  Time", True),
    "man": (c2, r5, "7  Manoeuvres", True),
    "cw": (c3, r5, "11  CW equations", False),
    "patch": (c1, r6, "8  Patched conics", True),
    "lam": (c2, r6, "9  Lambert's problem", False),
    "gve": (c4, r6, "14  Gauss equations", False),
    "uv": (c1, r7, "10  Universal variables", False),
    "chase": (c3, r7, "12  Chase manoeuvres", False),
    "j2": (c4, r7, "15  What $J_2$ does", False),
    "gt": (c4, r8, "16  Ground tracks", False),
    "topo": (c4, r9, "17  Topocentric frames", False),
    "cr3bp": (c2, r7, "18  Restricted 3-body", False),
}
h, w = H / 2, W / 2
EDGES = [                                   # poly-lines in the 760-wide frame, last point gets the arrow head
    [(c2, r1 + h), (c2, r2 - h)],
    [(c2 + w, r1), (c4, r1), (c4, r2 - h)],
    [(c2 - w, r2), (c1, r2), (c1, r3 - h)],
    [(c2, r2 + h), (c2, r3 - h)],
    [(c2, r3 + h), (c2, r4 - h)],
    [(c2 - w, r4), (c1, r4), (c1, r5 - h)],
    [(c2, r4 + h), (c2, r5 - h)],
    [(c2 + w, r2), (c3, r2), (c3, r5 - h)],
    [(c2, r5 + h), (c2, r5 + 35), (c1, r5 + 35), (c1, r6 - h)],
    [(c2, r5 + h), (c2, r6 - h)],
    [(c2, r6 + h), (c2, r6 + 35), (c1 + 32, r6 + 35), (c1 + 32, r7 - h)],
    [(c1, r6 + h), (c1, r7 - h)],
    [(c2, r6 + h), (c2, r6 + 35), (c3 - 24, r6 + 35), (c3 - 24, r7 - h)],
    [(c3, r5 + h), (c3, r7 - h)],
    [(c2 + w, r4), (c4 - 40, r4), (c4 - 40, r6 - h)],
    [(c3 + w, r5), (c4, r5), (c4, r6 - h)],
    [(c4 + w, r2), (754, r2), (754, r7), (c4 + w, r7)],
    [(c4, r6 + h), (c4, r7 - h)],
    [(c4, r7 + h), (c4, r8 - h)],
    [(c1 - w, r5), (10, r5), (10, r8), (c4 - w, r8)],
    [(c4, r8 + h), (c4, r9 - h)],
    [(c3 - w, r5), (c3 - w - 7, r5), (c3 - w - 7, r7), (c2 + w, r7)],
]

fig, ax = plt.subplots(figsize=(10.5, 9.9))
ax.set_xlim(0, 762)
ax.set_ylim(715, 50)
ax.set_aspect("equal")
ax.axis("off")
for pts in EDGES:
    xs, ys = zip(*pts)
    ax.plot(xs[:-1] + (xs[-1],), ys, color="#8a857c", lw=1.1, solid_capstyle="butt")
    ax.annotate("", xy=pts[-1], xytext=pts[-2],
                arrowprops=dict(arrowstyle="-|>", color="#8a857c", lw=1.1, shrinkA=0, shrinkB=0,
                                mutation_scale=11))
for cx, cy, lab, core in BOXES.values():
    ax.add_patch(FancyBboxPatch((cx - w, cy - h), W, H, boxstyle="round,pad=0,rounding_size=8",
                                fc=BLUE if core else "none", alpha=1, lw=1.4 if core else 1.1,
                                ec=BLUE if core else "#8a857c", zorder=3))
    if core:
        ax.add_patch(FancyBboxPatch((cx - w, cy - h), W, H, boxstyle="round,pad=0,rounding_size=8",
                                    fc="white", alpha=0.85, lw=0, zorder=3))
    ax.text(cx, cy, lab, ha="center", va="center", fontsize=10.5, color=INK, zorder=4)
ax.add_patch(FancyBboxPatch((20, 686), 18, 14, boxstyle="round,pad=0,rounding_size=3", fc="#dfe7f4",
                            ec=BLUE, lw=1.4))
ax.text(46, 693, "core of AE313 (stages 1 to 4)", va="center", fontsize=10, color=INK)
ax.add_patch(FancyBboxPatch((280, 686), 18, 14, boxstyle="round,pad=0,rounding_size=3", fc="none",
                            ec="#8a857c", lw=1.1))
ax.text(306, 693, "AE613 and beyond (stages 5 to 7)", va="center", fontsize=10, color=INK)
fig.savefig("img/fig-study-plan.png", dpi=150, bbox_inches="tight")
print("wrote img/fig-study-plan.png")