#
# Copyright (C) 2019 Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG),
# acting on behalf of its Max Planck Institute for Intelligent Systems and the
# Max Planck Institute for Biological Cybernetics. All rights reserved.
#
# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is holder of all proprietary rights on this computer program.
# You can only use this computer program if you have closed a license agreement with MPG or you get the right to use the computer program from someone who is authorized to grant you that right.
# Any use of the computer program without a valid license is prohibited and liable to prosecution.
# Contact: ps-license@tuebingen.mpg.de
#

#
# Usage:
#   + Blender Python integration: "Run Script"
#
# Known issues:
#   + When avatar shape is baked all pose shapes will also be removed.
#
# Requirements:
#   + Config files (female/male regressor + body definition) must be located in same folder as .py/.blend file
#   + Blender 2.79 Version 2018-04-10 or later for proper alphabetical shape key name import from FBX
#
# Version: 20180427
#

import argparse
import bpy
import json
from mathutils import Vector
import numpy as np
import os
import sys

############################################################
# Variables
############################################################

bakeShape = True
removePoseShapeKeys = True
exportToFBX = True

femaleJointRegressorName = "betas2Jnt_unity_f.json"
maleJointRegressorName = "betas2Jnt_unity_m.json"
bodyConfigName = "body.json"
fbxName = "SMPL-Blender.fbx"

# Dictionary for matching joint index to name
jointNames = {
    0:  'Pelvis',
    1:  'L_Hip',        4:  'L_Knee',            7:  'L_Ankle',           10: 'L_Foot',
    2:  'R_Hip',        5:  'R_Knee',            8:  'R_Ankle',           11: 'R_Foot',
    3:  'Spine1',       6:  'Spine2',            9:  'Spine3',            12: 'Neck',            15: 'Head',
    13: 'L_Collar',     16: 'L_Shoulder',       18: 'L_Elbow',            20: 'L_Wrist',         22: 'L_Hand',
    14: 'R_Collar',     17: 'R_Shoulder',       19: 'R_Elbow',            21: 'R_Wrist',         23: 'R_Hand',
}

############################################################
# Returns offset between lowest point on feet and floor
############################################################
def floorOffset(skinnedMesh):
    minZ = 999999

    # Important: Update scene to ensure that recent modifications of skin location are taken into account
    bpy.context.scene.update()

    bakedMesh = skinnedMesh.to_mesh(scene=bpy.context.scene, apply_modifiers=True, settings='PREVIEW')
    bakedMesh.transform(skinnedMesh.matrix_world)

    for vertex in bakedMesh.vertices:
        if vertex.co[2] < minZ:
            minZ = vertex.co[2]
    bpy.data.meshes.remove(bakedMesh)

    return minZ

############################################################
# Create avatar and optionally export to FBX
############################################################
def createAvatar():

    # TODO: parse command-line

    print("########################################")
    print("# Creating avatar")
    print("########################################\n")

    # Setup config file paths
    scriptDir = os.path.split(__file__)[0]

    # Fix path if we are running script from within Blender to point to parent directory of .blend file
    if scriptDir.endswith('.blend'):
        scriptDir = os.path.split(scriptDir)[0]

    configDir = scriptDir
    print("Config directory:", configDir)

    femaleJointRegressorPath = os.path.join(configDir, femaleJointRegressorName)
    maleJointRegressorPath = os.path.join(configDir, maleJointRegressorName)
    bodyConfigPath = os.path.join(configDir, bodyConfigName)
    fbxPath = os.path.join(configDir, fbxName)

    ############################################################
    # Load shape parameters and gender from local file
    ############################################################

    print("--------------------------------------------------")
    print("Loading shape parameters:", bodyConfigPath)

    if not os.path.exists(bodyConfigPath):
        print("ERROR: Missing body configuration: ", bodyConfigPath)
        return False

    bodyConfigData = json.load(open(bodyConfigPath))

    gender = bodyConfigData['gender']
    betaRange = bodyConfigData['betaRange']
    betas = np.array(bodyConfigData['betas'])

    print("Gender: " + gender)
    print("Beta range: [-%dSD, +%dSD]" % (betaRange, betaRange))

    np.set_printoptions(formatter={'float': '{: 0.3f}'.format})
    print("Betas: " + np.array2string(betas, separator=','))

    ############################################################
    # Setup gender
    ############################################################

    if gender == 'female':
        avatarName = 'Female'
        armatureName = 'ArmatureFemale'
        genderPrefix = 'f_avg'
        jointRegressorPath = femaleJointRegressorPath
    else:
        avatarName = 'Male'
        armatureName = 'ArmatureMale'
        genderPrefix = 'm_avg'
        jointRegressorPath = maleJointRegressorPath

    ############################################################
    # Load joint regressor
    ############################################################

    print("--------------------------------------------------")
    print("Loading joint regressor")


    if not os.path.exists(jointRegressorPath):
        print("ERROR: Missing joint regressor: ", jointRegressorPath)
        return False

    jointRegressorData = json.load(open(jointRegressorPath))
    jointFromBetaMatrix = np.array(jointRegressorData['betasJ_regr'])
    jointTemplate = np.array(jointRegressorData['template_J'])

    ############################################################
    # Apply shape key weights
    ############################################################
    bpy.ops.object.mode_set(mode='OBJECT')
    bpy.context.scene.objects.active = bpy.data.objects[genderPrefix]

    for beta in range(0, 10):
        # Map beta range to [0, 1]
        weight = betas[beta] / betaRange

        shapenamePos = "Shape%03d_pos" % (beta)
        shapenameNeg = "Shape%03d_neg" % (beta)

        if weight >= 0.0:
            shapename = shapenamePos
            shapenameReset = shapenameNeg
        else:
            shapename = shapenameNeg
            shapenameReset = shapenamePos

        bpy.context.object.data.shape_keys.key_blocks[shapename].value = abs(weight)
        bpy.context.object.data.shape_keys.key_blocks[shapenameReset].value = 0.0

    ############################################################
    # Determine new floor offset
    ############################################################

    # If the script was run before it introduced Armature location offset to
    # put the feet on the ground.
    # Reset any offset on skinned mesh during floor offset calculation so that we have proper default positions.
    skinnedMesh = bpy.data.objects[genderPrefix]
    skinnedMesh.location = Vector((0.0, 0.0, 0.0))
    feetFloorOffsetM = floorOffset(skinnedMesh)

    ############################################################
    # Store default bone positions
    ############################################################

    # We must select Armature as current selection before switching to Edit Mode
    bpy.context.scene.objects.active = bpy.data.objects[avatarName]
    bpy.ops.object.mode_set(mode='EDIT')

    defaultBonePositions = {}
    for bone in bpy.data.armatures[armatureName].edit_bones:
        defaultBonePositions[bone.name] = bone.head

    ############################################################
    # Calculate new bone world positions
    ############################################################
    newJoints = jointFromBetaMatrix.dot(betas)
    newJoints = np.add(newJoints, jointTemplate)

    # Convert regressor [m] units to [cm] for Blender joint positions
    newJoints *= 100

    ############################################################
    # Set new bone world positions
    ############################################################
    print("--------------------------------------------------")
    print("Set new bone positions")

    numJoints = newJoints.shape[0]
    for i in range(0, numJoints):
        boneName = genderPrefix + "_" + jointNames[i]
        newPosition = Vector(newJoints[i])
        oldPosition = defaultBonePositions[boneName]
        offset = newPosition - oldPosition

        # Translate bone (head and tail) to new target position
        bpy.data.armatures[armatureName].edit_bones[boneName].translate(offset)

        # Apply previously calculated floor offset to position feet on floor
        bpy.data.armatures[armatureName].edit_bones[boneName].translate(Vector((0.0, -feetFloorOffsetM*100, 0.0)))

    ############################################################
    # Place skinned mesh on ground
    ############################################################
    bpy.ops.object.mode_set(mode='OBJECT')
    skinnedMesh.location[1] = -feetFloorOffsetM * 100

    ############################################################
    # Remove pose shape keys (optional)
    ############################################################
    if removePoseShapeKeys:
        print("--------------------------------------------------")
        print("Performance optimization: Removing pose shape keys")
        bpy.context.scene.objects.active = bpy.data.objects[genderPrefix]
        numShapeKeys = len(bpy.context.object.data.shape_keys.key_blocks.keys())

        currentShapeKeyIndex = 0
        for index in range(0, numShapeKeys):
            bpy.context.object.active_shape_key_index = currentShapeKeyIndex

            if bpy.context.object.active_shape_key is not None:
                if bpy.context.object.active_shape_key.name.startswith('Pose'):
                    bpy.ops.object.shape_key_remove(all=False)
                else:
                    currentShapeKeyIndex = currentShapeKeyIndex + 1

    ############################################################
    # Bake shape by removing the Shape shape keys (optional)
    ############################################################
    if bakeShape:
        print("--------------------------------------------------")
        print("Bake body shape (removing all shape keys)")
        bpy.context.scene.objects.active = bpy.data.objects[genderPrefix]

        # Create shape mix for current shape
        bpy.ops.object.shape_key_add(from_mix=True)
        numShapeKeys = len(bpy.context.object.data.shape_keys.key_blocks.keys())

        #FIXME: Do not remove pose shapes if they still exist

        # Delete all shape keys except newly added one
        bpy.context.object.active_shape_key_index = 0
        for count in range(0, numShapeKeys):
            bpy.ops.object.shape_key_remove(all=False)

    ############################################################
    # Export to FBX (optional)
    ############################################################
    if exportToFBX:

        print("--------------------------------------------------")
        print("Exporting to FBX:", fbxPath)

        # Deselect all selected objects
        for obj in bpy.context.selected_objects:
            obj.select = False

        # Select bones and skin
        bpy.data.objects[avatarName].select = True
        bpy.data.objects[genderPrefix].select = True

        bpy.ops.export_scene.fbx(filepath=fbxPath, use_selection=True)
    else:
        print("--------------------------------------------------")
        print("FBX export disabled")

    return True

###############################################################################
# Main
###############################################################################

# Parse command line arguments for Python script
argv = sys.argv

if "--" not in argv:
    argv = []  # as if no args are passed
else:
    argv = argv[argv.index("--") + 1:]  # get all args after "--"

if len(argv) > 0:
    # Note: \n characters are not processed in epilog string
    parser = argparse.ArgumentParser(description='Create and export FBX from SMPL template FBX and body.json description.',
                                     epilog = 'Example: blender.exe createSMPLAvatar.blend --background --python createSMPLAavatar.py -- --body="body.json" --name="SMPL-Blender.fbx"')
    parser.add_argument('--body', dest="body", type=str, help='Body JSON file name')
    parser.add_argument('--name', dest="name", type=str, help='Output file name')
    parser.add_argument('--noposeshapes', dest="noposeshapes", help='Remove all pose shape keys', action="store_true")
    parser.add_argument('--nobakeshape', dest="nobakeshape", help='Keep all shape keys for the body shape', action="store_true")
    parser.add_argument('--test', dest="test", help='Do not create output file', action="store_true")

    args = parser.parse_args(argv)

    if args.name is not None:
        fbxName = args.name

    if args.body is not None:
        bodyConfigName = args.body

    if args.noposeshapes:
        removePoseShapeKeys = True

    if args.nobakeshape:
        bakeShape = False

    if args.test:
        exportToFBX = False

success = createAvatar()
if success:
    print("--------------------------------------------------")
    print("Finished\n")