import bpy
import sys
import os

mydir = "C:\\UntrackedGit\\BvhToDM_Blender"
if not mydir in sys.path:
    sys.path.append(mydir)
print(mydir)

import BvhToDM
import importlib
importlib.reload(BvhToDM)

from BvhToDM import BoneConstructionData, BvhToDM, BvhToDMSettings

target_name = "rotation_prob"
skel_name = "DeepMimicSkeleton"

# Couple every DM_bone to BVH bone
# Structure: (Target_bone_name: str, inherit_bone_length: bool, inherit_bone_pos: bool, use_connect: bool, bone_direction_vec: list[float])
bonedict = {
    "root": BoneConstructionData("Hip", False, False, False, [0.0, 0.0, 1.0]),
    "spine": BoneConstructionData("Hip", False, False, True, [0.0, 0.0, 1.0]),
    "chest": BoneConstructionData("ShoulderCenter", False, False, True, [0.0, 0.0, 1.0], use_DT=True),
    "head": BoneConstructionData("Head", False, False, True, [0.0, 0.0, 1.0], use_DT=True),
    "right_upper_arm": BoneConstructionData("ShoulderRight", True, True, False, [0.0, 0.0, -1.0], use_IK=True, IK_chain_length=1),
    "right_lower_arm": BoneConstructionData("ElbowRight", True, False, True, [0.0, 0.0, -1.0], use_IK=True, IK_chain_length=2),
    "right_hand": BoneConstructionData("WristRight", True, False, True, [0.0, 0.0, -1.0], use_DT=True),
    "left_upper_arm": BoneConstructionData("ShoulderLeft", True, True, False, [0.0, 0.0, -1.0], use_IK=True, IK_chain_length=1),
    "left_lower_arm": BoneConstructionData("ElbowLeft", True, False, True, [0.0, 0.0, -1.0], use_IK=True, IK_chain_length=2),
    "left_hand": BoneConstructionData("WristLeft", True, False, True, [0.0, 0.0, -1.0], use_DT=True),
    "right_thigh": BoneConstructionData("HipRight", True, True, False, [0.0, 0.0, -1.0], use_IK=True, IK_chain_length=1),
    "right_shin": BoneConstructionData("KneeRight", True, False, True, [0.0, 0.0, -1.0], use_IK=True, IK_chain_length=2),
    "right_foot": BoneConstructionData("AnkleRight", True, False, True, [0.0, -1.0, 0.0], use_DT=True),
    "left_thigh": BoneConstructionData("HipLeft", True, True, False, [0.0, 0.0, -1.0], use_IK=True, IK_chain_length=1),
    "left_shin": BoneConstructionData("KneeLeft", True, False, True, [0.0, 0.0, -1.0], use_IK=True, IK_chain_length=2),
    "left_foot": BoneConstructionData("AnkleLeft", True, False, True, [0.0, -1.0, 0.0], use_DT=True),
}

settings = BvhToDMSettings(bonedict, target_name)
settings.translation_scale = 1.2
settings.output_path = mydir
converter = BvhToDM(settings)
converter.run()