import torch
import numpy as np
from cmib.data.quaternion import qmul, qrot
import torch.nn as nn

amass_offsets = [
    [0.0, 0.0, 0.0],

    [0.058581, -0.082280, -0.017664],
    [0.043451, -0.386469, 0.008037],
    [-0.014790, -0.426874, -0.037428],
    [0.041054, -0.060286, 0.122042],
    [0.0, 0.0, 0.0],

    [-0.060310, -0.090513, -0.013543],
    [-0.043257, -0.383688, -0.004843],
    [0.019056, -0.420046, -0.034562],
    [-0.034840, -0.062106, 0.130323],
    [0.0, 0.0, 0.0],

    [0.004439, 0.124404, -0.038385],
    [0.004488, 0.137956, 0.026820],
    [-0.002265, 0.056032, 0.002855],
    [-0.013390, 0.211636, -0.033468],
    [0.010113, 0.088937, 0.050410],
    [0.0, 0.0, 0.0],

    [0.071702, 0.114000, -0.018898],
    [0.122921, 0.045205, -0.019046],
    [0.255332, -0.015649, -0.022946],
    [0.265709, 0.012698, -0.007375],
    [0.0, 0.0, 0.0],

    [-0.082954, 0.112472, -0.023707],
    [-0.113228, 0.046853, -0.008472],
    [-0.260127, -0.014369, -0.031269],
    [-0.269108, 0.006794, -0.006027],
    [0.0, 0.0, 0.0]
]

sk_offsets = [
    [-42.198200, 91.614723, -40.067841],

    [0.103456, 1.857829, 10.548506],
    [43.499992, -0.000038, -0.000002],
    [42.372192, 0.000015, -0.000007],
    [17.299999, -0.000002, 0.000003],
    [0.000000, 0.000000, 0.000000],

    [0.103457, 1.857829, -10.548503],
    [43.500042, -0.000027, 0.000008],
    [42.372257, -0.000008, 0.000014],
    [17.299992, -0.000005, 0.000004],
    [0.000000, 0.000000, 0.000000],

    [6.901968, -2.603733, -0.000001],
    [12.588099, 0.000002, 0.000000],
    [12.343206, 0.000000, -0.000001],
    [25.832886, -0.000004, 0.000003],
    [11.766620, 0.000005, -0.000001],
    [0.000000, 0.000000, 0.000000],

    [19.745899, -1.480370, 6.000108],
    [11.284125, -0.000009, -0.000018],
    [33.000050, 0.000004, 0.000032],
    [25.200008, 0.000015, 0.000008],
    [0.000000, 0.000000, 0.000000],

    [19.746099, -1.480375, -6.000073],
    [11.284138, -0.000015, -0.000012],
    [33.000092, 0.000017, 0.000013],
    [25.199780, 0.000135, 0.000422],
    [0.000000, 0.000000, 0.000000],
]

sk_parents = [
    -1,
    0,
    1,
    2,
    3,
    4,
    0,
    6,
    7,
    8,
    9,
    0,
    11,
    12,
    13,
    14,
    15,
    13,
    17,
    18,
    19,
    20,
    13,
    22,
    23,
    24,
    25,
]

sk_joints_to_remove = [5, 10, 16, 21, 26]

joint_names = [
    "Hips",
    "LeftUpLeg",
    "LeftLeg",
    "LeftFoot",
    "LeftToe",
    "RightUpLeg",
    "RightLeg",
    "RightFoot",
    "RightToe",
    "Spine",
    "Spine1",
    "Spine2",
    "Neck",
    "Head",
    "LeftShoulder",
    "LeftArm",
    "LeftForeArm",
    "LeftHand",
    "RightShoulder",
    "RightArm",
    "RightForeArm",
    "RightHand",
]


class Skeleton:
    def __init__(
        self,
        offsets,
        parents,
        joints_left=None,
        joints_right=None,
        bone_length=None,
        device=None,
    ):
        assert len(offsets) == len(parents)

        self._offsets = torch.Tensor(offsets).to(device)
        self._parents = np.array(parents)
        self._joints_left = joints_left
        self._joints_right = joints_right
        self._compute_metadata()

    def num_joints(self):
        return self._offsets.shape[0]

    def offsets(self):
        return self._offsets

    def parents(self):
        return self._parents

    def has_children(self):
        return self._has_children

    def children(self):
        return self._children

    def convert_to_global_pos(self, unit_vec_rerp):
        """
        Convert the unit offset matrix to global position.
        First row(root) will have absolute position value in global coordinates.
        """
        bone_length = self.get_bone_length_weight()
        batch_size = unit_vec_rerp.size(0)
        seq_len = unit_vec_rerp.size(1)
        unit_vec_table = unit_vec_rerp.reshape(batch_size, seq_len, 22, 3)
        global_position = torch.zeros_like(unit_vec_table, device=unit_vec_table.device)

        for i, parent in enumerate(self._parents):
            if parent == -1:  # if root
                global_position[:, :, i] = unit_vec_table[:, :, i]

            else:
                global_position[:, :, i] = global_position[:, :, parent] + (
                    nn.functional.normalize(unit_vec_table[:, :, i], p=2.0, dim=-1)
                    * bone_length[i]
                )

        return global_position

    def convert_to_unit_offset_mat(self, global_position):
        """
        Convert the global position of the skeleton to a unit offset matrix.
        First row(root) will have absolute position value in global coordinates.
        """

        bone_length = self.get_bone_length_weight()
        unit_offset_mat = torch.zeros_like(
            global_position, device=global_position.device
        )

        for i, parent in enumerate(self._parents):

            if parent == -1:  # if root
                unit_offset_mat[:, :, i] = global_position[:, :, i]
            else:
                unit_offset_mat[:, :, i] = (
                    global_position[:, :, i] - global_position[:, :, parent]
                ) / bone_length[i]

        return unit_offset_mat

    def remove_joints(self, joints_to_remove):
        """
        Remove the joints specified in 'joints_to_remove', both from the
        skeleton definition and from the dataset (which is modified in place).
        The rotations of removed joints are propagated along the kinematic chain.
        """
        valid_joints = []
        for joint in range(len(self._parents)):
            if joint not in joints_to_remove:
                valid_joints.append(joint)

        index_offsets = np.zeros(len(self._parents), dtype=int)
        new_parents = []
        for i, parent in enumerate(self._parents):
            if i not in joints_to_remove:
                new_parents.append(parent - index_offsets[parent])
            else:
                index_offsets[i:] += 1
        self._parents = np.array(new_parents)

        self._offsets = self._offsets[valid_joints]
        self._compute_metadata()

    def forward_kinematics(self, rotations, root_positions):
        """
        Perform forward kinematics using the given trajectory and local rotations.
        Arguments (where N = batch size, L = sequence length, J = number of joints):
         -- rotations: (N, L, J, 4) tensor of unit quaternions describing the local rotations of each joint.
         -- root_positions: (N, L, 3) tensor describing the root joint positions.
        """
        assert len(rotations.shape) == 4
        assert rotations.shape[-1] == 4

        positions_world = []
        rotations_world = []

        expanded_offsets = self._offsets.expand(
            rotations.shape[0],
            rotations.shape[1],
            self._offsets.shape[0],
            self._offsets.shape[1],
        )

        # Parallelize along the batch and time dimensions
        for i in range(self._offsets.shape[0]):
            if self._parents[i] == -1:
                positions_world.append(root_positions)
                rotations_world.append(rotations[:, :, 0])
            else:
                positions_world.append(
                    qrot(rotations_world[self._parents[i]], expanded_offsets[:, :, i])
                    + positions_world[self._parents[i]]
                )
                if self._has_children[i]:
                    rotations_world.append(
                        qmul(rotations_world[self._parents[i]], rotations[:, :, i])
                    )
                else:
                    # This joint is a terminal node -> it would be useless to compute the transformation
                    rotations_world.append(None)

        return torch.stack(positions_world, dim=3).permute(0, 1, 3, 2)

    def forward_kinematics_with_rotation(self, rotations, root_positions):
        """
        Perform forward kinematics using the given trajectory and local rotations.
        Arguments (where N = batch size, L = sequence length, J = number of joints):
         -- rotations: (N, L, J, 4) tensor of unit quaternions describing the local rotations of each joint.
         -- root_positions: (N, L, 3) tensor describing the root joint positions.
        """
        assert len(rotations.shape) == 4
        assert rotations.shape[-1] == 4

        positions_world = []
        rotations_world = []

        expanded_offsets = self._offsets.expand(
            rotations.shape[0],
            rotations.shape[1],
            self._offsets.shape[0],
            self._offsets.shape[1],
        )

        # Parallelize along the batch and time dimensions
        for i in range(self._offsets.shape[0]):
            if self._parents[i] == -1:
                positions_world.append(root_positions)
                rotations_world.append(rotations[:, :, 0])
            else:
                positions_world.append(
                    qrot(rotations_world[self._parents[i]], expanded_offsets[:, :, i])
                    + positions_world[self._parents[i]]
                )
                if self._has_children[i]:
                    rotations_world.append(
                        qmul(rotations_world[self._parents[i]], rotations[:, :, i])
                    )
                else:
                    # This joint is a terminal node -> it would be useless to compute the transformation
                    rotations_world.append(
                        torch.Tensor([1, 0, 0, 0])
                        .expand(rotations.shape[0], rotations.shape[1], 4)
                        .to(rotations.device)
                    )

        return torch.stack(positions_world, dim=3).permute(0, 1, 3, 2), torch.stack(
            rotations_world, dim=3
        ).permute(0, 1, 3, 2)

    def get_bone_length_weight(self):
        bone_length = []
        for i, parent in enumerate(self._parents):
            if parent == -1:
                bone_length.append(1)
            else:
                bone_length.append(
                    torch.linalg.norm(self._offsets[i : i + 1], ord="fro").item()
                )
        return torch.Tensor(bone_length)

    def joints_left(self):
        return self._joints_left

    def joints_right(self):
        return self._joints_right

    def _compute_metadata(self):
        self._has_children = np.zeros(len(self._parents)).astype(bool)
        for i, parent in enumerate(self._parents):
            if parent != -1:
                self._has_children[parent] = True

        self._children = []
        for i, parent in enumerate(self._parents):
            self._children.append([])
        for i, parent in enumerate(self._parents):
            if parent != -1:
                self._children[parent].append(i)