<reponame>seanschneeweiss/RoSeMotion
import re
import numpy as np
import pandas
import sys

from config.Skeleton import Skeleton
from config.AnybodyFirstFrame import AnybodyFirstFrame
from config.BasisFirstFrame import BasisFirstFrame
from resources.pymo.pymo.data import MocapData
from RotationUtil import rot2eul, get_order
from resources.LeapSDK.v41_python38 import Leap


class LeapData:
    """
    A class to convert LeapMotion frames to PyMO data structure (MocapData)

    Calculates translations (offsets) and rotation data for the joints
    """

    def __init__(self, channel_setting='rotation', frame_rate=0.033333, anybody_basis=True):
        self._skeleton = {}
        self._setting = Skeleton(channel_setting)
        self._motion_channels = []
        self._motions = []
        self._root_name = ''
        self.data = MocapData()

        self.first_frame = None
        self.anybody_first_frame = AnybodyFirstFrame()
        self.basis_first_frame = BasisFirstFrame()
        # anybody_reference = True -> Use Anybody basis from config/*.json (see AnybodyFirstFrame)
        # anybody_reference = False -> Use Leap Motion First Frame for basis
        self.anybody_basis = anybody_basis
        self.status = 0
        self._frame_rate = frame_rate

        self._skeleton = self._setting.skeleton
        # fill channels into skeleton in selected order (i.e. xyz)
        self._skeleton_apply_channels(self._setting.channel_setting)
        self._root_name = self._setting.root_name

        # initialize offsets for each joint
        for joint_name, joint_value in self._skeleton.items():
            joint_value['offsets'] = [0, 0, 0]
            for channel in joint_value['channels']:
                self._motion_channels.append((joint_name, channel))

    def parse(self):
        self.data.skeleton = self._skeleton
        self.data.channel_names = self._motion_channels
        if not self.first_frame:
            sys.exit("No data was recorded - will terminate now!")
        self.data.values = self._motion2dataframe()
        self.data.root_name = self._root_name
        self.data.framerate = self._frame_rate

        return self.data

    def _check_frame(self, frame):
        """
        Check whether frame and hand and fingers are valid, produce error when left hand is shown
        """
        status_no_hand = 1
        status_no_finger = 2
        status_left_hand = 3
        status_valid = 4

        if frame.hands.is_empty:
            if self.status != status_no_hand:
                print("-- No hand found. --")
                self.status = status_no_hand
            return False

        # Get the first hand
        hand = frame.hands[0]

        if hand.is_left:
            if self.status != status_left_hand:
                print("-- Please use your right hand. --")
                self.status = status_left_hand
            return False

        if not hand.is_right and not hand.is_valid:
            return False

        # Check if the hand has any fingers
        fingers = hand.fingers
        if fingers.is_empty:
            if self.status != status_no_finger:
                print("-- No valid fingers found. --")
                self.status = status_no_finger
            return False

        # frame_number = 0 if not self.first_frame else frame.id - self.first_frame.id
        if self.status != status_valid:
            print("-- Valid right hand found, recording data. --")
            self.status = status_valid
        return True

    def add_frame(self, frame):
        if not self._check_frame(frame):
            return None

        # Get the first hand
        hand = frame.hands[0]
        if not self.first_frame:
            self.first_frame = frame
            channel_values = self._get_channel_values(hand, firstframe=True)
            self._motions.append((0, channel_values))
            return

        channel_values = self._get_channel_values(hand)
        self._motions.append((frame.timestamp - self.first_frame.timestamp, channel_values))
        return frame

    def _get_channel_values(self, hand, firstframe=False):
        channel_values = []
        # export_basis = {}
        for joint_name, joint_value in self._skeleton.items():
            # motion data with rotations
            if joint_name == self._root_name:
                x_pos, y_pos, z_pos = LeapData._get_root_offset()
            elif joint_name == 'RightElbow':
                x_pos, y_pos, z_pos = LeapData._get_elbow_offset(hand)
            elif joint_name == 'RightHand':
                x_pos, y_pos, z_pos = LeapData._get_wrist_offset(hand)
            else:
                x_pos, y_pos, z_pos = LeapData._get_finger_offset(joint_name, hand)

            x_rot, y_rot, z_rot = self._calculate_euler_angles(hand, joint_name)

            x_rot *= Leap.RAD_TO_DEG
            y_rot *= Leap.RAD_TO_DEG
            z_rot *= Leap.RAD_TO_DEG

            # if joint_name == 'RightHand':
            #     print(x_rot, y_rot, z_rot)

            if firstframe:
                if self.anybody_basis:
                    x_pos, y_pos, z_pos = self._calculate_offset(joint_name, [x_pos, y_pos, z_pos])
                joint_value['offsets'] = [x_pos, y_pos, z_pos]

                # # dump the basis of leap motion bones
                # if 'End' not in joint_name and 'Root' not in joint_name:
                #     export_basis[joint_name] = np.ndarray.tolist(self._get_basis(hand, joint_name))

            for channel in joint_value['channels']:
                if channel == 'Xposition':
                    channel_values.append((joint_name, channel, x_pos))
                if channel == 'Yposition':
                    channel_values.append((joint_name, channel, y_pos))
                if channel == 'Zposition':
                    channel_values.append((joint_name, channel, z_pos))
                if channel == 'Xrotation':
                    channel_values.append((joint_name, channel, x_rot))
                if channel == 'Yrotation':
                    channel_values.append((joint_name, channel, y_rot))
                if channel == 'Zrotation':
                    channel_values.append((joint_name, channel, z_rot))

        # # dump the basis of leap motion bones
        # if firstframe:
        #     import json
        #     import datetime
        #     with open('../output/{}basis.json'.format(datetime.datetime.today().strftime('%Y%m%d_%H%M%S')), 'w') as o:
        #         json.dump(export_basis, o)

        return channel_values

    def _calculate_euler_angles(self, hand, joint_name):
        initial_hand = self.first_frame.hands[0]

        # special case for root and finger tip
        if joint_name == self._root_name or not self._skeleton[joint_name]['children']:
            return 0.0, 0.0, 0.0

        if self.anybody_basis:
            # compare basis to anybody basis
            parent_initial_basis = self._get_basis_first_frame(self._skeleton[joint_name]['parent'])
            initial_basis = self._get_basis_first_frame(joint_name)
        else:
            # compare basis to first frame from Leap Motion
            parent_initial_basis = self._get_basis(initial_hand, self._skeleton[joint_name]['parent'])
            initial_basis = self._get_basis(initial_hand, joint_name)

        parent_basis = self._get_basis(hand, self._skeleton[joint_name]['parent'])
        basis = self._get_basis(hand, joint_name)

        # if joint_name == 'RightHand':
        #     print(basis)

        # calculation of local rotation matrix - important!!!
        rot = np.matmul(
                np.matmul(
                    initial_basis, np.transpose(basis)
                ),
                np.transpose(
                    np.matmul(
                        parent_initial_basis, np.transpose(parent_basis)
                    )
                )
        )
        return rot2eul(rot)

    def _get_basis(self, hand, joint_name):
        if joint_name == self._root_name:
            return np.array([[1, 0, 0],
                             [0, 1, 0],
                             [0, 0, 1]])
        if joint_name == 'RightElbow':
            return LeapData._basismatrix(hand.arm.basis)
        if joint_name == 'RightHand':
            return LeapData._basismatrix(hand.basis)

        # else, return basis of the finger
        finger, bone_number = LeapData._split_key(joint_name)
        fingerlist = hand.fingers.finger_type(LeapData._get_finger_type(finger))
        bone = fingerlist[0].bone(LeapData._get_bone_type(bone_number))
        return LeapData._basismatrix(bone.basis)

    def _get_basis_first_frame(self, joint_name):
        if joint_name == self._root_name:
            return np.array([[1, 0, 0],
                             [0, 1, 0],
                             [0, 0, 1]])
        return self.anybody_first_frame.get_basis(joint_name)

    def _calculate_offset(self, joint_name, leap_offset):
        if joint_name == self._root_name:
            return 0, 0, 0

        leap_length = np.linalg.norm(leap_offset)

        position = self._get_anybody_position(joint_name)
        position_parent = self._get_anybody_position(self._skeleton[joint_name]['parent'])
        offset_anybody = position - position_parent
        # make a unit vector
        offset_anybody = np.divide(offset_anybody, np.linalg.norm(offset_anybody))
        offset = offset_anybody * leap_length
        return offset[0], offset[1], offset[2]

    def _get_anybody_position(self, joint_name):
        if joint_name == self._root_name:
            return np.array([0, 0, 0])
        return self.anybody_first_frame.get_position(joint_name)

    @staticmethod
    def _get_root_offset():
        return 0, 0, 0

    @staticmethod
    def _get_elbow_offset(hand):
        arm = hand.arm
        return arm.elbow_position.x, arm.elbow_position.y, arm.elbow_position.z

    @staticmethod
    def _get_wrist_offset(hand):
        arm = hand.arm
        x_wrist = hand.wrist_position.x - arm.elbow_position.x
        y_wrist = hand.wrist_position.y - arm.elbow_position.y
        z_wrist = hand.wrist_position.z - arm.elbow_position.z

        return x_wrist, y_wrist, z_wrist

    @staticmethod
    def _get_finger_offset(key, hand):
        key, bone_number = LeapData._split_key(key)

        fingerlist = hand.fingers.finger_type(LeapData._get_finger_type(key))

        # vector between wrist and joint metacarpal proximal (length of carpals)
        if bone_number == 1 or ('Thumb' in key and bone_number == 2):
            bone = fingerlist[0].bone(LeapData._get_bone_type(bone_number))
            return \
                bone.prev_joint.x - hand.wrist_position.x, \
                bone.prev_joint.y - hand.wrist_position.y, \
                bone.prev_joint.z - hand.wrist_position.z

        # vector for bones metacarpal, proximal, intermediate, distal
        bone = fingerlist[0].bone(LeapData._get_bone_type(bone_number - 1))
        return \
            bone.next_joint.x - bone.prev_joint.x, \
            bone.next_joint.y - bone.prev_joint.y, \
            bone.next_joint.z - bone.prev_joint.z

    @staticmethod
    def _split_key(key):
        key_split = re.split('(\d)', key)
        key = key_split[0]
        if key_split[-1] == '_Nub':
            return key, 5
        else:
            return key, int(key_split[1])

    @staticmethod
    def _get_finger_type(key):
        if key == 'RightHandThumb':
            return Leap.Finger.TYPE_THUMB
        if key == 'RightHandIndex':
            return Leap.Finger.TYPE_INDEX
        if key == 'RightHandMiddle':
            return Leap.Finger.TYPE_MIDDLE
        if key == 'RightHandRing':
            return Leap.Finger.TYPE_RING
        if key == 'RightHandPinky':
            return Leap.Finger.TYPE_PINKY
        else:
            raise Exception('Key ({}) did not match'.format(key))

    @staticmethod
    def _get_bone_type(bone_number):
        if bone_number == 4:
            return Leap.Bone.TYPE_DISTAL
        if bone_number == 3:
            return Leap.Bone.TYPE_INTERMEDIATE
        if bone_number == 2:
            return Leap.Bone.TYPE_PROXIMAL
        if bone_number == 1:
            return Leap.Bone.TYPE_METACARPAL
        else:
            raise Exception('bone number ({}) did not match'.format(bone_number))

    @staticmethod
    def _basismatrix(basis):
        return np.array([[basis.x_basis.x, basis.y_basis.x, basis.z_basis.x],
                        [basis.x_basis.y, basis.y_basis.y, basis.z_basis.y],
                        [basis.x_basis.z, basis.y_basis.z, basis.z_basis.z]])

    def _get_channels(self, joint_name, channel_setting):
        if '_Nub' in joint_name:
            return []

        channels_position = ['Xposition', 'Yposition', 'Zposition']
        channels_rotation = ['Xrotation', 'Yrotation', 'Zrotation']

        order = get_order()  # rotation order, i.e. xyz
        channels_rotation = \
            [channels_rotation[order[0]]] + [channels_rotation[order[1]]] + [channels_rotation[order[2]]]

        if channel_setting == 'position' or joint_name == self._root_name:
            return channels_position + channels_rotation

        return channels_rotation

    def _skeleton_apply_channels(self, channel_setting):
        for joint_name, joint_dict in self._skeleton.items():
            joint_dict['channels'] = self._get_channels(joint_name, channel_setting)

    def _motion2dataframe(self):
        """Returns all of the channels parsed from the LeapMotion sensor as a pandas DataFrame"""

        time_index = pandas.to_timedelta([f[0] for f in self._motions], unit='s')
        frames = [f[1] for f in self._motions]
        channels = np.asarray([[channel[2] for channel in frame] for frame in frames])
        column_names = ['%s_%s' % (c[0], c[1]) for c in self._motion_channels]

        return pandas.DataFrame(data=channels, index=time_index, columns=column_names)
