import numpy as np
import cv2
import plotly.graph_objects as go
from math import cos, sin
from scipy.spatial.transform import Rotation as R


def plot_3D_rotation(rotation_matrix):
    fig = go.Figure()

    # Original axis orientation
    axes_points = np.array([
        [1, 0, 0, 0],
        [0, 1, 0, 0],
        [0, 0, 1, 0]
    ], dtype=np.float64)

    # Plot original axes
    fig.add_trace(go.Scatter3d(
        x=[0, axes_points[0, 0]],
        y=[0, axes_points[1, 0]],
        z=[0, axes_points[2, 0]],
        mode='lines+text',
        line=dict(color='blue', width=6),
        name='Canonical X-axis',
        text=['', 'X axis'],
        textposition='middle center',
    ))

    fig.add_trace(go.Scatter3d(
        x=[0, axes_points[0, 1]],
        y=[0, axes_points[1, 1]],
        z=[0, axes_points[2, 1]],
        mode='lines+text',
        line=dict(color='blue', width=6),
        name='Canonical Z-axis',
        text=['', 'Z axis'],
        textposition='middle center',
    ))

    fig.add_trace(go.Scatter3d(
        x=[0, axes_points[0, 2]],
        y=[0, axes_points[1, 2]],
        z=[0, axes_points[2, 2]],
        mode='lines+text',
        line=dict(color='blue', width=6),
        name='Canonical Y-axis',
        text=['', 'Y axis'],
        textposition='middle center',
    ))

     # Apply rotation
    axes_points = rotation_matrix @ axes_points

    # Plot rotated axes
    fig.add_trace(go.Scatter3d(
        x=[0, axes_points[0, 0]],
        y=[0, axes_points[1, 0]],
        z=[0, axes_points[2, 0]],
        mode='lines+text',
        line=dict(color='red', width=6),
        name='Rotated X\'-axis',
        text=['', 'Rotated X axis'],
        textposition='middle center',
    ))

    fig.add_trace(go.Scatter3d(
        x=[0, axes_points[0, 1]],
        y=[0, axes_points[1, 1]],
        z=[0, axes_points[2, 1]],
        mode='lines+text',
        line=dict(color='red', width=6),
        name='Rotated Z\'-axis',
        text=['', 'Rotated Z axis'],
        textposition='middle center',
    ))

    fig.add_trace(go.Scatter3d(
        x=[0, axes_points[0, 2]],
        y=[0, axes_points[1, 2]],
        z=[0, axes_points[2, 2]],
        mode='lines+text',
        line=dict(color='red', width=6),
        name='Rotated Y\'-axis',
        text=['', 'Rotated Y axis'],
        textposition='middle center',
    ))

    # Retrieve pitch, yaw, roll from rotation matrix
    r = R.from_matrix(rotation_matrix)
    pitch, yaw, roll = r.as_euler('xzy', degrees=True)

    # Set layout
    fig.update_layout(
        scene=dict(
            xaxis=dict(title='X-axis', range=[-1.2, 1.2]),
            yaxis=dict(title='Z-axis', range=[-1.2, 1.2]),
            zaxis=dict(title='Y-axis', range=[-1.2, 1.2]),
            xaxis_tickvals=np.arange(-1.2, 1.2, 0.6),
            yaxis_tickvals=np.arange(-1.2, 1.2, 0.5),
            zaxis_tickvals=np.arange(-1.2, 1.2, 0.5),
            aspectmode='cube',
            aspectratio=dict(x=1, y=1, z=1),
        ),
        margin=dict(l=0, r=0, t=0, b=30),
    )
    # add annotation
    fig.add_annotation(dict(font=dict(color='black',size=15),
        x=-30,
        y=50,
        showarrow=False,
        text=f"Pitch: {int(pitch)} - Yaw: {int(yaw)} - Roll: {int(roll)}",
        textangle=0,
        xanchor='left',
        xref="paper",
        yref="paper"))
    return fig


def draw_2D_axes(img, roll, pitch, yaw, tdx=None, tdy=None, size=150.):
    # Input is a cv2 image
    # pose_params: (pitch, yaw, roll, tdx, tdy)
    # Where (tdx, tdy) is the translation of the face.
    # For pose we have [pitch yaw roll tdx tdy tdz scale_factor]

    p = pitch * np.pi / 180
    y = (yaw * np.pi / 180)
    r = -roll * np.pi / 180
    if tdx != None and tdy != None:
        face_x = tdx - 0.50 * size 
        face_y = tdy - 0.50 * size

    else:
        height, width = img.shape[:2]
        face_x = width / 2 - 0.5 * size
        face_y = height / 2 - 0.5 * size

    x1 = size * (cos(y) * cos(r)) + face_x
    y1 = size * (cos(p) * sin(r) + cos(r) * sin(p) * sin(y)) + face_y 
    x2 = size * (-cos(y) * sin(r)) + face_x
    y2 = size * (cos(p) * cos(r) - sin(p) * sin(y) * sin(r)) + face_y
    x3 = size * (sin(y)) + face_x
    y3 = size * (-cos(y) * sin(p)) + face_y

    # Draw base in red
    cv2.line(img, (int(face_x), int(face_y)), (int(x1),int(y1)),(0,0,255),3)
    cv2.line(img, (int(face_x), int(face_y)), (int(x2),int(y2)),(0,0,255),3)
    cv2.line(img, (int(x2), int(y2)), (int(x2+x1-face_x),int(y2+y1-face_y)),(0,0,255),3)
    cv2.line(img, (int(x1), int(y1)), (int(x1+x2-face_x),int(y1+y2-face_y)),(0,0,255),3)
    # Draw pillars in blue
    cv2.line(img, (int(face_x), int(face_y)), (int(x3),int(y3)),(255,0,0),2)
    cv2.line(img, (int(x1), int(y1)), (int(x1+x3-face_x),int(y1+y3-face_y)),(255,0,0),2)
    cv2.line(img, (int(x2), int(y2)), (int(x2+x3-face_x),int(y2+y3-face_y)),(255,0,0),2)
    cv2.line(img, (int(x2+x1-face_x),int(y2+y1-face_y)), (int(x3+x1+x2-2*face_x),int(y3+y2+y1-2*face_y)),(255,0,0),2)
    # Draw top in green
    cv2.line(img, (int(x3+x1-face_x),int(y3+y1-face_y)), (int(x3+x1+x2-2*face_x),int(y3+y2+y1-2*face_y)),(0,255,0),2)
    cv2.line(img, (int(x2+x3-face_x),int(y2+y3-face_y)), (int(x3+x1+x2-2*face_x),int(y3+y2+y1-2*face_y)),(0,255,0),2)
    cv2.line(img, (int(x3), int(y3)), (int(x3+x1-face_x),int(y3+y1-face_y)),(0,255,0),2)
    cv2.line(img, (int(x3), int(y3)), (int(x3+x2-face_x),int(y3+y2-face_y)),(0,255,0),2)

    return img