#!BPY
# -*- encoding: utf-8 -*-
""" Registration info for Blender menus:
Name: 'BioVision Hierarchy (.bvh)...'
Blender: 242
Group: 'Export'
Tooltip: 'Export to BVH file (.bvh)'
"""

#==============================
# To execute bs.py, move your mouse cursor
# into this window. And push ...
#
#    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
#
#==============================

__author__ = ("Kenji Saito")
__url__ = ["http://acerola3d.sourceforge.jp/bs/"]
__email__ = ["ksaito.sourceforge.jp@gmail.com"]
__version__ = "2010/08/12"
__bpydoc__ = """\
This script exports to BVH format.
"""

#==============================
# BriskSkeletons v0.9
# http://acerola3d.sourceforge.jp/bs/
#==============================
# BriskSkeletons is Copyright (C) 2004-2010 Kenji Saito.
#
#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.
#
#A copy of the GNU General Public License is also available at
#<URL:http://www.gnu.org/copyleft/gpl.html>. You may also obtain
#it by writing to the Free Software Foundation, Inc., 59 Temple
#Place - Suite 330, Boston, MA 02111-1307, USA. 
#
#Kenji Saito <ksaito.sourceforge.jp@gmail.com>
#==============================

import Blender
import math
from Blender import Ipo, Draw, BGL, Mathutils

outputFile = "output.bvh"
scale = 1.0
startFrame = None
endFrame = None
armatures = None
targetArmature = None
action = None

EVENT_SELECT_ARMATURE = 1
EVENT_CHANGE_SCALE = 2
EVENT_CHANGE_START_FRAME = 3
EVENT_CHANGE_END_FRAME = 4
EVENT_CHANGE_FILENAME = 5
EVENT_EXPORT = 6
EVENT_EXIT = 7

guiArmatureMenu = None
guiScaleNumber = None
guiStartFrameNumber = None
guiEndFrameNumber = None
guiFileNameString = None
message = "status: READY"


#============================================================
# INIT
#============================================================
#==============================
def init():
#==============================
    global startFrame,endFrame,armatures,targetArmature,action,outputFile
    scene = Blender.Scene.GetCurrent()
    context = scene.getRenderingContext()
    startFrame = context.startFrame()
    endFrame = context.endFrame()
    #
    targetArmature = None
    aDic = Blender.Armature.Get()
    armatures = []
    for v in aDic.values():
        armatures.append(v)
    for a in armatures:
        if (targetArmature == None):
            targetArmature = a
            action = searchAction(targetArmature)
            break
    if (targetArmature != None):
        o = searchObject(targetArmature)
        outputFile = o.getName() + '.bvh'

def searchAction(armature):
    for o in Blender.Object.Get():
        if o.getData() == armature:
            return o.getAction()
    return None

def searchObject(armature):
    for o in Blender.Object.Get():
        if o.getData() == armature:
            return o
    return None

#==============================
def changeParam():
#==============================
    global outputFile, scale, startFrame, endFrame
    global armatures, targetArmature, action
    menuitems = 'Select the armature.%t'
    for a in armatures:
        o = searchObject(a)
        menuitems = menuitems + '|' + o.getName()
        #menuitems = menuitems + '|' + a.name
    result = Blender.Draw.PupMenu(menuitems)
    if result == -1:
        return None
    targetArmature = armatures[result-1]
    action = searchAction(targetArmature)
    o = searchObject(targetArmature)
    outputFile = o.getName() + '.bvh'

    f = Blender.Draw.Create(1.0)
    block = []
    block.append(("Value: ",f,0.0,100.0))
    result = Blender.Draw.PupBlock("PupBlock test", block)

    #outputFile = None
    #
    #if not outputFile:
    #    return None

    return True

#============================================================
# EXPORT BVH
#============================================================
#==============================
def getRootBones():
#==============================
    global targetArmature
    rootBones = []
    bs = targetArmature.bones
    for b in bs.values():
        if (b.parent == None):
            rootBones.append(b)
    return rootBones

#==============================
def exportBVH(filename):
#==============================
    global outputFile
    outputFile = filename
    file = open(outputFile,'w')
    rootBones = getRootBones()
    printBVH(rootBones,file)
    file.close()

#==============================
#vecをquatで回転した結果を返す
def rotate(vec,quat):
#==============================
    vQuat = Mathutils.Quaternion([0.0,vec[0],vec[1],vec[2]])
    cquat = Mathutils.Quaternion(quat.w,quat.x,quat.y,quat.z)
    cquat.conjugate()
    a = Mathutils.CrossQuats(quat,vQuat)
    b = Mathutils.CrossQuats(a,cquat)
    return Mathutils.Vector([b[1],b[2],b[3]])

#==============================
def needRotChannels(bone):
#==============================
    global action
    ipo = action.getChannelIpo(bone.name)
    ipoCurve = ipo[Ipo.PO_QUATW]
    if (ipoCurve != None):
        return True
    else:
        return False

#==============================
def needLocChannels(bone):
#==============================
    global action
    ipo = action.getChannelIpo(bone.name)
    ipoCurve = ipo[Ipo.PO_LOCX]
    if (ipoCurve != None):
        return True
    else:
        return False

#==============================
def needScaleChannels(bone):
#==============================
    global action
    ipo = action.getChannelIpo(bone.name)
    ipoCurve = ipo[Ipo.PO_SCALEX]
    if (ipoCurve != None):
        return True
    else:
        return False

#==============================
def printBVH(rootBones,file):
#==============================
    printHierarchy(rootBones,file)
    printMotion(rootBones,file)

#==============================
def printHierarchy(rootBones,file):
#==============================
    file.write('HIERARCHY\n')
    file.write('ROOT BSRootBone\n')
    file.write('{\n')
    file.write('  OFFSET 0.0 0.0 0.0\n')
    file.write('  CHANNELS 3 Zrotation Xrotation Yrotation\n')
    for bone in rootBones:
        printBlock(bone,1,file)
    file.write('}\n')

#==============================
def printBlock(bone,i,file):
#==============================
    indent = ''
    for j in range(i):
        indent = indent + '  '
    file.write(indent + 'JOINT ' + bone.name + '\n')
    file.write(indent + '{\n')
    file.write(indent + '  ' + offsetString(bone) + '\n')
    file.write(indent + '  ' + channelsString(bone) + '\n')
    bs = bone.children
    if (bs == None):
        printEndSite(bone,i+1,file)
    elif (len(bs) == 0):
        printEndSite(bone,i+1,file)
    else:
        for b in bs:
            printBlock(b,i+1,file)
    file.write(indent + '}\n')

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

#==============================
def channelsString(bone):
#==============================
    r = needRotChannels(bone)
    l = needLocChannels(bone)
    s = needScaleChannels(bone)
    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

#==============================
def printEndSite(bone,i,file):
#==============================
    global scale
    boneHead = bone.head['ARMATURESPACE']
    boneTail = bone.tail['ARMATURESPACE']
    offset = [boneTail[j]-boneHead[j] for j in range(3)]
    offset[0] = scale * offset[0]
    offset[1] = scale * offset[1]
    offset[2] = scale * offset[2]
    indent = ''
    for j in range(i):
        indent = indent + '  '
    file.write(indent + 'End Site\n')
    file.write(indent + '{\n')
    file.write(indent + '  OFFSET %f %f %f\n' % (offset[0],offset[1],offset[2]))
    file.write(indent + '}\n')

#==============================
def printMotion(rootBones,file):
#==============================
    global startFrame,endFrame
    file.write('MOTION\n')
    file.write('Frames: ' + str(endFrame - startFrame +1) + '\n')
    file.write('Frame Time: 0.0333\n')

    for f in range(startFrame,endFrame+1):
        file.write('0.0 0.0 0.0 ') # <-- for BSRootBone
        for bone in rootBones:
            printMotion2(bone,f,file)
        file.write('\n')

#==============================
def printMotion2(bone,frame,file):
#==============================
    global scale
    gQuat = bone.matrix['ARMATURESPACE'].toQuat()
    if (needLocChannels(bone)):
        loc = getLoc(bone,frame)
        loc = rotate(loc,gQuat)
        loc = [scale*loc[i] for i in range(3)]
        file.write('%f %f %f ' % (loc[0],loc[1],loc[2]))
#    if (needRotChannels(bone)):
    if (True):
        quat = getQuat(bone,frame)
        quat = rotateQuat(quat,gQuat)
        [eX,eY,eZ] = quat.toEuler()
        eX, eY, eZ = ZYXToZXY(eX, eY, eZ)
        file.write('%f %f %f ' % (eZ,eX,eY))
    if (needScaleChannels(bone)):
        lScale = getSize(bone,frame)
        #lScale = rotate(lScale,gQuat)
        file.write('%f %f %f ' % (lScale[0],lScale[1],lScale[2]))
    bs = bone.children
    if bs:
        for b in bs:
            printMotion2(b,frame,file)

#==============================
def getLoc(bone,frame):
#==============================
    global action
    ipo = action.getChannelIpo(bone.name)
    cx = ipo[Ipo.PO_LOCX]
    cy = ipo[Ipo.PO_LOCY]
    cz = ipo[Ipo.PO_LOCZ]
    x = cx.evaluate(frame)
    y = cy.evaluate(frame)
    z = cz.evaluate(frame)
    return Mathutils.Vector(x,y,z)

#==============================
def getQuat(bone,frame):
#==============================
    global action
    ipo = action.getChannelIpo(bone.name)
    cw = ipo[Ipo.PO_QUATW]
    cx = ipo[Ipo.PO_QUATX]
    cy = ipo[Ipo.PO_QUATY]
    cz = ipo[Ipo.PO_QUATZ]
    if (cw == None or cx == None or cy == None or cz == None):
        w = 1.0
        x = y = z = 0.0
    else:
        w = cw.evaluate(frame)
        x = cx.evaluate(frame)
        y = cy.evaluate(frame)
        z = cz.evaluate(frame)
    return Mathutils.Quaternion(w,x,y,z)

#==============================
def getSize(bone,frame):
#==============================
    global action
    ipo = action.getChannelIpo(bone.name)
    cx = ipo[Ipo.PO_SCALEX]
    cy = ipo[Ipo.PO_SCALEY]
    cz = ipo[Ipo.PO_SCALEZ]
    x = cx.evaluate(frame)
    y = cy.evaluate(frame)
    z = cz.evaluate(frame)
    return Mathutils.Vector(x,y,z)

#==============================
#回転の回転。回転の合成とは違うことに注意。
#座標系Aと座標系Bがあり、座標系Bで計測された座標を
#座標系Aの座標に変換してくれる回転をqbaとする。
#このとき座標系Bで計測された回転(qとする)を
#座標系Aにおける回転に変換してくれる関数
def rotateQuat(q,qba):
#==============================
    qv1 = Mathutils.Vector([q.x,q.y,q.z])
    qv2 = rotate(qv1,qba)
    return Mathutils.Quaternion([q.w,qv2[0],qv2[1],qv2[2]])

#==============================
# bvh_export.pyから借用して手を加えました．
def ZYXToZXY(x, y, z):
#==============================
    RAD_TO_DEG = 180.0/3.14159265359
    x = x/RAD_TO_DEG
    y = y/RAD_TO_DEG
    z = z/RAD_TO_DEG
    A,B = math.cos(x),math.sin(x)
    C,D = math.cos(y),math.sin(y)
    E,F = math.cos(z),math.sin(z)

    x = math.asin(-B*C)
    y = math.atan2(D, A*C)
    z = math.atan2(-B*D*E + A*F, B*D*F + A*E)

    x = -x
    RAD_TO_DEG = 180.0/3.14159265359
    return x*RAD_TO_DEG, y*RAD_TO_DEG, z*RAD_TO_DEG


#============================================================
# GUI
#============================================================
#==============================
def event(evt, val):
#==============================
#    print '[event]-- evt:', evt, ', val:', val, ' --'
    if evt == Draw.ESCKEY:
        Draw.Exit()
    return

#==============================
def button_event(evt):
#==============================
    global outputFile, scale, startFrame, endFrame, targetArmature, armatures, action
    global guiArmatureMenu, guiScaleNumber, guiStartFrameNumber, guiEndFrameNumber, guiFileNameString
    if (evt == EVENT_EXIT):
        print 'EVENT_EXIT'
        Draw.Exit()
    elif (evt == EVENT_EXPORT):
        print 'EVENT_EXPORT'
        exportBVH(outputFile)
    elif (evt == EVENT_SELECT_ARMATURE):
        print 'EVENT_SELECT_ARMATURE'
        print(guiArmatureMenu.val)
        targetArmature = armatures[guiArmatureMenu.val-1]
        action = searchAction(targetArmature)
        o = searchObject(targetArmature)
        outputFile = o.getName() + '.bvh'
        Draw.Redraw(1)
        print targetArmature
    elif (evt == EVENT_CHANGE_SCALE):
        print 'EVENT_CHANGE_SCALE'
        print(guiScaleNumber.val)
        scale = guiScaleNumber.val
    elif (evt == EVENT_CHANGE_START_FRAME):
        print 'EVENT_CHANGE_START_FRAME'
        print(guiStartFrameNumber.val)
        startFrame = guiStartFrameNumber.val
    elif (evt == EVENT_CHANGE_END_FRAME):
        print 'EVENT_CHANGE_END_FRAME'
        print(guiEndFrameNumber.val)
        endFrame = guiEndFrameNumber.val
    elif (evt == EVENT_CHANGE_FILENAME):
        print 'EVENT_CHANGE_FILENAME'
        print(guiFileNameString.val)
        outputFile = guiFileNameString.val
        #Draw.PupMenu("Not implemented yet%t|OK")

#==============================
def gui():
#==============================
    global outputFile, scale, startFrame, endFrame, armatures, message
    global guiArmatureMenu, guiScaleNumber, guiStartFrameNumber
    global guiEndFrameNumber, guiFileNameString, targetArmature
    BGL.glClearColor(0.8,0.8,1,1)
    BGL.glClear(BGL.GL_COLOR_BUFFER_BIT)
    BGL.glColor3f(0,0.5,0)
    BGL.glRasterPos2i(10,210)
    Draw.Text("BriskSkeletons")
    BGL.glRasterPos2i(10,195)
    Draw.Text("USAGE: --> http://acerola3d.sourceforge.jp/bs/")
    name = "Root Bone List%t"
    for a in armatures:
        o = searchObject(a)
        name = name + "|" + o.getName()
        #name = name + "|" + a.name
        #name = name + "|" + a.getName()
    guiArmatureMenu = Draw.Menu(name,EVENT_SELECT_ARMATURE,10,160,200,20,armatures.index(targetArmature)+1,"Select an Armature.")
    guiScaleNumber = Draw.Number("Scale: ",EVENT_CHANGE_SCALE,10,130,200,20,scale,0.0,100000.0,"scale.")
    guiStartFrameNumber = Draw.Number("Start: ",EVENT_CHANGE_START_FRAME,10,100,200,20,startFrame,0,100000,"start.")
    guiEndFrameNumber = Draw.Number("End: ",EVENT_CHANGE_END_FRAME,10,70,200,20,endFrame,0,100000,"end.")
    guiFileNameString = Draw.String("File Name: ",EVENT_CHANGE_FILENAME,10,40,200,20,outputFile,100,"Input output file name.")
    Draw.Button("Export",EVENT_EXPORT,10,10,50,20,"Export.")
    Draw.Button("Exit",EVENT_EXIT,70,10,50,20,"Quit.")
    BGL.glColor3f(0,0.5,0)
    BGL.glRasterPos2i(130,15)
    Draw.Text(message)

#==============================
def errorGUI():
#==============================
    BGL.glClearColor(0.8,0.8,1,1)
    BGL.glClear(BGL.GL_COLOR_BUFFER_BIT)
    BGL.glColor3f(0,0.5,0)
    BGL.glRasterPos2i(10,150)
    Draw.Text("There is no abailable bone!")
    BGL.glRasterPos2i(10,130)
    Draw.Text("USAGE: --> http://acerola3d.sourceforge.jp/bs/")
    Draw.Button("Exit",EVENT_EXIT,10,100,200,20,"Quit.")


#============================================================
# MAIN
#============================================================
init()
#if changeParam():
#    Blender.Window.FileSelector(exportBVH,"Export BVH",outputFile)
if (targetArmature == None):
    Draw.Register(errorGUI, event, button_event)
else:
    Draw.Register(gui, event, button_event)
