"""Figure for notation.qmd (img/fig-frames.png): the five frames the notes use, in one picture.

Inertial (I, J, K) and Earth-fixed (x_ECEF turned by the Greenwich angle theta_G about K) at the Earth's
centre; perifocal (p, q, w) at the focus of an inclined orbit; the orbit-following frame (R, S, W, also
called LVLH and r, theta, h) on the satellite; and the topocentric frame (S, E, Z) at a ground site.

Run from the sfm-slides directory:   python3 codes/fig-frames.py
"""
import numpy as np
import _figfont                                # all figure text at one scale
import matplotlib.pyplot as plt
from _manstyle import INK, MAROON, BLUE, GREEN, AMBER, PURPLE


def R1(a):
    return np.array([[1, 0, 0], [0, np.cos(a), -np.sin(a)], [0, np.sin(a), np.cos(a)]])


def R3(a):
    return np.array([[np.cos(a), -np.sin(a), 0], [np.sin(a), np.cos(a), 0], [0, 0, 1]])


Om, inc, w, f_sat, thG = np.radians(35), np.radians(40), np.radians(60), np.radians(-70), np.radians(70)
a_, e_ = 2.0, 0.25
p_ = a_ * (1 - e_ ** 2)
Q = R3(Om) @ R1(inc) @ R3(w)
fig = plt.figure(figsize=(10, 8.4))
ax = fig.add_subplot(111, projection="3d")
u, v = np.meshgrid(np.linspace(0, 2 * np.pi, 40), np.linspace(-np.pi / 2, np.pi / 2, 20))
ax.plot_surface(np.cos(v) * np.cos(u), np.cos(v) * np.sin(u), np.sin(v), color="#cfdcec", alpha=0.35, lw=0)
t = np.linspace(0, 2 * np.pi, 300)
ax.plot(np.cos(t), np.sin(t), 0 * t, color="#9db4d0", lw=0.8)


def triad(o, M, cols, labs, L=0.9, fs=12):
    for k in range(3):
        ax.quiver(*o, *(L * M[:, k]), color=cols[k], lw=2, arrow_length_ratio=0.12)
        ax.text(*(o + 1.12 * L * M[:, k]), labs[k], color=cols[k], fontsize=fs)


O = np.zeros(3)
triad(O, np.eye(3) * 1.0, ["#6f685e"] * 3, [r"$\hat{\mathbf{I}}$", r"$\hat{\mathbf{J}}$", r"$\hat{\mathbf{K}}$"], L=2.1)
E = R3(thG)
ax.quiver(*O, *(1.6 * E[:, 0]), color=AMBER, lw=2, arrow_length_ratio=0.1)
ax.quiver(*O, *(1.6 * E[:, 1]), color=AMBER, lw=2, arrow_length_ratio=0.1)
ax.text(*(1.75 * E[:, 0]), "ECEF $x$", color=AMBER, fontsize=10)
ax.text(*(1.75 * E[:, 1]), "ECEF $y$", color=AMBER, fontsize=10)
arc = np.array([np.cos(s) * np.array([1, 0, 0]) + np.sin(s) * np.array([0, 1, 0]) for s in np.linspace(0, thG, 30)])
ax.plot(*(0.7 * arc).T, color=AMBER, lw=1.4)
ax.text(*(0.75 * arc[15] + np.array([0, 0, -0.15])), r"$\theta_G$", color=AMBER, fontsize=11)
# orbit and perifocal triad
ff = np.linspace(0, 2 * np.pi, 400)
rr = p_ / (1 + e_ * np.cos(ff))
orbit = Q @ np.array([rr * np.cos(ff), rr * np.sin(ff), 0 * ff])
ax.plot(*orbit, color=BLUE, lw=1.8)
triad(O, Q, [MAROON, GREEN, PURPLE], [r"$\hat{\mathbf{p}}$", r"$\hat{\mathbf{q}}$", r"$\hat{\mathbf{w}}$"], L=1.3)
# satellite and RSW
rs = p_ / (1 + e_ * np.cos(f_sat))
S = Q @ np.array([rs * np.cos(f_sat), rs * np.sin(f_sat), 0])
Rh = S / np.linalg.norm(S); Wh = Q[:, 2]; Sh = np.cross(Wh, Rh)
ax.scatter(*S, color=INK, s=30)
triad(S, np.column_stack([Rh, Sh, Wh]), [MAROON, GREEN, PURPLE], [r"$\hat{\mathbf{R}}$", r"$\hat{\mathbf{S}}$", r"$\hat{\mathbf{W}}$"], L=0.7, fs=11)
ax.text(*(S + np.array([0.1, 0.1, -0.35])), "RSW = LVLH", color=INK, fontsize=10)
# topocentric SEZ at a site
lat, lon = np.radians(-35), np.radians(-75)
Zh = np.array([np.cos(lat) * np.cos(lon), np.cos(lat) * np.sin(lon), np.sin(lat)])
Eh = np.array([-np.sin(lon), np.cos(lon), 0])
Shh = np.array([np.sin(lat) * np.cos(lon), np.sin(lat) * np.sin(lon), -np.cos(lat)])
G = Zh * 1.0
ax.scatter(*G, color=INK, s=20)
triad(G, np.column_stack([Shh, Eh, Zh]), ["#8a5a2b"] * 3, [r"$\hat{\mathbf{S}}$", r"$\hat{\mathbf{E}}$", r"$\hat{\mathbf{Z}}$"], L=0.55, fs=11)
ax.text(*(G + np.array([0.45, 0.25, -0.5])), "SEZ at a site", color="#8a5a2b", fontsize=10)
ax.set_xlim(-2.2, 2.2); ax.set_ylim(-2.2, 2.2); ax.set_zlim(-1.8, 2.2)
ax.set_box_aspect((1, 1, 0.95), zoom=1.2)
ax.view_init(20, -50)
ax.set_axis_off()
ax.set_title("The five frames: inertial, Earth-fixed, perifocal, orbit-following (RSW/LVLH) and topocentric", fontsize=11.5)
fig.savefig("img/fig-frames.png", dpi=150, bbox_inches="tight")
print("wrote img/fig-frames.png")