from pathlib import Path
import RLPy

# export .obj to which folder
output_dir = "path/to/your/target/folder"
# which arkit bs to export. If empty list is given, all 52 arkit bs will be exported.
arkit_names = []

# ARKit bs mapping to RL CC4 Extended expressions
arkit_mapping = {
    "eyeBlinkLeft": {"Eye_Blink_L": 100},
    "eyeLookDownLeft": {"Eye_L_Look_Down": 100},
    "eyeLookInLeft": {"Eye_L_Look_R": 100},
    "eyeLookOutLeft": {"Eye_L_Look_L": 100},
    "eyeLookUpLeft": {"Eye_L_Look_Up": 100},
    "eyeSquintLeft": {"Eye_Squint_L": 100},
    "eyeWideLeft": {"Eye_Wide_L": 100},
    "eyeBlinkRight": {"Eye_Blink_R": 100},
    "eyeLookDownRight": {"Eye_R_Look_Down": 100},
    "eyeLookInRight": {"Eye_R_Look_L": 100},
    "eyeLookOutRight": {"Eye_R_Look_R": 100},
    "eyeLookUpRight": {"Eye_R_Look_Up": 100},
    "eyeSquintRight": {"Eye_Squint_R": 100},
    "eyeWideRight": {"Eye_Wide_R": 100},
    "jawForward": {"Jaw_Forward": 100},
    "jawLeft": {"Jaw_L": 100},
    "jawRight": {"Jaw_R": 100},
    "jawOpen": {"Jaw_Open": 80},
    "mouthClose": {"Jaw_Open": 60, "Mouth_Close": 50},
    "mouthFunnel": {"Mouth_Funnel_Up_L": 80, "Mouth_Funnel_Up_R": 80, "Mouth_Funnel_Down_L": 100, "Mouth_Funnel_Down_R": 100},
    "mouthPucker": {"Mouth_Pucker_Up_L": 120, "Mouth_Pucker_Up_R": 120, "Mouth_Pucker_Down_L": 120, "Mouth_Pucker_Down_R": 120},
    "mouthLeft": {"Mouth_L": 100},
    "mouthRight": {"Mouth_R": 100},
    "mouthSmileLeft": {"Mouth_Smile_L": 100},
    "mouthSmileRight": {"Mouth_Smile_R": 100},
    "mouthFrownLeft": {"Mouth_Frown_L": 80},
    "mouthFrownRight": {"Mouth_Frown_R": 80},
    "mouthDimpleLeft": {"Mouth_Dimple_L": 120},
    "mouthDimpleRight": {"Mouth_Dimple_R": 120},
    "mouthStretchLeft": {"Mouth_Stretch_L": 100},
    "mouthStretchRight": {"Mouth_Stretch_R": 100},
    "mouthRollLower": {"Mouth_Roll_In_Lower_L": 80, "Mouth_Roll_In_Lower_R": 80},
    "mouthRollUpper": {"Mouth_Roll_In_Upper_L": 80, "Mouth_Roll_In_Upper_R": 80, "Mouth_Frown_L": 20, "Mouth_Frown_R": 20},
    "mouthShrugLower": {"Mouth_Shrug_Lower": 100},
    "mouthShrugUpper": {"Mouth_Shrug_Upper": 120},
    "mouthPressLeft": {"Mouth_Press_L": 120},
    "mouthPressRight": {"Mouth_Press_R": 120},
    "mouthLowerDownLeft": {"Mouth_Down_Lower_L": 100},
    "mouthLowerDownRight": {"Mouth_Down_Lower_R": 100},
    "mouthUpperUpLeft": {"Mouth_Up_Upper_L": 100},
    "mouthUpperUpRight": {"Mouth_Up_Upper_R": 100},
    "browDownLeft": {"Brow_Drop_L": 120},
    "browDownRight": {"Brow_Drop_R": 120},
    "browInnerUp": {"Brow_Raise_Inner_L": 100, "Brow_Raise_Inner_R": 100},
    "browOuterUpLeft": {"Brow_Raise_Outer_L": 100},
    "browOuterUpRight": {"Brow_Raise_Outer_R": 100},
    "cheekPuff": {"Cheek_Puff_L": 100, "Cheek_Puff_R": 100},
    "cheekSquintLeft": {"Cheek_Raise_L": 120},
    "cheekSquintRight": {"Cheek_Raise_R": 120},
    "noseSneerLeft": {"Nose_Sneer_L": 120},
    "noseSneerRight": {"Nose_Sneer_R": 120},
    "tongueOut": {"Tongue_Out": 80}
}


def export_arkit_objs(output_dir, arkit_names=[]):
    output_path = Path(output_dir)

    # verifiy what arkit bs need to export
    user_arkit_names = list(arkit_mapping.keys())
    if arkit_names:
        user_arkit_names = list(set(arkit_names) & set(user_arkit_names))
    print(f"+ arkit to export: {user_arkit_names}")

    # get the avatar and face component
    avatar_list = RLPy.RScene.GetAvatars()
    avatar = avatar_list[0]
    face_component = avatar.GetFaceComponent()

    # RL export optioins
    export_options = RLPy.EExport3DFileOption_FullBodyPart
    export_options += RLPy.EExport3DFileOption_AxisYUp
    export_options += RLPy.EExport3DFileOption_ExportFacialAnimation

    for bs in user_arkit_names:
        print(f"+ start processing: {bs}...")
        expr_names = list(arkit_mapping[bs].keys())
        strength_keys = list(arkit_mapping[bs].values())
        print(f"- RL expressions: {arkit_mapping[bs]}")

        normalized_keys = [i/100. for i in strength_keys[:]]

        # set strenght values for RL expressions
        current_time = RLPy.RTime.FromValue(0)
        time_inv = RLPy.RTime.FromValue(0)
        face_component.BeginKeyEditing()
        face_component.AddExpressionKeys(current_time, expr_names, normalized_keys, time_inv)
        face_component.EndKeyEditing()

        # export .obj
        obj_path = output_path.joinpath(f"{bs}.obj")
        RLPy.RFileIO.ExportObjFile(avatar, obj_path.as_posix(), export_options)

        # reset expressions
        strength_keys = [0.] * len(expr_names)
        face_component.BeginKeyEditing()
        face_component.AddExpressionKeys(current_time, expr_names, strength_keys, time_inv)
        face_component.EndKeyEditing()

# main func
export_arkit_objs(output_dir, arkit_names)
print("+ export tasks finished.")