import numpy as np
import matplotlib.pyplot as plt

# ==========================================================
# Rodrigues Rotation
# ==========================================================
def rotate(v, axis, angle):
    axis = axis / np.linalg.norm(axis)

    return (
        v * np.cos(angle)
        + np.cross(axis, v) * np.sin(angle)
        + axis * np.dot(axis, v) * (1 - np.cos(angle))
    )

# ==========================================================
# Parameters
# ==========================================================

pitch_deg = 30

# ----------------------------------------------------------
# Change this value
#
# 45
# 90
# 135
# 180
#
# to observe different yaw angles
# ----------------------------------------------------------

yaw_deg = 45

# ==========================================================
# Initial Attitude
# ==========================================================

pitch = np.deg2rad(pitch_deg)

nose0 = np.array([
    np.cos(pitch),
    0,
    np.sin(pitch)
])

tail0 = -nose0

earth_z = np.array([0, 0, 1])

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

body_z /= np.linalg.norm(body_z)

# ==========================================================
# Betaflight Convention
# Yaw+ = nose turns right
# ==========================================================
angles = np.linspace(
    0,
    -np.deg2rad(yaw_deg),
    120
)

earth_nose = []
body_nose = []

earth_tail = []
body_tail = []

for a in angles:

    earth_nose.append(
        rotate(nose0, earth_z, a)
    )

    body_nose.append(
        rotate(nose0, body_z, a)
    )

    earth_tail.append(
        rotate(tail0, earth_z, a)
    )

    body_tail.append(
        rotate(tail0, body_z, a)
    )

earth_nose = np.array(earth_nose)
body_nose = np.array(body_nose)

earth_tail = np.array(earth_tail)
body_tail = np.array(body_tail)

earth_final = earth_nose[-1]
body_final = body_nose[-1]

# ==========================================================
# Figure
# ==========================================================
fig = plt.figure(
    figsize=(14, 12)
)

# ==========================================================
# Large 3D View
# ==========================================================
ax3d = plt.subplot2grid(
    (3, 2),
    (0, 0),
    rowspan=2,
    colspan=2,
    projection='3d'
)

# ----------------------------------------------------------
# z = 0 plane
# ----------------------------------------------------------
xx, yy = np.meshgrid(
    np.linspace(-1.2, 1.2, 10),
    np.linspace(-1.2, 1.2, 10)
)

zz = np.zeros_like(xx)

ax3d.plot_surface(
    xx,
    yy,
    zz,
    color='lightgray',
    alpha=0.25,
    edgecolor='none'
)

# ----------------------------------------------------------
# Earth Z
# ----------------------------------------------------------
ax3d.plot(
    [0, 0],
    [0, 0],
    [0, 1.2],
    color='blue',
    linewidth=2,
    label='Earth Z'
)

# ----------------------------------------------------------
# Body Z'
# ----------------------------------------------------------
ax3d.plot(
    [0, body_z[0]],
    [0, body_z[1]],
    [0, body_z[2]],
    '--',
    color='red',
    linewidth=2,
    label="Body Z'"
)

# ----------------------------------------------------------
# Reference Circle
# Earth-Yaw should stay here
# ----------------------------------------------------------

theta = np.linspace(
    0,
    2*np.pi,
    300
)

r = np.cos(pitch)

ref_x = r * np.cos(theta)
ref_y = r * np.sin(theta)
ref_z = np.ones_like(theta) * np.sin(pitch)

ax3d.plot(
    ref_x,
    ref_y,
    ref_z,

    '--',
    color='black',
    linewidth=1.5,
    alpha=0.6,

    label='Reference Circle'
)

# ----------------------------------------------------------
# Nose Paths
# ----------------------------------------------------------
# ----------------------------------------------------------
# Earth Path
# ----------------------------------------------------------
ax3d.plot(
    earth_nose[:,0],
    earth_nose[:,1],
    earth_nose[:,2],
    color='blue',
    linewidth=3,
    label='Earth Yaw Path'
)

# ----------------------------------------------------------
# Body Path
# Above z=0 : red
# Below z=0 : orange
# ----------------------------------------------------------

above = body_nose[:,2] >= 0
below = body_nose[:,2] < 0

ax3d.plot(
    body_nose[above,0],
    body_nose[above,1],
    body_nose[above,2],
    color='red',
    linewidth=3,
    label='Body Yaw (z>=0)'
)

ax3d.plot(
    body_nose[below,0],
    body_nose[below,1],
    body_nose[below,2],
    color='orange',
    linewidth=3,
    label='Body Yaw (z<0)'
)

# ----------------------------------------------------------
# Tail Paths
# ----------------------------------------------------------
ax3d.plot(
    earth_tail[:,0],
    earth_tail[:,1],
    earth_tail[:,2],
    color='blue',
    linewidth=1,
    alpha=0.5
)

ax3d.plot(
    body_tail[:,0],
    body_tail[:,1],
    body_tail[:,2],
    color='red',
    linewidth=1,
    alpha=0.5
)

# ----------------------------------------------------------
# Pose History
# ----------------------------------------------------------
history_idx = np.linspace(
    0,
    len(angles)-1,
    7
).astype(int)

for i in history_idx:

    # Earth history
    ax3d.plot(
        [earth_tail[i,0], earth_nose[i,0]],
        [earth_tail[i,1], earth_nose[i,1]],
        [earth_tail[i,2], earth_nose[i,2]],
        color='blue',
        alpha=0.15,
        linewidth=2
    )

    # Body history
    ax3d.plot(
        [body_tail[i,0], body_nose[i,0]],
        [body_tail[i,1], body_nose[i,1]],
        [body_tail[i,2], body_nose[i,2]],
        color='red',
        alpha=0.15,
        linewidth=2
    )

# ----------------------------------------------------------
# Initial Aircraft
# ----------------------------------------------------------
ax3d.plot(
    [tail0[0], nose0[0]],
    [tail0[1], nose0[1]],
    [tail0[2], nose0[2]],
    color='black',
    linewidth=3,
    linestyle='--',
    label='Initial Aircraft'
)

# ----------------------------------------------------------
# Final Aircraft
# ----------------------------------------------------------
earth_tail_final = earth_tail[-1]
body_tail_final = body_tail[-1]

ax3d.plot(
    [earth_tail_final[0], earth_final[0]],
    [earth_tail_final[1], earth_final[1]],
    [earth_tail_final[2], earth_final[2]],
    color='gray',
    linewidth=1,
    linestyle='--',
    alpha=0.5,
    label='Final Aircraft'
)

ax3d.plot(
    [body_tail_final[0], body_final[0]],
    [body_tail_final[1], body_final[1]],
    [body_tail_final[2], body_final[2]],
    color='gray',
    linewidth=1,
    linestyle='--',
    alpha=0.5,
    label='Body Final'
)

# ----------------------------------------------------------
# Direction Arrows
# ----------------------------------------------------------
mid = 60

offset = 0.08

ax3d.quiver(
    earth_nose[mid,0],
    earth_nose[mid,1],
    earth_nose[mid,2] + offset,

    earth_nose[mid+2,0] - earth_nose[mid,0],
    earth_nose[mid+2,1] - earth_nose[mid,1],
    earth_nose[mid+2,2] - earth_nose[mid,2],

    color='blue',
    length=0.25,
    normalize=True
)

ax3d.quiver(
    body_nose[mid,0],
    body_nose[mid,1],
    body_nose[mid,2] - offset,

    body_nose[mid+2,0] - body_nose[mid,0],
    body_nose[mid+2,1] - body_nose[mid,1],
    body_nose[mid+2,2] - body_nose[mid,2],

    color='red',
    length=0.25,
    normalize=True
)

# ----------------------------------------------------------
# Markers
# ----------------------------------------------------------
ax3d.scatter(
    nose0[0],
    nose0[1],
    nose0[2],
    color='black',
    s=60
)

ax3d.scatter(
    earth_final[0],
    earth_final[1],
    earth_final[2],
    color='blue',
    s=60
)

ax3d.scatter(
    body_final[0],
    body_final[1],
    body_final[2],
    color='red',
    s=60
)

ax3d.set_title(
    "3D View (Main Teaching View)"
)

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

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

ax3d.legend()

# ==========================================================
# Top View
# ==========================================================
ax_top = plt.subplot2grid(
    (3,2),
    (2,0)
)

ax_top.plot(
    earth_nose[:,0],
    earth_nose[:,1],
    color='blue',
    linewidth=3
)

ax_top.plot(
    body_nose[above,0],
    body_nose[above,1],
    color='red',
    linewidth=3
)

ax_top.plot(
    body_nose[below,0],
    body_nose[below,1],
    color='orange',
    linewidth=3
)

ax_top.grid(True)
ax_top.set_aspect('equal')
ax_top.set_title("Top View")

# ==========================================================
# Side View
# ==========================================================
ax_side = plt.subplot2grid(
    (3,2),
    (2,1)
)

ax_side.plot(
    earth_nose[:,0],
    earth_nose[:,2],
    color='blue',
    linewidth=3
)

ax_side.plot(
    body_nose[above,0],
    body_nose[above,2],
    color='red',
    linewidth=3
)

ax_side.plot(
    body_nose[below,0],
    body_nose[below,2],
    color='orange',
    linewidth=3
)

ax_side.axhline(
    0,
    color='gray',
    linestyle='--'
)

ax_side.axhline(
    np.sin(pitch),
    color='black',
    linestyle='--',
    linewidth=1.5,

    label='Reference Height'
)

ax_side.grid(True)
ax_side.set_title("Side View")

fig.suptitle(
    "Pitch Flip Demonstration (Yaw = 180°)",
    fontsize=16
)

plt.tight_layout()
plt.show()