import torch
import numpy as np

from typing import Union, List, overload
from wis3d import Wis3D

from lib.platform import PM
from lib.utils.geometry.rotation import axis_angle_to_matrix


class HWis3D(Wis3D):
    ''' Abstraction of Wis3D for human motion. '''

    def __init__(
        self,
        out_path    : str   = PM.outputs / 'wis3d',
        seq_name    : str   = 'debug',
        xyz_pattern : tuple = ('x', 'y', 'z'),
    ):
        seq_name = seq_name.replace('/', '-')
        super().__init__(out_path, seq_name, xyz_pattern)


    def add_text(self, text:str):
        '''
        Add an item of vertices whose name is used to put the text message. *Dirty use!*

        ### Args
        - text: str
        '''
        fake_verts = np.array([[0, 0, 0]])
        self.add_point_cloud(
            vertices = fake_verts,
            colors   = None,
            name     = text,
        )


    def add_text_seq(self, texts:List[str], offset:int=0):
        '''
        Add an item of vertices whose name is used to put the text message. *Dirty use!*

        ### Args
        - texts: List[str]
            - The list of text messages.
        - offset: int, default = 0
            - The offset for the sequence index.
        '''
        fake_verts = np.array([[0, 0, 0]])
        for i, text in enumerate(texts):
            self.set_scene_id(i + offset)
            self.add_point_cloud(
                vertices = fake_verts,
                colors   = None,
                name     = text,
            )

    def add_image_seq(self, imgs:List[np.ndarray], name:str, offset:int=0):
        '''
        Add an item of vertices whose name is used to put the image. *Dirty use!*

        ### Args
        - imgs: List[np.ndarray]
            - The list of images.
        - offset: int, default = 0
            - The offset for the sequence index.
        '''
        for i, img in enumerate(imgs):
            self.set_scene_id(i + offset)
            self.add_image(
                image = img,
                name  = name,
            )

    def add_motion_mesh(
        self,
        verts : Union[torch.Tensor, np.ndarray],
        faces : Union[torch.Tensor, np.ndarray],
        name  : str,
        offset: int = 0,
    ):
        '''
        Add sequence of vertices and face(s) to the wis3d viewer.

        ### Args
        - verts: torch.Tensor or np.ndarray, (L, V, 3), L ~ sequence length, V ~ number of vertices
        - faces: torch.Tensor or np.ndarray, (F, 3) or (L, F, 3), F ~ number of faces, L ~ sequence length
        - name: str
            - The name of the point cloud.
        - offset: int, default = 0
            - The offset for the sequence index.
        '''
        assert (len(verts.shape) == 3), 'The input `verts` should have 3 dimensions: (L, V, 3).'
        assert (verts.shape[-1] == 3), 'The last dimension of `verts` should be 3.'
        if isinstance(verts, np.ndarray):
            verts = torch.from_numpy(verts)
        if isinstance(faces, torch.Tensor):
            faces = faces.detach().cpu().numpy()
        if len(faces.shape) == 2:
            faces = faces[None].repeat(verts.shape[0], 0)
        assert (len(faces.shape) == 3), 'The input `faces` should have 2 or 3 dimensions: (F, 3) or (L, F, 3).'
        assert (faces.shape[-1] == 3), 'The last dimension of `faces` should be 3.'
        assert (verts.shape[0] == faces.shape[0]), 'The first dimension of `verts` and `faces` should be the same.'

        L, _, _ = verts.shape
        verts = verts.detach().cpu()

        # Add vertices frame by frame.
        for i in range(L):
            self.set_scene_id(i + offset)
            self.add_mesh(
                vertices = verts[i],
                faces    = faces[i],
                name     = name,
            )  # type: ignore

        # Reset Wis3D scene id.
        self.set_scene_id(0)


    def add_motion_verts(
        self,
        verts : Union[torch.Tensor, np.ndarray],
        name  : str,
        offset: int = 0,
    ):
        '''
        Add sequence of vertices to the wis3d viewer.

        ### Args
        - verts: torch.Tensor or np.ndarray, (L, V, 3), L ~ sequence length, V ~ number of vertices
        - name: str
            - The name of the point cloud.
        - offset: int, default = 0
            - The offset for the sequence index.
        '''
        assert (len(verts.shape) == 3), 'The input `verts` should have 3 dimensions: (L, V, 3).'
        assert (verts.shape[-1] == 3), 'The last dimension of `verts` should be 3.'
        if isinstance(verts, np.ndarray):
            verts = torch.from_numpy(verts)

        L, _, _ = verts.shape
        verts = verts.detach().cpu()

        # Add vertices frame by frame.
        for i in range(L):
            self.set_scene_id(i + offset)
            self.add_point_cloud(
                vertices = verts[i],
                colors   = None,
                name     = name,
            )

        # Reset Wis3D scene id.
        self.set_scene_id(0)


    def add_motion_skel(
        self,
        joints    : Union[torch.Tensor, np.ndarray],
        bones     : Union[list, torch.Tensor],
        colors    : Union[list, torch.Tensor],
        name      : str,
        offset    : int = 0,
        threshold : float = 0.5,
    ):
        '''
        Add sequence of joints with specific skeleton to the wis3d viewer.

        ### Args
        - joints: torch.Tensor or np.ndarray, shape = (L, J, 3) or (L, J, 4), L ~ sequence length, J ~ number of joints
        - bones: list
            - A list of bones of the skeleton, i.e. the edge in the kinematic trees.
        - colors: list
        - name: str
            - The name of the point cloud.
        - offset: int, default = 0
            - The offset for the sequence index.
        - threshold: float, default = 0.5
            - Threshold to filter the confidence of the joints. It's useless when no confidence provided.
        '''
        assert (len(joints.shape) == 3), 'The input `joints` should have 3 dimensions: (L, J, 3).'
        assert (joints.shape[-1] == 3 or joints.shape[-1] == 4), 'The last dimension of `joints` should be 3 or 4.'
        if isinstance(joints, np.ndarray):
            joints = torch.from_numpy(joints)
        if isinstance(bones, List):
            bones = torch.tensor(bones)
        if isinstance(colors, List):
            colors = torch.tensor(colors)

        # Get the sequence length.
        joints = joints.detach().cpu() # (L, J, 3) or (L, J, 4)
        L, J, D = joints.shape
        if D == 4:
            conf = joints[:, :, 3]
            joints = joints[:, :, :3]
        else:
            conf = None

        # Add vertices frame by frame.
        for i in range(L):
            self.set_scene_id(i + offset)
            bones_s = joints[i][bones[:, 0]]
            bones_e = joints[i][bones[:, 1]]
            if conf is not None:
                mask = torch.logical_and(conf[i][bones[:, 0]] > threshold, conf[i][bones[:, 1]] > threshold)
                bones_s, bones_e = bones_s[mask], bones_e[mask]
            if len(bones_s) > 0:
                self.add_lines(
                    start_points = bones_s,
                    end_points   = bones_e,
                    colors       = colors,
                    name         = name,
                )

        # Reset Wis3D scene id.
        self.set_scene_id(0)


    def add_vec_seq(
        self,
        vecs    : torch.Tensor,
        name    : str,
        offset  : int = 0,
        seg_num : int = 16,
    ):
        '''
        Add directional line sequence to the wis3d viewer.

        The line will be gradient colored, and the direction of the vector is visualized as from dark to light.

        ### Args
        - vecs: torch.Tensor, (L, 2, 3) or (L, N, 2, 3), L ~ sequence length, N ~ vectors counts in one frame,
              then give the start 3D point and end 3D point.
        - name: str
            - The name of the vector.
        - offset: int, default = 0
            - The offset for the sequence index.
        - seg_num: int, default = 16
            - The number of segments for gradient color, will just change the visualization effect.
        '''
        if len(vecs.shape) == 3:
            vecs = vecs[:, None, :, :] # (L, 2, 3) -> (L, 1, 2, 3)
        assert (len(vecs.shape) == 4), 'The input `vecs` should have 3 or 4 dimensions: (L, 2, 3) or (L, N, 2, 3).'
        assert (vecs.shape[-2:] == (2, 3)), f'The last two dimension of `vecs` should be (2, 3), but got vecs.shape = {vecs.shape}.'

        # Get the sequence length.
        L, N, _, _ = vecs.shape
        vecs = vecs.detach().cpu()

        # Cut the line into segments.
        steps_delta = (vecs[:, :, [1]] - vecs[:, :, [0]]) / (seg_num + 1) # (L, N, 1, 3)
        steps_cnt   = torch.arange(seg_num + 1).reshape((1, 1, seg_num + 1, 1)) # (1, 1, seg_num+1, 1)
        segs = steps_delta * steps_cnt + vecs[:, :, [0]] # (L, N, seg_num+1, 3)
        start_pts = segs[:, :, :-1] # (L, N, seg_num, 3)
        end_pts   = segs[:, :, 1:] # (L, N, seg_num, 3)

        # Prepare the gradient colors.
        grad_colors = torch.linspace(0, 255, seg_num).reshape((1, seg_num, 1)).repeat(N, 1, 3) # (N, seg_num, 3)

        # Add vertices frame by frame.
        for i in range(L):
            self.set_scene_id(i + offset)
            self.add_lines(
                start_points = start_pts[i].reshape(-1, 3),
                end_points   = end_pts[i].reshape(-1, 3),
                colors       = grad_colors.reshape(-1, 3),
                name         = name,
            )

        # Reset Wis3D scene id.
        self.set_scene_id(0)


    def add_traj(
        self,
        positions : torch.Tensor,
        name      : str,
        offset    : int = 0,
    ):
        '''
        Visualize the the positions change across the time as trajectory. The newer position will be brighter.

        ### Args
        - positions: torch.Tensor, (L, 3), L ~ sequence length
        - name: str
            - The name of the trajectory.
        - offset: int, default = 0
            - The offset for the sequence index.
        '''
        assert (len(positions.shape) == 2), 'The input `positions` should have 2 dimensions: (L, 3).'
        assert (positions.shape[-1] == 3), 'The last dimension of `positions` should be 3.'

        # Get the sequence length.
        L, _ = positions.shape
        positions = positions.detach().cpu()
        traj = positions[[0]] # (1, 3)

        # Prepare the gradient colors.
        grad_colors = torch.linspace(208, 48, L).reshape((L, 1)).repeat(1, 3) # (L, 3)

        for i in range(L):
            traj = torch.cat((traj, positions[[i]]), dim=0) # (i+2, 3)
            self.set_scene_id(i + offset)
            self.add_lines(
                start_points = traj[:-1],
                end_points   = traj[1:],
                colors       = grad_colors[-(i+1):],
                name         = name,
            )

        # Reset Wis3D scene id.
        self.set_scene_id(0)


    def add_sphere_sensors(
        self,
        positions  : torch.Tensor,
        radius     : Union[torch.Tensor, float],
        activities : torch.Tensor,
        name       : str,
    ):
        '''
        Draw the sphere sensors with different colors to represent the activities. The color is from white to red.

        ### Args
        - positions: torch.Tensor, (N, 3), N ~ number of sensors
        - radius: torch.Tensor or float, (N,), N ~ number of sensors
        - activities: torch.Tensor, (N)
            - The activities of the sensors, from 0 to 1.
        - name: str
            - The name of the spheres.
        '''
        assert (len(positions.shape) == 2), 'The input `positions` should have 2 dimensions: (N, 3).'
        assert (positions.shape[-1] == 3), 'The last dimension of `positions` should be 3.'
        N, _ = positions.shape
        if isinstance(radius, float):
            radius = torch.Tensor(radius).reshape(1).repeat(N) # (N)
        elif len(radius.shape) == 0:
            radius = radius.reshape(1).repeat(N)

        colors = torch.ones(size=(N, 3)).float()
        colors[:, 0] = 255
        colors[:, 1] = (1 - activities) ** 2 * 255
        colors[:, 2] = (1 - activities) ** 2 * 255
        self.add_spheres(
            centers = positions,
            radius  = radius,
            colors  = colors,
            name    = name,
        )


    def add_sphere_sensors_seq(
        self,
        positions  : torch.Tensor,
        radius     : Union[torch.Tensor, float],
        activities : torch.Tensor,
        name       : str,
        offset     : int = 0,
    ):
        '''
        Draw the sphere sensors with different colors to represent the activities. The color is from white to red.

        ### Args
        - positions: torch.Tensor, (L, N, 3), N ~ number of sensors
        - radius: torch.Tensor or float, (L, N,), N ~ number of sensors
        - activities: torch.Tensor, (L, N)
            - The activities of the sensors, from 0 to 1.
        - name: str
            - The name of the spheres.
        - offset: int, default = 0
            - The offset for the sequence index.
        '''
        assert (len(positions.shape) == 3), 'The input `positions` should have 3 dimensions: (L, N, 3).'
        assert (positions.shape[-1] == 3), 'The last dimension of `positions` should be 3.'
        L, N, _ = positions.shape

        for i in range(L):
            self.set_scene_id(i + offset)
            self.add_sphere_sensors(
                positions  = positions[i],
                radius     = radius,
                activities = activities[i],
                name       = name,
            )


    # ===== Overriding methods from original Wis3D. =====


    def add_lines(
        self,
        start_points: torch.Tensor,
        end_points  : torch.Tensor,
        colors      : Union[list, torch.Tensor] = None,
        name        : str   = None,
        thickness   : float = 0.01,
        resolution  : int   = 4,
    ):
        '''
        Add lines by points. Overriding the original `add_lines` method to use mesh to provide browser from crash.

        ### Args
        - start_points: torch.Tensor, (N, 3), N ~ number of lines
        - end_points: torch.Tensor, (N, 3), N ~ number of lines
        - colors: list or torch.Tensor, (N, 3)
            - The color of the lines, from 0 to 255.
        - name: str
            - The name of the vector.
        - thickness: float, default = 0.01
            - The thickness of the lines.
        - resolution: int, default = 3
            - The 'line' was actually a poly-cylinder, and the resolution how it looks like a cylinder.
        '''
        if isinstance(colors, List):
            colors = torch.tensor(colors)

        assert (len(start_points.shape) == 2), 'The input `start_points` should have 2 dimensions: (N, 3).'
        assert (len(end_points.shape) == 2), 'The input `end_points` should have 2 dimensions: (N, 3).'
        assert (start_points.shape == end_points.shape), 'The input `start_points` and `end_points` should have the same shape.'

        # ===== Prepare the data. =====
        N, _ = start_points.shape
        device = start_points.device
        dir = end_points - start_points # (N, 3)
        dis = torch.norm(dir, dim=-1, keepdim=True) # (N, 1)
        dir = dir / dis # (N, 3)
        K = resolution + 1 # the first & the last point share the position
        # Find out directions that are negative to the y-axis.
        vec_y = torch.Tensor([[0, 1, 0]]).float().to(device) # (1, 3)
        neg_mask = (dir @ vec_y.transpose(-1, -2) < 0).squeeze() # (N,)

        # ===== Get the ending surface vertices of the cylinder. =====
        # 1. Get the surface vertices template in x-z plain.
        radius = torch.linspace(0, 2*torch.pi, K) # (K,)
        v_ending_temp = \
            torch.stack(
                [torch.cos(radius), torch.zeros_like(radius), torch.sin(radius)],
                dim = -1
            ) # (K, 3)
        v_ending_temp *= thickness # (K, 3)
        v_ending_temp = v_ending_temp[None].repeat(N, 1, 1) # (N, K, 3)

        # 2. Rotate the template plane to the direction of the line.
        rot_axis = torch.linalg.cross(vec_y, dir) # (N, 3)
        rot_axis[neg_mask] *= -1
        rot_mat = axis_angle_to_matrix(rot_axis) # (N, 3, 3)
        v_ending_temp = v_ending_temp @ rot_mat.transpose(-1, -2)
        v_ending_temp = v_ending_temp.to(device)

        # 3. Move the template plane to the start and end points and get the cylinder vertices.
        v_cylinder_start = v_ending_temp + start_points[:, None] # (N, K, 3)
        v_cylinder_end = v_ending_temp + end_points[:, None] # (N, K, 3)
        #    Swap the start and end points for the negative direction to adjust the normal direction.
        v_cylinder_start[neg_mask], v_cylinder_end[neg_mask] = v_cylinder_end[neg_mask], v_cylinder_start[neg_mask]
        v_cylinder = torch.cat([v_cylinder_start, v_cylinder_end], dim=1) # (N, 2*K, 3)

        # ===== Calculate the face index. =====
        idx = torch.arange(0, 2*K, device=device).to(device) # (2*K,)
        idx_s, idx_e = idx[:K], idx[K:]
        f_cylinder = torch.cat([
            # Two ending surface.
            torch.stack([idx_s[0].repeat(K-2), idx_s[1:-1], idx_s[2:]], dim=-1),
            torch.stack([idx_e[0].repeat(K-2), idx_e[2:], idx_e[1:-1]], dim=-1),
            # The side surface.
            torch.stack([idx_e[:-1], idx_s[1:], idx_s[:-1]], dim=-1),
            torch.stack([idx_e[:-1], idx_e[1:], idx_s[1:]], dim=-1),
        ], dim=0) # (4*K-4, 3)
        f_cylinder = f_cylinder[None].repeat(N, 1, 1) # (N, 4*K-4, 3)

        # ===== Calculate the face index. =====
        if colors is not None:
            c_cylinder = colors / 255.0 # (N, 3)
            c_cylinder = c_cylinder[:, None].repeat(1, 2*K, 1) # (N, 2*K, 3)
        else:
            c_cylinder = None

        N, V = v_cylinder.shape[:2]
        v_cylinder = v_cylinder.reshape(-1, 3) # (N*(2*K), 3)

        # ===== Manually match the points index before flatten. =====
        f_cylinder = f_cylinder + torch.arange(0, N, device=device).unsqueeze(1).unsqueeze(1) * V
        f_cylinder = f_cylinder.reshape(-1, 3) # (N*(4*K-4), 3)
        if c_cylinder is not None:
            c_cylinder = c_cylinder.reshape(-1, 3) # (N*(2*K), 3)

        self.add_mesh(
            vertices      = v_cylinder,
            vertex_colors = c_cylinder,
            faces         = f_cylinder,
            name          = name,
        )