# export_anim_bvh_a3.py

bl_addon_info = {
    "name": "Export BVH (.bvh)",
    "author": "Kenji Saito",
    "version": (0, 1),
    "blender": (2, 5, 6),
    "api": 33047,
    "location": "File > Export",
    "description": "Export Animation to .bvh",
    "warning": "This is a test release.",
    "wiki_url": "http://wiki.blender.org/index.php/Extensions:2.5/Py/"\
        "Scripts/Import-Export/BVH_Exporter_A3",
    "tracker_url": "https://projects.blender.org/tracker/index.php?"\
        "func=detail&aid=26496&group_id=153&atid=467",
    "category": "Import-Export"}


#========================================
# Move your mouse cursor into this window,
# and Press Alt-P. Then you will find
# "Motion Capture A3 (.bvh)" item in your
# "Export" menu items(it's in File menu).
#
#    AAA    ll     t          PPPPPP
#   A   A    l   tttttt       P     P
#  A     A   l     t          P     P
#  AAAAAAA   l     t    ----  PPPPPP
#  A     A   l     t          P
#  A     A   l     t          P
#  A     A   ll     ttt       P
#
#========================================

import bpy
from bpy.props import *
from mathutils import *

# class for all processing
class ExpBVH:

    def __init__(self,file=None,armature_object=None,scale=1.0,start_frame=1,end_frame=500,scale_data=False):
        self.filePath = file
        self.file = open(file,"w")
        self.scale = scale
        self.start_frame = start_frame
        self.end_frame = end_frame
        self.includeScaleData = scale_data
        self.scn = bpy.context.scene
        self.armObj = armature_object
        self.arm = self.armObj.data
        self.bones = self.arm.bones
        self.pose_bones = self.armObj.pose.bones
        self.select_armature_object()
        self.getRootBones()

    def select_armature_object(self):
        if self.armObj:
            if self.armObj.type=='ARMATURE':
                return
        for o in bpy.context.selected_objects:
            if o.type=='ARMATURE':
                self.armObj = o
                return
        for o in bpy.data.objects:
            if o.type=='ARMATURE':
                self.armObj = o
                return
        print('ERROR: no armature!!!')

    # list all bones which has no parent
    def getRootBones(self):
        self.root_bones = []
        for b in self.bones:
            if (b.parent == None):
                self.root_bones.append(b)

    # determine if the bone needs rot channels
    # (kludgy...)
    def needRotChannels(self,bone):
        actionGroups = self.armObj.animation_data.action.groups.keys()
        if not bone.name in actionGroups:
            return False
        channels = self.armObj.animation_data.action.groups[bone.name].channels
        for c in channels:
            if c.data_path.find('rotation')==-1:
                continue
            else:
                return True
        return False

    # determine if the bone needs loc channels
    # (kludgy...)
    def needLocChannels(self,bone):
        actionGroups = self.armObj.animation_data.action.groups.keys()
        if not bone.name in actionGroups:
            return False
        channels = self.armObj.animation_data.action.groups[bone.name].channels
        for c in channels:
            if c.data_path.find('location')==-1:
                continue
            else:
                return True
        return False

    # determine if the bone needs scale channels
    # (kludgy...)
    def needScaleChannels(self,bone):
        actionGroups = self.armObj.animation_data.action.groups.keys()
        if not bone.name in actionGroups:
            return False
        channels = self.armObj.animation_data.action.groups[bone.name].channels
        for c in channels:
            if c.data_path.find('scale')==-1:
                continue
            else:
                return True
        return False

    # export the Hierarchy part
    def exportHierarchy(self):
        self.file.write('HIERARCHY\n')
        self.file.write('ROOT BSRootBone\n')
        self.file.write('{\n')
        self.file.write('  OFFSET 0.0 0.0 0.0\n')
        self.file.write('  CHANNELS 3 Zrotation Xrotation Yrotation\n')
        for bone in self.root_bones:
            self.exportBlock(bone,1)
        self.file.write('}\n')

    # export the Block
    def exportBlock(self,bone,iLevel):
        indent = ''
        for j in range(iLevel):
            indent = indent + '  '
        self.file.write(indent + 'JOINT ' + bone.name + '\n')
        self.file.write(indent + '{\n')
        self.file.write(indent + '  ' + self.offsetString(bone) + '\n')
        self.file.write(indent + '  ' + self.channelsString(bone) + '\n')
        bs = bone.children
        if (bs == None):
            self.exportEndSite(bone,iLevel+1)
        elif (len(bs) == 0):
            self.exportEndSite(bone,iLevel+1)
        else:
            for b in bs:
                self.exportBlock(b,iLevel+1)
        self.file.write(indent + '}\n')

    # make a OFFSET string
    def offsetString(self,bone):
        parentBone = bone.parent
        if (parentBone == None):
            parentBoneHead = [0,0,0]
        else:
            parentBoneHead = parentBone.head_local
        boneHead = bone.head_local
        offset = [boneHead[i]-parentBoneHead[i] for i in range(3)]
        offset[0] = self.scale * offset[0]
        offset[1] = self.scale * offset[1]
        offset[2] = self.scale * offset[2]
        return 'OFFSET %f %f %f' % (offset[0],offset[1],offset[2])

    # make a CHANNELS string
    def channelsString(self,bone):
    #==============================
        r = self.needRotChannels(bone)
        l = self.needLocChannels(bone)
        s = self.needScaleChannels(bone)
        if not self.includeScaleData:
            s = False
        c = 0
    #    if (r):
        if (True):
            c = c + 3
        if (l):
            c = c + 3
        if (s):
            c = c + 3
        ret = ''
        if (l):
            ret = ret + 'Xposition Yposition Zposition '
    #    if (r):
        if (True):
            ret = ret + 'Zrotation Xrotation Yrotation '
        if (s):
            ret = ret + 'Xscale Yscale Zscale '
        ret = ('CHANNELS %i ' % (c)) + ret
        return ret

    # export a EndSite
    def exportEndSite(self,bone,iLevel):
        boneHead = bone.head_local
        boneTail = bone.tail_local
        offset = [boneTail[j]-boneHead[j] for j in range(3)]
        offset[0] = self.scale * offset[0]
        offset[1] = self.scale * offset[1]
        offset[2] = self.scale * offset[2]
        indent = ''
        for j in range(iLevel):
            indent = indent + '  '
        self.file.write(indent + 'End Site\n')
        self.file.write(indent + '{\n')
        self.file.write(indent + '  OFFSET %f %f %f\n' % (offset[0],offset[1],offset[2]))
        self.file.write(indent + '}\n')

    # export the MOTION part
    def exportMotion(self):
        self.file.write('MOTION\n')
        self.file.write('Frames: ' + str(self.end_frame - self.start_frame +1) + '\n')
        self.file.write('Frame Time: 0.0333\n')

        for f in range(self.start_frame,self.end_frame+1):
            self.scn.frame_set(f)
            self.file.write('0.0 0.0 0.0 ') # <-- for BSRootBone
            for bone in self.root_bones:
                self.exportMotion2(bone)
            self.file.write('\n')

    # search pose_bone correspond to the bone
    # (Is this OK?)
    def getPoseBone(self,bone):
        for pb in self.pose_bones:
            if pb.name == bone.name:
                return pb
        return None

    # actually export the MOTION part. (recursion)
    def exportMotion2(self,bone):
        gMatrix = bone.matrix_local.rotation_part()
        pose_bone = self.getPoseBone(bone)
        if self.needLocChannels(bone):
            loc = pose_bone.location
            loc = loc * gMatrix
            loc = [self.scale*loc[i] for i in range(3)]
            self.file.write('%f %f %f ' % (loc[0],loc[1],loc[2]))
        #if self.needRotChannels(bone):
        if True:
            if pose_bone.rotation_mode=='QUATERNION':
                quat = pose_bone.rotation_quaternion
            elif pose_bone.rotation_mode=='AXIS_ANGLE':
                aa = pose_bone.rotation_axis_angle
                quat = RotationMatrix(aa[0],3,aa[1:]).to_quat()
            else:
                quat = pose_bone.rotation_euler.to_quat()
            v = Vector((quat[1],quat[2],quat[3])) * gMatrix
            quat = Quaternion((quat[0],v[0],v[1],v[2]))
            euler = quat.to_euler('YZX') # why 'YZX' ???
            euler[0] = euler[0]*180.0/3.141592
            euler[1] = euler[1]*180.0/3.141592
            euler[2] = euler[2]*180.0/3.141592
            self.file.write('%f %f %f ' % (euler[2],euler[0],euler[1]))
        if self.needScaleChannels(bone) and self.includeScaleData:
            scale = pose_bone.scale
            #scale = gMatrix * scale
            self.file.write('%f %f %f ' % (scale[0],scale[1],scale[2]))
        bs = bone.children
        if bs:
            for b in bs:
                self.exportMotion2(b)

    # export all
    def export(self):
        self.exportHierarchy()
        self.exportMotion()

def export_bvh(file=None,armature_object=None,scale=1.0,start_frame=1,end_frame=500,scale_data=False):
    eb = ExpBVH(file,armature_object,scale,start_frame,end_frame,scale_data)
    eb.export()
    pass

class BvhExporter(bpy.types.Operator):
    '''Save a Armature Motion to File'''
    bl_idname ="export.bvh"
    bl_label = "Export BVH"

    filepath = StringProperty(name="File Path", description="Filepath used for exporting the BVH file",maxlen=1024,default="",subtype='FILE_PATH')
    armature = StringProperty(name="armature",description="input armature name",default='Armature')
    scale = FloatProperty(name="Scale", description="Scale", default=1.0)
    start_frame = IntProperty(name="Start Frame", description="Start Frame", default=1, min=1)
    end_frame = IntProperty(name="End Frame", description="End Frame", default=35, min=1)
    include_scale_data = BoolProperty(name="Include Scale Data", description="include scale data?", default=False)

    def execute(self, context):
        file = self.properties.filepath
        armature_name = self.properties.armature
        armature_object = bpy.data.objects[armature_name]
        scale = self.properties.scale
        start_frame = self.properties.start_frame
        end_frame = self.properties.end_frame
        scale_data = self.properties.include_scale_data
        export_bvh(file,armature_object,scale,start_frame,end_frame,scale_data)
        return {'FINISHED'}

    def invoke(self, context, event):
        WindowManager = context.window_manager
        WindowManager.fileselect_add(self)
        return {'RUNNING_MODAL'}

def menu_func(self, context):
    oProp = self.layout.operator(BvhExporter.bl_idname, text="Motion Capture A3 (.bvh)")
    oProp.filepath = bpy.data.filepath.replace(".blend", ".bvh")
    oProp.scale = 1.0
    oProp.start_frame = context.scene.frame_start
    oProp.end_frame = context.scene.frame_end
    for o in context.selected_objects:
        if o.type=='ARMATURE':
            oProp.armature = o.name
            return
    for o in bpy.data.objects:
        if o.type=='ARMATURE':
            oProp.armature = o.name
            return

def register():
    bpy.types.INFO_MT_file_export.append(menu_func)

def unregister():
    bpy.types.INFO_MT_file_export.remove(menu_func)

if __name__ == "__main__":
    register()
