import torch
from . import bvh_io as bvh_io
import numpy as np
from ..geometries.quaternion import Quaternions

from .skeleton_presets import SkeletonPresets
from ..kinematics import ForwardKinematicsJoint
from ..geometry import quat2repr6d, quat2mat
from .bvh_io import WriterWrapper
from alfred import logger


def velocity(pos, padding=False):
    velo = pos[1:, ...] - pos[:-1, ...]
    velo_norm = torch.norm(velo, dim=-1)
    if padding:
        pad = torch.zeros_like(velo_norm[:1, :])
        velo_norm = torch.cat([pad, velo_norm], dim=0)
    return velo_norm


def foot_contact(pos, ref_height=1.0, threshold=0.018):
    velo_norm = velocity(pos)
    contact = velo_norm < threshold
    contact = contact.int()
    padding = torch.zeros_like(contact)
    contact = torch.cat([padding[:1, :], contact], dim=0)
    return contact


class Skeleton:
    def __init__(self, names, parent, offsets, joint_reduction=True):
        self._names = names
        self.original_parent = parent
        self._offsets = offsets
        self._parent = None
        self._ee_id = None
        self.contact_names = []
        self.prefix = ""

        for i, name in enumerate(self._names):
            if ":" in name:
                self._names[i] = name.split(":")[-1]
                p = name.split(":")[0]
                self.prefix = p

        if joint_reduction:
            self.skeleton_type, match_num = SkeletonPresets.match(names)
            corps_names = SkeletonPresets.corps_names[self.skeleton_type]
            self.contact_names = SkeletonPresets.contact_names[self.skeleton_type]
            self.contact_threshold = SkeletonPresets.contact_thresholds[
                self.skeleton_type
            ]

            self.contact_id = []
            for i in self.contact_names:
                self.contact_id.append(corps_names.index(i))
        else:
            self.skeleton_type = -1
            corps_names = self._names

        self.details = (
            []
        )  # joints that does not belong to the corps (we are not interested in them)
        for i, name in enumerate(self._names):
            if name not in corps_names:
                self.details.append(i)

        self.corps = []
        self.simplified_name = []
        self.simplify_map = {}
        self.inverse_simplify_map = {}

        # Repermute the skeleton id according to the databse
        for name in corps_names:
            for j in range(len(self._names)):
                if name in self._names[j]:
                    self.corps.append(j)
                    break
        if len(self.corps) != len(corps_names):
            for i in self.corps:
                print(self._names[i], end=" ")
            print(self.corps, self.skeleton_type, len(self.corps), sep="\n")
            raise Exception("Problem in this skeleton")

        self.joint_num_simplify = len(self.corps)
        for i, j in enumerate(self.corps):
            self.simplify_map[j] = i
            self.inverse_simplify_map[i] = j
            self.simplified_name.append(self._names[j])
        self.inverse_simplify_map[0] = -1
        for i in range(len(self._names)):
            if i in self.details:
                self.simplify_map[i] = -1

    @property
    def parent(self):
        if self._parent is None:
            self._parent = self.original_parent[self.corps].copy()
            for i in range(self._parent.shape[0]):
                if i >= 1:
                    self._parent[i] = self.simplify_map[self._parent[i]]
            self._parent = tuple(self._parent)
        return self._parent

    @property
    def offsets(self):
        return torch.tensor(self._offsets[self.corps], dtype=torch.float)

    @property
    def names(self):
        return self.simplified_name

    @property
    def names_full(self):
        # for mixamo only
        return [self.prefix + ":" + i for i in self.simplified_name]
    
    @property
    def names_mixamo(self):
        # for mixamo only
        return ["mixamorig:" + i for i in self.simplified_name]

    @property
    def ee_id(self):
        raise Exception("Abaddoned")
        # if self._ee_id is None:
        #     self._ee_id = []
        #     for i in SkeletonDatabase.ee_names[self.skeleton_type]:
        #         self.ee_id._ee_id(corps_names[self.skeleton_type].index(i))


class BVH_file:
    def __init__(
        self, file_path, no_scale=False, requires_contact=False, joint_reduction=True
    ):
        """
        if joint reduction, we only using a subset from the original BVH
        """
        logger.info(
            f"Loading BVH file: {file_path}, joint_subsampling: {joint_reduction}"
        )
        self.anim = bvh_io.load(file_path)
        self._names = self.anim.names
        print(self._names)
        self.frametime = self.anim.frametime
        self.skeleton = Skeleton(
            self.anim.names, self.anim.parent, self.anim.offsets, joint_reduction
        )

        # Downsample to 30 fps for our application
        if self.frametime < 0.0084:
            self.frametime *= 2
            self.anim.positions = self.anim.positions[::2]
            self.anim.rotations = self.anim.rotations[::2]
        if self.frametime < 0.017:
            self.frametime *= 2
            self.anim.positions = self.anim.positions[::2]
            self.anim.rotations = self.anim.rotations[::2]

        # Scale by 1/100 if it's raw exported bvh from blender
        if not no_scale and self.skeleton.offsets[0, 1] > 10:
            self.scale(1.0 / 100)

        self.requires_contact = requires_contact

        if requires_contact:
            self.contact_names = self.skeleton.contact_names
        else:
            self.contact_names = []

        self.fk = ForwardKinematicsJoint(self.skeleton.parent, self.skeleton.offsets)
        self.writer = WriterWrapper(self.skeleton.parent, self.skeleton.offsets)
        if self.requires_contact:
            gl_pos = self.joint_position()
            self.contact_label = foot_contact(
                gl_pos[:, self.skeleton.contact_id],
                threshold=self.skeleton.contact_threshold,
            )
            self.gl_pos = gl_pos

    def local_pos(self):
        gl_pos = self.joint_position()
        local_pos = gl_pos - gl_pos[:, 0:1, :]
        return local_pos[:, 1:]

    def scale(self, ratio):
        self.anim.offsets *= ratio
        self.anim.positions *= ratio

    def to_tensor(self, repr="euler", rot_only=False):
        if repr not in ["euler", "quat", "quaternion", "repr6d"]:
            raise Exception("Unknown rotation representation")
        positions = self.get_position()
        rotations = self.get_rotation(repr=repr)

        if rot_only:
            return rotations.reshape(rotations.shape[0], -1)

        if self.requires_contact:
            virtual_contact = torch.zeros_like(
                rotations[:, : len(self.skeleton.contact_id)]
            )
            virtual_contact[..., 0] = self.contact_label
            rotations = torch.cat([rotations, virtual_contact], dim=1)

        rotations = rotations.reshape(rotations.shape[0], -1)
        return torch.cat((rotations, positions), dim=-1)

    def joint_position(self):
        positions = torch.tensor(self.anim.positions[:, 0, :], dtype=torch.float)
        rotations = self.anim.rotations[:, self.skeleton.corps, :]
        rotations = Quaternions.from_euler(np.radians(rotations)).qs
        rotations = torch.tensor(rotations, dtype=torch.float)
        j_loc = self.fk.forward(rotations, positions)
        return j_loc

    def get_rotation(self, repr="quat"):
        rotations = self.anim.rotations[:, self.skeleton.corps, :]
        if repr == "quaternion" or repr == "quat" or repr == "repr6d":
            rotations = Quaternions.from_euler(np.radians(rotations)).qs
            rotations = torch.tensor(rotations, dtype=torch.float)
        if repr == "repr6d":
            rotations = quat2repr6d(rotations)
        if repr == "euler":
            rotations = torch.tensor(rotations, dtype=torch.float)
        return rotations

    def get_position(self):
        return torch.tensor(self.anim.positions[:, 0, :], dtype=torch.float)

    def dfs(self, x, vis, dist):
        fa = self.skeleton.parent
        vis[x] = 1
        for y in range(len(fa)):
            if (fa[y] == x or fa[x] == y) and vis[y] == 0:
                dist[y] = dist[x] + 1
                self.dfs(y, vis, dist)

    def get_neighbor(self, threshold, enforce_contact=False):
        fa = self.skeleton.parent
        neighbor_list = []
        for x in range(0, len(fa)):
            vis = [0 for _ in range(len(fa))]
            dist = [0 for _ in range(len(fa))]
            self.dfs(x, vis, dist)
            neighbor = []
            for j in range(0, len(fa)):
                if dist[j] <= threshold:
                    neighbor.append(j)
            neighbor_list.append(neighbor)

        contact_list = []
        if self.requires_contact:
            for i, p_id in enumerate(self.skeleton.contact_id):
                v_id = len(neighbor_list)
                neighbor_list[p_id].append(v_id)
                neighbor_list.append(neighbor_list[p_id])
                contact_list.append(v_id)

        root_neighbor = neighbor_list[0]
        id_root = len(neighbor_list)

        if enforce_contact:
            root_neighbor = root_neighbor + contact_list
            for j in contact_list:
                neighbor_list[j] = list(set(neighbor_list[j]))

        root_neighbor = list(set(root_neighbor))
        for j in root_neighbor:
            neighbor_list[j].append(id_root)
        root_neighbor.append(id_root)
        neighbor_list.append(root_neighbor)  # Neighbor for root position
        return neighbor_list
