import torch
import bvh.bvh_io as bvh_io
import numpy as np
from bvh.Quaternions import Quaternions
from bvh.skeleton_database import SkeletonDatabase
from models.kinematics import ForwardKinematicsJoint
from models.transforms import quat2repr6d, quat2mat
from models.contact import foot_contact
from bvh.bvh_writer import WriterWrapper


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 = []

        for i, name in enumerate(self._names):
            if ':' in name:
                self._names[i] = name[name.find(':')+1:]

        if joint_reduction:
            self.skeleton_type, match_num = SkeletonDatabase.match(names)
            corps_names = SkeletonDatabase.corps_names[self.skeleton_type]
            self.contact_names = SkeletonDatabase.contact_names[self.skeleton_type]
            self.contact_threshold = SkeletonDatabase.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 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):
        self.anim = bvh_io.load(file_path)
        self._names = self.anim.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. / 100)
        # Scale by 0.14 for Xia's dataset
        if not no_scale and self.skeleton.skeleton_type == 3:
            self.scale(0.14)

        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