#!BPY

"""
Name: 'IK Capture (.bvh)...'
Blender: 237
Group: 'Export'
Tip: 'Export a (.bvh) IK capture file'
"""

__author__ = "Jean-Baptiste PERIN, Vincent BILLET"
__url__ = ("blender", "elysiun")
__version__ = "0.4 05/12/01"
__email__= ["M. PERIN Jean-Baptiste, jb_perin@yahoo.fr", "scripts"]
__bpydoc__ = """\
This script exports Empty animation to bvh file format.<br>

Missing:<br>

Known issues:<br>

Notes:<br>
jb_perin@yahoo.fr and vbillet@yahoo.fr
"""

# -------------------------------------------------------------------------- 
# IKBaker.py 
# -------------------------------------------------------------------------- 
# ***** BEGIN GPL LICENSE BLOCK ***** 
# 
# This program is free software; you can redistribute it and/or 
# modify it under the terms of the GNU General Public License 
# as published by the Free Software Foundation; either version 2 
# of the License, or (at your option) any later version. 
# 
# This program is distributed in the hope that it will be useful, 
# but WITHOUT ANY WARRANTY; without even the implied warranty of 
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the 
# GNU General Public License for more details. 
# 
# You should have received a copy of the GNU General Public License 
# along with this program; if not, write to the Free Software Foundation, 
# Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA. 
# 
# ***** END GPL LICENCE BLOCK ***** 
# -------------------------------------------------------------------------- 


import Blender
import math
empty_rest_ori = {}




def getRootBone(armature):
	rootBone = None
	for bone in armature.bones.values():
		if bone.parent == None :
			if rootBone == None:
				rootBone = bone
			else:
				print"Multiple root armature not handled"
	return rootBone

def getBoneCurPos(armature, bone):
	lempty = Blender.Object.Get(bone.name)
	return (lempty.getLocation())

def getBoneCurRot(armature, bone):
	lempty = Blender.Object.Get(bone.name)
	return (lempty.getMatrix('worldspace'))

def getBoneRestPos(armature, bone):
	return (bone.head)

def getBoneRestRot(armature, bone):
	lempty = Blender.Object.Get(bone.name)
	return (bone.matrix)
    
def indentString(i):
    string = ""
    for v in range(0,i):
        string += "\t"
    return string

def printBoneHierarchie (armature, bone, level):
	str = "%sJOINT %s\n"%(indentString(level), bone.name)
	str += "%s{\n"%indentString(level)
	str += "%s\tOFFSET %f %f %f\n"%(indentString(level),
                                     getBoneRestPos(armature, bone)['ARMATURESPACE'][0] - getBoneRestPos(armature, bone.parent)['ARMATURESPACE'][0],
                                    getBoneRestPos(armature, bone)['ARMATURESPACE'][1]- getBoneRestPos(armature, bone.parent)['ARMATURESPACE'][1],
                                    getBoneRestPos(armature, bone)['ARMATURESPACE'][2]- getBoneRestPos(armature, bone.parent)['ARMATURESPACE'][2])
	empty_rest_ori [bone.name] = getBoneRestRot(armature, bone)['ARMATURESPACE']
	print empty_rest_ori [bone.name]
	str += "%s\tCHANNELS 3 Zrotation Xrotation Yrotation\n"%indentString(level)
	if bone.children != None:
            for b in bone.children:
                if bone.name == b.parent.name:
                    str += printBoneHierarchie (armature,b, level+1)
        else:
            str += "%s\tEnd Site\n"%indentString(level)
            str += "%s\t{\n"%indentString(level)
            str += "%s\t\tOFFSET %f %f %f\n"%(indentString(level),
                                           bone.tail['ARMATURESPACE'][0]-bone.head['ARMATURESPACE'][0],
                                           bone.tail['ARMATURESPACE'][1]-bone.head['ARMATURESPACE'][1],
                                           bone.tail['ARMATURESPACE'][2]-bone.head['ARMATURESPACE'][2])
            str += "%s\t}\n"%indentString(level)
            
	str += "%s}\n"%indentString(level)
	return (str)
	
def printArmatureHierarchie(armature):
	rootBone = getRootBone(armature)
	str = "ROOT %s\n{\n"%rootBone.name
	str += "\tOFFSET %f %f %f\n"%(getBoneRestPos(armature, rootBone)['ARMATURESPACE'][0],
                                    getBoneRestPos(armature, rootBone)['ARMATURESPACE'][1],
                                    getBoneRestPos(armature, rootBone)['ARMATURESPACE'][2])
	str += "\tCHANNELS 6 Xposition Yposition Zposition Zrotation Xrotation Yrotation\n"
	empty_rest_ori [rootBone.name] = getBoneRestRot(armature, rootBone)['ARMATURESPACE']
	for b in rootBone.children:
            if rootBone.name == b.parent.name:
                str += printBoneHierarchie (armature,b,1)
	str += "}\n"
	return (str)
    
def cumulatedAngle(unbone):
    lunbone = unbone
    angle = Blender.Mathutils.Euler(0.0,0.0,0.0)
    if lunbone.parent != None:
        print "coucou", (lunbone.parent)
    while lunbone.parent != None:
        print "dans la boucle"
        angle += getBoneCurRot(armature, lunbone.parent).rotationPart().toEuler()
        print "angle = ",angle
        lunbone = lunbone.parent
    #angle -=  getBoneCurRot(armature, unbone).rotationPart().toEuler()
    return angle

def cumulatedAngle2(unbone):
    lunbone = unbone
    angle = Blender.Mathutils.Euler(0.0,0.0,0.0)
    print (lunbone.parent)
    while lunbone.parent == None:
        print "dans la boucle"
        angle += getBoneCurRot(armature, lunbone.parent).rotationPart().toEuler()
        print "angle = ",angle
        lunbone = lunbone.parent
    #angle -=  getBoneCurRot(armature, unbone).rotationPart().toEuler()
    return angle

def printBoneFrame (armature, bone, level):
    
##    mat1inv = Blender.Mathutils.Matrix(empty_rest_ori [bone.name][0],
##                                    empty_rest_ori [bone.name][1],
##                                    empty_rest_ori [bone.name][2],
##                                    empty_rest_ori [bone.name][3])
##    mat1inv.invert()
##    mat2 = Blender.Mathutils.Matrix(getBoneCurRot(armature, bone)[0],
##                                    getBoneCurRot(armature, bone)[1],
##                                    getBoneCurRot(armature, bone)[2],
##                                    getBoneCurRot(armature, bone)[3])
##    mat3=mat1inv*mat2
##    eul = mat3.rotationPart().toEuler()
##    str = "%f %f %f "%(
##        eul[2],
##        eul[0],
##        eul[1])
##
    
    #eul_parent = getBoneCurRot(armature, bone.parent).rotationPart().toEuler()
    print "Trrace ", bone, " child of", bone.parent #, " ", cumulatedAngle(bone)
    unbone = bone
    angle = Blender.Mathutils.Euler(0.0,0.0,0.0)
    print (unbone.parent)
    ok = True 
    while ok == True:
        if unbone.parent == None:
            ok = False
        else: 
            print "dans la boucle"
            angle[0] += getBoneCurRot(armature, unbone.parent).rotationPart().toEuler()[0]
            angle[1] += getBoneCurRot(armature, unbone.parent).rotationPart().toEuler()[1]
            angle[2] += getBoneCurRot(armature, unbone.parent).rotationPart().toEuler()[2]
            print "angle = ",angle

            unbone = unbone.parent
        

    
    z = (getBoneCurRot(armature, bone).rotationPart().toEuler()[2]) - (empty_rest_ori [bone.name].rotationPart().toEuler()[2])
    x = getBoneCurRot(armature, bone).rotationPart().toEuler()[0] - empty_rest_ori [bone.name].rotationPart().toEuler()[0]
    y = getBoneCurRot(armature, bone).rotationPart().toEuler()[1] - empty_rest_ori [bone.name].rotationPart().toEuler()[1]

    str = "%f %f %f "%(
        z-angle[2],
        x-angle[0],
        y-angle[1])


    
    if bone.children != None:
        for b in bone.children:
            if bone.name == b.parent.name:
                str += printBoneFrame (printBoneFrame,b, level+1)
    return (str)


def printArmatureFrame (armature):
    rootBone = getRootBone(armature)
    str = "%f %f %f %f %f %f "%(
        getBoneCurPos(armature, getRootBone(armature))[0],
        getBoneCurPos(armature, getRootBone(armature))[1],
        getBoneCurPos(armature, getRootBone(armature))[2],
        getBoneCurRot(armature, getRootBone(armature)).rotationPart().toEuler()[2],
        getBoneCurRot(armature, getRootBone(armature)).rotationPart().toEuler()[0],
        getBoneCurRot(armature, getRootBone(armature)).rotationPart().toEuler()[1]
       )	
    for b in rootBone.children:
        if rootBone.name == b.parent.name:
            str += printBoneFrame (armature,b,1)
    str += "\n"
    return (str)
		
def exportBVH(filename):
    global armaturename
    print "Exporting armature %s to file %s "%( armaturename, filename)
    startFrame = Blender.Get('staframe')
    endFrame = Blender.Get('endframe')

    outputfile = file(filename,'w')

    outputfile.write("HIERARCHY\n")

    armData = Blender.Object.Get(armaturename).getData()

    str_res = printArmatureHierarchie (armData)
    outputfile.write(str_res)
    outputfile.write("MOTION\n")
    outputfile.write("Frames: %d\n"%(endFrame-startFrame))
    outputfile.write("Frame Time: %f\n"%(1.0/Blender.Scene.getCurrent().getRenderingContext().framesPerSec()))

    ctx = Blender.Scene.getCurrent().getRenderingContext()
    Blender.Set('curframe', startFrame)
    while Blender.Get('curframe') < endFrame:
        print "frame : %d"%(Blender.Get('curframe'))
        str_res = printArmatureFrame(armData)
        outputfile.write(str_res)
        Blender.Set('curframe',Blender.Get('curframe')+1)

    outputfile.close()

    Blender.Set('curframe', startFrame)

    Blender.Redraw()

print "**** START ****"

bfilename = Blender.Get('filename')
armaturename = ""
nfilename = bfilename.replace('.blend', '.bvh')
print nfilename
obj_sel = Blender.Object.GetSelected()
if (obj_sel != None) and (len (obj_sel) == 1) and (type(obj_sel[0].getData()) == Blender.Types.ArmatureType):
    print obj_sel[0].getName()
    print type(obj_sel[0].getData())
    armaturename = obj_sel[0].getName()
    Blender.Window.FileSelector(exportBVH, 'Export BVH',nfilename)
else:
     error_txt = "Error%t|Select one and only one armature before running this script"
     Blender.Draw.PupMenu(error_txt)
     print error_txt

print "**** END ****"