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))
    )

# ==========================================================
# Initial Attitude
# ==========================================================
pitch_deg = 30

pitch = np.deg2rad(pitch_deg)

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

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)

# ==========================================================
# Sample Yaw Angles
# ==========================================================
yaw_samples = [
    0,
    15,
    30,
    45,
    60,
    75,
    90,
    105,
    120,
    135,
    150,
    165,
    180
]

earth_points = []
body_points = []

for yaw_deg in yaw_samples:

    yaw = np.deg2rad(yaw_deg)

    earth_points.append(
        rotate(
            nose0,
            earth_z,
            -yaw
        )
    )

    body_points.append(
        rotate(
            nose0,
            body_z,
            -yaw
        )
    )

earth_points = np.array(earth_points)
body_points = np.array(body_points)

# ==========================================================
# Pitch Flip Prediction
# ==========================================================

pitch_flip_points = earth_points.copy()

pitch_flip_points[:,2] *= -1

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

# ==========================================================
# 3D View
# ==========================================================
ax3d = plt.subplot2grid(
    (3,2),
    (0,0),
    rowspan=2,
    colspan=2,
    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)

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
# ----------------------------------------------------------
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',
    alpha=0.4
)

# ----------------------------------------------------------
# Earth Final Points
# ----------------------------------------------------------
ax3d.plot(
    earth_points[:,0],
    earth_points[:,1],
    earth_points[:,2],
    color='blue',
    linewidth=2,
    label='Earth Yaw Final'
)

# ----------------------------------------------------------
# Body Final Points
# ----------------------------------------------------------
ax3d.plot(
    body_points[:,0],
    body_points[:,1],
    body_points[:,2],
    color='red',
    linewidth=2,
    label='Body Yaw Final'
)

# ----------------------------------------------------------
# Pitch Flip Prediction
# ----------------------------------------------------------
ax3d.plot(
    pitch_flip_points[:,0],
    pitch_flip_points[:,1],
    pitch_flip_points[:,2],

    color='green',
    linewidth=2,

    label='Pitch Flip Prediction'
)

# ----------------------------------------------------------
# Connect Error Lines
# ----------------------------------------------------------
for i in range(len(yaw_samples)):

    ax3d.plot(
        [
            earth_points[i,0],
            body_points[i,0]
        ],
        [
            earth_points[i,1],
            body_points[i,1]
        ],
        [
            earth_points[i,2],
            body_points[i,2]
        ],
        color='gray',
        alpha=0.5
    )

# ----------------------------------------------------------
# Earth -> Pitch Flip
# ----------------------------------------------------------
for i in range(len(yaw_samples)):

    ax3d.plot(
        [
            earth_points[i,0],
            pitch_flip_points[i,0]
        ],
        [
            earth_points[i,1],
            pitch_flip_points[i,1]
        ],
        [
            earth_points[i,2],
            pitch_flip_points[i,2]
        ],

        color='green',
        linestyle='--',
        alpha=0.4
    )

# ----------------------------------------------------------
# Labels
# ----------------------------------------------------------
important = [
    45,
    90,
    135,
    180
]

for yaw_deg in important:

    idx = yaw_samples.index(yaw_deg)

    ax3d.text(
        earth_points[idx,0],
        earth_points[idx,1],
        earth_points[idx,2]+0.05,
        f"{yaw_deg}°",
        color='blue'
    )

    ax3d.text(
        body_points[idx,0],
        body_points[idx,1],
        body_points[idx,2]-0.05,
        f"{yaw_deg}°",
        color='red'
    )
    
    ax3d.text(
        pitch_flip_points[idx,0],
        pitch_flip_points[idx,1],
        pitch_flip_points[idx,2],
        f"P{yaw_deg}",
        color='green'
    )

# ----------------------------------------------------------
# Start Point
# ----------------------------------------------------------
ax3d.scatter(
    nose0[0],
    nose0[1],
    nose0[2],
    color='black',
    s=80
)

ax3d.text(
    nose0[0],
    nose0[1],
    nose0[2]+0.05,
    "Start"
)

# ----------------------------------------------------------
# Axis
# ----------------------------------------------------------
ax3d.set_title(
    "Pitch Flip Failure Curve"
)

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_points[:,0],
    earth_points[:,1],
    '-o',
    color='blue'
)

ax_top.plot(
    body_points[:,0],
    body_points[:,1],
    '-o',
    color='red'
)

ax_top.plot(
    pitch_flip_points[:,0],
    pitch_flip_points[:,1],
    '-o',
    color='green'
)

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_points[:,0],
    earth_points[:,2],
    '-o',
    color='blue'
)

ax_side.plot(
    body_points[:,0],
    body_points[:,2],
    '-o',
    color='red'
)

ax_side.plot(
    pitch_flip_points[:,0],
    pitch_flip_points[:,2],
    '-o',
    color='green'
)

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

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

ax_side.grid(True)

ax_side.set_title(
    "Side View"
)

plt.tight_layout()
plt.show()