from typing import Tuple

import numpy as np
from scipy import io
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

from arm3d import Robot3D


def cylinder(R: float = 1.0, N: int = 20):
    R = np.atleast_1d(R).astype(float)
    if R.size == 1:
        # Match MATLAB: 2 rows (z=0 and z=1) when R is scalar
        R = np.array([R.item(), R.item()], dtype=float)

    M = R.size
    theta = np.linspace(0.0, 2.0*np.pi, N+1)  # include 2π to close seam
    z = np.linspace(0.0, 1.0, M)

    X = np.outer(R, np.cos(theta))
    Y = np.outer(R, np.sin(theta))
    Z = np.tile(z[:, None], (1, N+1))

    return X, Y, Z

def plotCylinderWithCaps(r, cnt, height, nSides, color, orientation, ax):
    def fill3(x, y, z, color, ax):
        # Matplotlib lacks equivalent fill3
        # Need to fill with polygons
        vertices = [list(zip(x, y, z))]
        h = Poly3DCollection(vertices, facecolors=color, edgecolors='none')
        ax.add_collection3d(h)
    
    [X,Y,Z] = cylinder(r, nSides)
    X = X + cnt[0]
    Y = Y + cnt[1] 
    Z = Z * height + cnt[2]
    if orientation == 'x':
        ax.plot_surface(Z, Y, X, color=color)
        fill3(Z[0, :], Y[0, :], X[0, :], color, ax) 
        fill3(Z[1, :], Y[1, :], X[1, :], color, ax)
    else:
        ax.plot_surface(X, Y, Z, color=color)
        fill3(X[0, :], Y[0, :], Z[0, :], color, ax) 
        fill3(X[1, :], Y[1, :], Z[1, :], color, ax)
    return

def draw_cuboid(ax, L, W, H, offset, color):
    vert = np.array([
        [0, 0, 0],
        [L, 0, 0],
        [L, W, 0],
        [0, W, 0],
        [0, 0, H],
        [L, 0, H],
        [L, W, H],
        [0, W, H]
    ])
    vert[:, 0] = vert[:, 0] + offset[0]
    vert[:, 1] = vert[:, 1] + offset[1]
    vert[:, 2] = vert[:, 2] + offset[2]

    ax.bar3d(offset[0], offset[1], offset[2], L, W, H, color=color, alpha=0.5)


class ConveyorBeltPickPlace:
    def __init__(
        self,
        figsize: Tuple[int, int] = (5, 5),
        theta0: np.ndarray = np.array([3 * np.pi / 4, np.pi / 2, np.pi])
        ):
        self.robot = Robot3D(figsize=figsize)
        self.ax = self.robot.ax
        self.theta = theta0

        # Environment
        self.n_cans = 9
        can_pos = io.loadmat("can_positions.mat")
        self.start_can_pos, self.end_can_pos = can_pos["start_can_pos"], can_pos["end_can_pos"]

        # Keep track of status
        self.obj_picked = 0
        self.obj_placed = 0
        self.in_hand = 0
        self.deliver = 0
        self.steps = 6

    def plot_cans(self):
        nSides = 12           # number of "sides" of the cyl
        color = [1, 0, 0]     # color of each cyl
        r = 1/10
        height = 0.2

        n_objs = self.start_can_pos.shape[0]
        for i in range(n_objs):
            if i >= self.obj_picked:
                print(i)
                plotCylinderWithCaps(r, self.start_can_pos[i], height, nSides, color, "z", self.ax)
            else:
                plotCylinderWithCaps(r, self.end_can_pos[i], height, nSides, color, "z", self.ax)

    def plot_static(self):
        brown = (0.588235, 0.294118, 0.0)      # (150, 75, 0) / 255
        grey  = (0.501961, 0.501961, 0.501961) # (128,128,128) / 255
        black = (0.0, 0.0, 0.0)

        draw_cuboid(self.ax,1,1,0.2,[-0.1, -1.3, -0.4], brown) # table
        draw_cuboid(self.ax,4,4,0.2,[-2, -2, -1], grey) # floor
        draw_cuboid(self.ax,1,2,0.2,[-1.3, -1.5, 0], black) # conveyor belt
        plotCylinderWithCaps(0.1, [0.1, -1.5, -1.3], 1, 12, [0, 0, 0], 'x', self.ax); # end of belt

    def __call__(
        self, 
        ):
        """Plot the robot.
        """
        self.plot_cans()
        self.plot_static()
        self.robot(self.theta, plot_now=False)
        self.ax.scatter(self.start_can_pos[:, 0], self.start_can_pos[:, 1], self.start_can_pos[:, 2])
        self.ax.scatter(self.end_can_pos[:, 0], self.end_can_pos[:, 1], self.end_can_pos[:, 2])
        
        plt.show(block=False)