from typing import Union, NewType, Optional, List, Dict, Final

import ipywidgets as widgets
import k3d
import numpy as np
from k3d.objects import Group, Line, Points, Text, Drawable
from k3d.plot import Plot

from .jointset import JointSet, MuPoTSJoints, OpenPoseJoints, CocoExJoints, PanopticJoints, Common14Joints

SkeletonPlot = NewType('SkeletonPlot', Optional[Plot])
Color = NewType('Color', Union[str, int])

DEFAULT_COLORS: Final[Dict[str, int]] = {
    'red': 0xFF0000, 'green': 0x00FF00,  'blue': 0x0000FF,
    'yellow': 0xFFFF00, 'teal': 0x00FFFF, 'purple': 0xFF00FF,
    'white': 0xFFFFFF, 'black': 0x000000
}


class DrawableSkeleton(Group):
    def __init__(self, joint_points: Points, joint_lines: List[Line], joint_names: List[Text]):
        drawable_objects: List[Drawable] = [joint_points] + joint_lines + joint_names
        super().__init__(drawable_objects)
        self.joint_points: Points = joint_points
        self.joint_lines: List[Line] = joint_lines
        self.joint_names: List[Text] = joint_names

    def hide_joint_names(self):
        for joint_name in self.joint_names:
            joint_name.visible = False

    def show_joint_names(self):
        for joint_name in self.joint_names:
            joint_name.visible = True


class Skeleton:
    def __init__(self, joint_coordinates: np.ndarray,
                 joint_set: JointSet, part_size: float, color: Color):
        self.joint_coordinates: np.ndarray = joint_coordinates
        self.joint_set: JointSet = joint_set
        self.part_size: float = part_size
        self.color: Color = color

    def to_drawable_skeleton(self) -> DrawableSkeleton:
        joint_points = self.get_joint_points()
        joint_lines: List[Drawable] = self.get_joint_lines()
        joint_names: List[Drawable] = self.get_joint_names()
        drawable_skeleton: DrawableSkeleton = DrawableSkeleton(joint_points, joint_lines, joint_names)
        return drawable_skeleton

    def get_joint_colors(self) -> np.ndarray:
        joint_colors_shape = (self.joint_set.number_of_joints, )
        if self.color == 'default':
            joint_colors = np.zeros(shape=joint_colors_shape, dtype='uint32')
            joint_colors[self.joint_set.left_joint_indices] = 0xFF0000
            joint_colors[self.joint_set.right_joint_indices] = 0x0000FF
            joint_colors[self.joint_set.center_joint_indices] = 0xFFFFFF
            return joint_colors
        else:
            return np.full(shape=joint_colors_shape,
                           fill_value=self.color, dtype='uint32')

    def get_line_colors(self) -> np.ndarray:
        line_colors_shape = (self.joint_set.number_of_joints - 1, )
        if self.color == 'default':
            line_colors = np.zeros(shape=line_colors_shape, dtype='uint32')
            line_colors[self.joint_set.left_line_indices] = 0xFF0000
            line_colors[self.joint_set.right_line_indices] = 0x0000FF
            line_colors[self.joint_set.center_line_indices] = 0xFFFFFF
            return line_colors
        else:
            return np.full(shape=line_colors_shape,
                           fill_value=self.color, dtype='uint32')

    def get_joint_points(self) -> Points:
        skeleton_joint_colors = self.get_joint_colors()
        return k3d.points(
            positions=self.joint_coordinates, point_size=self.part_size,
            shader='mesh', colors=skeleton_joint_colors
        )

    def get_joint_lines(self) -> List[Line]:
        skeleton_line_colors = self.get_line_colors()
        return [self.__create_line_between_joints(
            start=self.joint_coordinates[self.joint_set.limb_graph[i][0]],
            end=self.joint_coordinates[self.joint_set.limb_graph[i][1]],
            width=self.part_size / 4.0,
            color=int(skeleton_line_colors[i]))
            for i in range(len(self.joint_set.limb_graph))
        ]

    def get_joint_names(self) -> List[Text]:
        return [k3d.text(
            text=self.joint_set.names[i],
            position=self.joint_coordinates[i],
            size=0.45, label_box=False, color=DEFAULT_COLORS.get('black'))
            for i in range(self.joint_set.number_of_joints)
        ]

    @staticmethod
    def __create_line_between_joints(start: float, end: float, width: float, color: int) -> Line:
        return k3d.line(
            vertices=[start, end], shader='mesh',
            width=width, color=color
        )


class SkeletonVisualizer:
    def __init__(self, joint_set: JointSet, size_scalar: float = 1.0):
        self.joint_set: JointSet = joint_set
        self.plot: SkeletonPlot = None
        self.size_scalar: float = size_scalar
        self.skeletons: List[DrawableSkeleton] = []

    def visualize(self, skeletons: np.ndarray, colors: List[Color] = None):
        self.__assert_visualization_arguments(skeletons, colors)
        number_of_skeletons = skeletons.shape[0]
        if colors is None:
            colors = ['default' for i in range(number_of_skeletons)]
        self.plot = self.__create_skeleton_plot(skeletons, colors)
        for skeleton in self.skeletons:
            skeleton.hide_joint_names()
        self.plot.display()

    def __assert_visualization_arguments(self, skeletons: np.ndarray, colors: List[Color]):
        assert len(skeletons.shape) == 3, 'The \'skeletons\' parameter should be a 3 dimensional numpy array.'
        if colors is not None:
            assert skeletons.shape[0] == len(colors),\
                'The \'skeletons\' and \'colors\' parameters must be the same length.'
        assert skeletons.shape[1] == self.joint_set.number_of_joints,\
            'The number of joints of skeletons and the number of joints in the specified joint set must be the same.'
        assert skeletons.shape[2] == 3, 'The skeleton joint coordinates must be 3 dimensional'

    def __create_skeleton_plot(self, skeletons: np.ndarray, colors: List[Color]) -> SkeletonPlot:
        skeleton_part_size = self.__calculate_skeleton_part_size(skeletons)
        skeleton_plot = k3d.plot()
        for skeleton in map(
                lambda skeleton_color_tuple: Skeleton(
                    joint_coordinates=skeleton_color_tuple[0],
                    joint_set=self.joint_set,
                    part_size=skeleton_part_size,
                    color=skeleton_color_tuple[1]),
                zip(skeletons, colors)):
            drawable_skeleton = skeleton.to_drawable_skeleton()
            skeleton_plot += drawable_skeleton
            self.skeletons.append(drawable_skeleton)
        return skeleton_plot

    def __calculate_skeleton_part_size(self, skeletons: np.ndarray) -> float:
        max_values = [abs(skeleton).max() for skeleton in skeletons]
        return (min(max_values) / 50.0) * self.size_scalar


class MuPoTSVisualizer(SkeletonVisualizer):
    def __init__(self, size_scalar: float = 1.0):
        super().__init__(joint_set=MuPoTSJoints(), size_scalar=size_scalar)


class OpenPoseVisualizer(SkeletonVisualizer):
    def __init__(self, size_scalar: float = 1.0):
        super().__init__(joint_set=OpenPoseJoints(), size_scalar=size_scalar)


class CocoExVisualizer(SkeletonVisualizer):
    def __init__(self, size_scalar: float = 1.0):
        super().__init__(joint_set=CocoExJoints(), size_scalar=size_scalar)


class PanopticVisualizer(SkeletonVisualizer):
    def __init__(self, size_scalar: float = 1.0):
        super().__init__(joint_set=PanopticJoints(), size_scalar=size_scalar)


class Common14Visualizer(SkeletonVisualizer):
    def __init__(self, size_scalar: float = 1.0):
        super().__init__(joint_set=Common14Joints(), size_scalar=size_scalar)