from typing import Tuple

import numpy as np
import matplotlib.pyplot as plt


def cylinder2P(R, N, r1, r2):
    """
    Args:
        R - radius of cylinder 
        N - number of dense points taken on the surface. Ex: 100
        r1 - starting coordinate of cylinder (x1, y1, z1)
        r2 - starting coordinate of cylinder (x2, y2, z2)
    Output: 
        Use this (X,Y,Z) to plot the surface of the cylinder between
        the points r1 and r2.
    """
    theta = np.linspace(0, 2 * np.pi, N)
    # m = len(R)
    # if m == 1:
    R = np.array([R, R])                
    m = 2

    X = np.zeros((m, N))             
    Y = np.zeros((m, N))
    Z = np.zeros((m, N))
    
    v = (r2 - r1) / np.sqrt((r2-r1) @ (r2-r1).T)
    R2 = np.random.rand(1, 3)             
    x2= v - R2 / (R2 @ v.T)
    x2 = x2 / np.sqrt(x2 @ x2.T)    
    x3 = np.cross(v, x2)
    x3 = x3 / np.sqrt(x3 @ x3.T)
    
    r1x, r1y , r1z = r1[0], r1[1], r1[2]
    r2x, r2y , r2z = r2[0], r2[1], r2[2]
    vx, vy, vz = v[0], v[1], v[2]
    x2x, x2y, x2z = x2[0, 0], x2[0, 1], x2[0, 2]
    x3x, x3y, x3z = x3[0, 0], x3[0, 1], x3[0, 2]
    
    time = np.linspace(0, 1, m)
    for j in range(m):
        t = time[j]
        X[j, :] = r1x + (r2x - r1x) * t + R[j] * np.cos(theta) * x2x + R[j] * np.sin(theta) * x3x
        Y[j, :] = r1y + (r2y - r1y) * t + R[j] * np.cos(theta) * x2y + R[j] * np.sin(theta) * x3y
        Z[j, :] = r1z + (r2z - r1z) * t + R[j] * np.cos(theta) * x2z + R[j] * np.sin(theta) * x3z

    return X, Y, Z


def sphere(n_faces=20, radius=1, center=(0, 0, 0)):
    """Generates x, y, z coordinates for a sphere.
    Args:
        n_faces: Number of faces along the azimuth and elevation.
                       Determines the resolution of the sphere.
        radius: Radius of the sphere.
        center: (cx, cy, cz) center of the sphere.
    """
    theta = np.linspace(0, 2 * np.pi, n_faces + 1)
    phi = np.linspace(0, np.pi, n_faces + 1)

    # Create 2D arrays for theta and phi
    theta_grid, phi_grid = np.meshgrid(theta, phi)

    # Convert spherical coordinates to Cartesian coordinates
    x = radius * np.sin(phi_grid) * np.cos(theta_grid) + center[0]
    y = radius * np.sin(phi_grid) * np.sin(theta_grid) + center[1]
    z = radius * np.cos(phi_grid) + center[2]

    return x, y, z

def getSphere(r, offset, ax):
    [Sx, Sy, Sz] = sphere()
    Sx = Sx * r + offset[0]
    Sy = Sy * r + offset[1]
    Sz = Sz * r + offset[2]
    ax.plot_surface(Sx, Sy, Sz, alpha=0.5, color=[0, 0, 1])

def getHemisphere(r, offset, angle, ax):
    def rotate(
        theta: Tuple[float, float, float],
        ):
        # Roll
        Rx = np.array([
            [1, 0, 0, 0],
            [0, np.cos(theta[0]), -np.sin(theta[0]), 0],
            [0, np.sin(theta[0]), np.cos(theta[0]), 0],
            [0, 0, 0, 1]
        ])

        # Pitch
        Ry = np.array([
            [np.cos(theta[1]), 0, np.sin(theta[1]), 0],
            [0, 1, 0, 0],
            [-np.sin(theta[1]), 0, np.cos(theta[1]), 0],
            [0, 0, 0, 1],
        ])

        # Yaw
        Rz = np.array([
            [np.cos(theta[2]), -np.sin(theta[2]), 0, 0],
            [np.sin(theta[2]), np.cos(theta[2]), 0, 0],
            [0, 0, 1, 0],
            [0, 0, 0, 1],
        ])

        return Rz @ Ry @ Rx
    
    def translate(
        t: Tuple[float, float, float],
        ):
        T = np.array([
            [1, 0, 0, t[0]],
            [0, 1, 0, t[1]],
            [0, 0, 1, t[2]],
            [0, 0, 0, 1],
        ])
        return T
    
    # Initial placement
    [x, y, z] = sphere()
    x, y, z = x[:, 5:16], y[:, 5:16], z[:, 5:16]
    Sx = x*r + offset[0] + r * np.cos(angle[0] + angle[1]) * np.cos(angle[2])
    Sy = y*r + offset[1] + r * np.cos(angle[0] + angle[1]) * np.sin(angle[2])
    Sz = z*r + offset[2] + r * np.sin(angle[0] + angle[1])

    # Origin of rotation
    axis = np.array([
        offset[0] + r * np.cos(angle[0] + angle[1]) * np.cos(angle[2]),
        offset[1] + r * np.cos(angle[0] + angle[1]) * np.sin(angle[2]),
        offset[2] + r * np.sin(angle[0] + angle[1]),
    ])

    # Rotation along origin
    R = rotate([0, -(angle[0] + angle[1]), angle[2]])

    # Flatten -> translate -> rotate Y -> rotate Z -> translate back
    P = np.stack([Sx, Sy, Sz, np.ones(Sz.shape)], axis=0).reshape(4, -1)          # (m, n, 3)
    Pf = translate(axis) @ R @ translate(-axis) @ P

    Xf = Pf[0, :].reshape(Sx.shape)
    Yf = Pf[1, :].reshape(Sy.shape)
    Zf = Pf[2, :].reshape(Sz.shape)
    ax.plot_surface(Xf, Yf, Zf, color=[0, 0, 0])

class Robot3D:
    def __init__(
        self,
        ls: np.ndarray = np.array([0.8, 0.7]),      # NOTE: Do not change link length for 3d experiment(s)
        figsize: Tuple[int, int] = (5, 5),
        ):
        self.ls = ls    # Link length

        # Initialize figure once
        self.fig = plt.figure(figsize=figsize)
        self.ax = self.fig.add_subplot(projection="3d")
        self.fig.canvas.mpl_connect('close_event', self._on_close)
        self.terminate = False
    
    def _on_close(self, event):
        """Event handler, should the figure be manually closed.
        """
        self.terminate = True

    def __call__(
        self, 
        theta: np.ndarray, 
        plot_now: bool = True,
        ):
        """Plot the robot.
        """
        # Restrict workspace & figure size
        lim = 2.0      # Small overflow
        self.ax.set_adjustable("datalim")
        self.ax.set_aspect('equal')
        self.ax.set_box_aspect([1, 1, 1])
        self.ax.set(xlim=(-lim, lim), ylim=(-lim, lim), zlim=(-lim, lim))
        self.ax.grid(False)

        l1, l2 = self.ls

        [X1, Y1, Z1] = cylinder2P(0.2,100, np.array([0, 0, -1]), np.array([0, 0, 0]))
        # surf(X1,Y1,Z1,'FaceColor',[0.3010 0.7450 0.9330],'LineStyle','none')
        self.ax.plot_surface(X1, Y1, Z1, color=[0.3010, 0.7450, 0.9330])

        # Define the joint angles as q1, q2 and q3
        q1, q2, q3 = theta
        xp1 = l1 * np.cos(q1)
        xp2 = xp1 + l2 * np.cos(q1 + q2)
        # Make 2D robot above 3D by projecting xp-z plane at different angles
        x1 = xp1 * np.cos(q3);         # x coordinate of the second joint
        y1 = xp1 * np.sin(q3);         # y coordinate of the second joint
        z1 = l1 * np.sin(q1);            # z coordinate of the second joint
        x2 = xp2 * np.cos(q3);         # X coordinate of the end effector
        y2 = xp2 * np.sin(q3);         # y coordinate of the end effector
        z2 = z1 + l2 * np.sin(q1 + q2);   # z coordinate of the end effector

        # A sphere depicts a joint in our simulation so let's draw a sphere at our
        # first joint [0 0 0]
        getSphere(0.2, np.array([0, 0, 0]), self.ax) # Joint 1

        # We have the coordinates of all the joints and the end effector above.
        # Let's start plotting the remainig links and joints

        [X2, Y2, Z2] = cylinder2P(0.1, 100, np.array([0, 0, 0]), np.array([x1, y1, z1]))
        # surf(X2,Y2,Z2,'FaceColor',[0.3010 0.7450 0.9330],'LineStyle','none');
        self.ax.plot_surface(X2, Y2, Z2, color=[0.3010, 0.7450, 0.9330])
        getSphere(0.15, np.array([x1, y1, z1]), self.ax) # Joint 2

        [X3, Y3, Z3] = cylinder2P(0.1, 100, np.array([x1, y1, z1]), np.array([x2, y2, z2])) # link 2
        # surf(X3,Y3,Z3,'FaceColor',[0.3010 0.7450 0.9330],'LineStyle','none');
        self.ax.plot_surface(X3, Y3, Z3, color=[0.3010, 0.7450, 0.9330])

        getHemisphere(0.13, np.array([x2, y2, z2]), theta, self.ax)

        # Non-blocking show
        if plot_now:
            plt.show(block=True)