import numpy as np
import matplotlib.pyplot as plt

# ==========================================================
# Script03
#
# Earth Z Projection onto Body X
#
# Goal:
# Explain why:
#
# Projection
#
# =
#
# sin(Pitch)
#
# (Geometry only. No math proof here.)
# ==========================================================

# ==========================================================
# Aircraft Attitude
# ==========================================================

roll_deg = 30
pitch_deg = 45
yaw_deg = 20

roll = np.deg2rad(roll_deg)
pitch = np.deg2rad(pitch_deg)
yaw = np.deg2rad(yaw_deg)

# ==========================================================
# Earth Coordinate
# ==========================================================

earth_x = np.array([1.0, 0.0, 0.0])
earth_y = np.array([0.0, 1.0, 0.0])
earth_z = np.array([0.0, 0.0, 1.0])

# ==========================================================
# Rotation Matrix
#
# Positive Pitch
#
# =
#
# Nose Up
# ==========================================================

Rx = np.array([
    [1, 0, 0],
    [0, np.cos(roll), -np.sin(roll)],
    [0, np.sin(roll),  np.cos(roll)]
])

Ry = np.array([
    [ np.cos(pitch), 0, -np.sin(pitch)],
    [0,              1, 0],
    [ np.sin(pitch), 0,  np.cos(pitch)]
])

Rz = np.array([
    [ np.cos(yaw), -np.sin(yaw), 0],
    [ np.sin(yaw),  np.cos(yaw), 0],
    [0,             0,           1]
])

R = Rz @ Ry @ Rx

# ==========================================================
# Body Coordinate
# ==========================================================

body_x = R @ earth_x
body_y = R @ earth_y
body_z = R @ earth_z

# ==========================================================
# Aircraft
# ==========================================================

nose = body_x
tail = -body_x

# ==========================================================
# Figure Layout
# ==========================================================

fig = plt.figure(
    figsize=(14,7)
)

# ==========================================================
# Reality View
# ==========================================================

ax_real = plt.subplot2grid(
    (1,2),
    (0,0),
    projection='3d'
)

# ==========================================================
# Body Reference View
# ==========================================================

ax_body = plt.subplot2grid(
    (1,2),
    (0,1),
    projection='3d'
)

# ==========================================================
# Ground Plane
# ==========================================================

xx, yy = np.meshgrid(
    np.linspace(-1.2,1.2,10),
    np.linspace(-1.2,1.2,10)
)

zz = np.zeros_like(xx)

for ax in [ax_real, ax_body]:

    ax.plot_surface(
        xx,
        yy,
        zz,

        color='lightgray',
        alpha=0.12,
        edgecolor='none'
    )

# ==========================================================
# Axis Drawing Helper
# ==========================================================

def draw_axis(
        ax,
        vec,
        color,
        label,

        lw=2,
        alpha=1.0,
        ls='-'
):

    ax.plot(
        [0, vec[0]],
        [0, vec[1]],
        [0, vec[2]],

        color=color,
        linewidth=lw,
        alpha=alpha,
        linestyle=ls
    )

    ax.text(
        vec[0]*1.06,
        vec[1]*1.06,
        vec[2]*1.06,

        label,

        fontsize=9,

        color=color,
        alpha=alpha
    )
    
# ==========================================================
# Reality View
# ==========================================================

# ----------------------------------------------------------
# Earth Coordinate (Background)
# ----------------------------------------------------------

draw_axis(
    ax_real,
    earth_x,
    "red",
    "Earth X",

    lw=1.5,
    ls="--"
)

draw_axis(
    ax_real,
    earth_y,
    "green",
    "Earth Y",

    lw=1.5,
    ls="--"
)

draw_axis(
    ax_real,
    earth_z,
    "blue",
    "Earth Z",

    lw=1.5,
    ls="--"
)

# ----------------------------------------------------------
# Body Coordinate
# ----------------------------------------------------------

draw_axis(
    ax_real,
    body_x,
    "cyan",
    "Body X",

    lw=2.5
)

draw_axis(
    ax_real,
    body_y,
    "gray",
    "Body Y",

    lw=2.5
)

draw_axis(
    ax_real,
    body_z,
    "purple",
    "Body Z",

    lw=2.5
)

# ==========================================================
# Projection of Body X onto Earth XY Plane
# ==========================================================

proj_exy = np.array([
    body_x[0],
    body_x[1],
    0.0
])

# Projection Vector

ax_real.plot(
    [0, proj_exy[0]],
    [0, proj_exy[1]],
    [0, proj_exy[2]],
    color="orange",
    linewidth=2
)

ax_real.scatter(
    proj_exy[0],
    proj_exy[1],
    proj_exy[2],
    color="orange",
    s=40
)

ax_real.text(
    proj_exy[0],
    proj_exy[1],
    proj_exy[2]-0.05,
    "ProjEXY",
    color="orange",
    fontsize=9
)

# Vertical Line

ax_real.plot(
    [body_x[0], proj_exy[0]],
    [body_x[1], proj_exy[1]],
    [body_x[2], proj_exy[2]],
    "--",
    color="gray",
    linewidth=1.5
)

# ==========================================================
# Pitch Arc (Body X -> ProjEXY)
# ==========================================================

pitch_arc_radius = 0.25

# Body X 单位向量
v1 = body_x / np.linalg.norm(body_x)

# ProjEXY 单位向量
v2 = proj_exy / np.linalg.norm(proj_exy)

# 圆弧采样
theta = np.linspace(
    0,
    pitch,
    80
)

arc = []

for t in theta:

    vec = (
        np.sin(pitch - t) * v2 +
        np.sin(t) * v1
    ) / np.sin(pitch)

    vec = vec / np.linalg.norm(vec)

    arc.append(
        pitch_arc_radius * vec
    )

arc = np.array(arc)

ax_real.plot(
    arc[:,0],
    arc[:,1],
    arc[:,2],
    color="black",
    linewidth=1.5
)

mid = arc[len(arc)//2]

ax_real.text(
    mid[0],
    mid[1],
    mid[2]+0.03,
    f"{pitch_deg:.0f}°",
    fontsize=10
)

ax_real.set_title(
    "Reality View"
)

# ==========================================================
# Body Reference View
# ==========================================================

# ----------------------------------------------------------
# Body Coordinate (Fixed)
# ----------------------------------------------------------

draw_axis(
    ax_body,
    np.array([1,0,0]),

    "cyan",
    "Body X",

    lw=2.5
)

draw_axis(
    ax_body,
    np.array([0,1,0]),

    "gray",
    "Body Y",

    lw=2.5
)

draw_axis(
    ax_body,
    np.array([0,0,1]),

    "purple",
    "Body Z",

    lw=2.5
)

# ----------------------------------------------------------
# Earth Coordinate
#
# Body Frame
#
# Earth rotates backwards
# ----------------------------------------------------------

earth_x_body = R.T @ earth_x
earth_y_body = R.T @ earth_y
earth_z_body = R.T @ earth_z

draw_axis(
    ax_body,
    earth_x_body,

    "gray",
    "Earth X",

    lw=1,
    alpha=0.40,
    ls="--"
)

draw_axis(
    ax_body,
    earth_y_body,

    "gray",
    "Earth Y",

    lw=1,
    alpha=0.40,
    ls="--"
)

draw_axis(
    ax_body,
    earth_z_body,

    "blue",
    "EZ",

    lw=1.5,
    ls="--"
)

# ==========================================================
# Projection
#
# Earth Z
#
# projected onto
#
# Body X
# ==========================================================

projection = np.dot(
    earth_z_body,
    np.array([1,0,0])
)

proj_point = np.array([
    projection,
    0,
    0
])

# ----------------------------------------------------------
# Pitch Angle (EZ <-> ProjYZ)
# ----------------------------------------------------------

pitch_arc_radius = 0.32

# EZ方向单位向量
v1 = earth_z_body / np.linalg.norm(earth_z_body)

# ProjYZ方向单位向量
v2 = proj_yz / np.linalg.norm(proj_yz)

theta = np.linspace(
    0,
    1,
    80
)

arc = []

for t in theta:

    vec = (
        (1-t)*v2 +
        t*v1
    )

    vec = vec / np.linalg.norm(vec)

    arc.append(
        pitch_arc_radius * vec
    )

arc = np.array(arc)

ax_body.plot(
    arc[:,0],
    arc[:,1],
    arc[:,2],
    color="black",
    linewidth=1.5
)

mid = arc[len(arc)//2]

ax_body.text(
    mid[0],
    mid[1],
    mid[2]+0.03,
    f"{pitch_deg:.0f}°",
    fontsize=10
)

# ----------------------------------------------------------
# Earth Z Projection on Body YZ Plane
# ----------------------------------------------------------

earth_z_body = R.T @ np.array([0.0, 0.0, 1.0])

proj_yz = np.array([
    0.0,
    earth_z_body[1],
    earth_z_body[2]
])

# ============================
# Body X End Point
# ============================

body_x_point = np.array([
    projection,
    0,
    0
])

ax_body.scatter(
    body_x_point[0],
    body_x_point[1],
    body_x_point[2],
    color="gray",
    s=40
)

ax_body.text(
    body_x_point[0],
    body_x_point[1],
    body_x_point[2]-0.05,
    "BodyXPt",
    color="black",
    fontsize=9
)

# ==========================================================
# 绘制 Earth Z 到 Body YZ 的垂线
# ==========================================================
ax_body.scatter(
    earth_z_body[0],
    earth_z_body[1],
    earth_z_body[2],
    color="orange",
    s=40
)

ax_body.plot(
    [earth_z_body[0], 0],
    [earth_z_body[1], earth_z_body[1]],
    [earth_z_body[2], earth_z_body[2]],
    color="gray",
    linestyle="--",
    linewidth=1.5
)

mid = (
    earth_z_body +
    np.array([0, earth_z_body[1], earth_z_body[2]])
) / 2

ax_body.text(
    mid[0],
    mid[1],
    mid[2] + 0.03,
    "ProjLineYZ",
    color="gray",
    fontsize=8
)

# Projection vector
ax_body.plot(
    [0, proj_yz[0]],
    [0, proj_yz[1]],
    [0, proj_yz[2]],
    color="orange",
    linewidth=2
)

# Projection point
ax_body.scatter(
    proj_yz[0],
    proj_yz[1],
    proj_yz[2],
    color="orange",
    s=40
)

ax_body.text(
    proj_yz[0],
    proj_yz[1],
    proj_yz[2] + 0.05,
    "ProjYZ",
    color="orange",
    fontsize=9
)

# ----------------------------------------------------------
# Top Edge of Parallelogram
# ----------------------------------------------------------

ax_body.plot(
    [earth_z_body[0], body_x_point[0]],
    [earth_z_body[1], body_x_point[1]],
    [earth_z_body[2], body_x_point[2]],
    "--",
    color="gray",
    linewidth=1.3
)

ax_body.set_title(
    "Body Reference View"
)

# ==========================================================
# 3D Axis
# ==========================================================

for ax in [ax_real, ax_body]:

    ax.set_xlim(-1.2,1.2)
    ax.set_ylim(-1.2,1.2)
    ax.set_zlim(-1.2,1.2)

    ax.set_box_aspect([1,1,1])

# ==========================================================
# Better View Angle
# ==========================================================

ax_real.view_init(
    elev=24,
    azim=-55
)

ax_body.view_init(
    elev=18,
    azim=-45
)

# ==========================================================
# Figure Title
# ==========================================================

fig.suptitle(
    "Script03 : Earth Z Projection onto Body X",
    fontsize=16
)

plt.tight_layout()

plt.show()