#!/usr/bin/python
# -*- coding: utf-8 -*-
#最新のAPIを使って大幅に書換えることに
#したのだけど，役にたつこともあるかもしれない
#ので残しておく．
#==============================
# 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
#
#==============================
# BriskSkeletons v0.6
# http://vesma.sf.net/bs/
#==============================
# BriskSkeletons is Copyright (C) 2004,2005,2006 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@do-johodai.ac.jp>
#==============================

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

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

outputFile = "output.bvh"
scale = 1.0
startFrame = None
endFrame = None
armatures = None
targetArmature = None
guiArmatureMenu = None
guiScaleNumber = None
guiStartFrameNumber = None
guiEndFrameNumber = None
guiFileNameString = None
message = "status: READY"

#Boneから得られるhead,tailの座標は、親の骨のtailを原点とするローカル座標系での
#座標なので，それをグローバル座標系に直して保存しておく。
#それと、後で使えそうな回転のデータを保存しておくためのクラス
class BoneBase:
	head = None # グローバル座標系での骨の頭
	tail = None # グローバル座標系での骨の尾
	gQuat = None # 骨のローカル座標系の点をグローバル座標系の点に変換するための回転とする

#============================================================
# INIT
#============================================================
#==============================
def init():
#==============================
	global startFrame,endFrame,armatures,targetArmature
	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
			break

#============================================================
# EXPORT BVH
#============================================================
#==============================
def getRootBones():
#==============================
	global targetArmature, message
	message = "status: EXPORTING..."
	Draw.Redraw(1)
	rootBones = []
        bs = targetArmature.bones
        for b in bs.values():
		if (b.parent == None):
			rootBones.append(b)
	message = "status: EXPORTING... DONE!"
	Draw.Redraw(1)
	return rootBones

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

#==============================
def calBoneBases(rootBones):
#==============================
	boneBases = {}
	bb = BoneBase()
	bb.head = Mathutils.Vector([0.0,0.0,0.0])
	bb.tail = Mathutils.Vector([0.0,0.0,0.0])
	bb.gQuat = Mathutils.Quaternion([1.0,0.0,0.0,0.0])
	boneBases['BSRootBone'] = bb
	for bone in rootBones:
		calBoneBases2(bone,bb,boneBases)
	return boneBases

#==============================
def calBoneBases2(bone,parentBB,boneBases):
#==============================
	h = bone.head['BONESPACE']
	t = bone.tail['BONESPACE']
	hh = rotate(h,parentBB.gQuat)
	tt = rotate(t,parentBB.gQuat)
	head = [parentBB.tail[i]+hh[i] for i in range(3)]
	tail = [parentBB.tail[i]+tt[i] for i in range(3)]
	vecY = [t[i]-h[i] for i in range(3)]
	pQuat = vec_vec2quat([0,1,0],vecY)
	roll = bone.roll['BONESPACE']
	if (roll > 0.1):
		print "roll("+bone.name+")="+str(roll)
        roll = 3.14159265359*roll/180.0
	rQuat = Mathutils.Quaternion([math.cos(roll/2.0),0.0,-1.0*math.sin(roll/2.0),0.0])
	gQuat = Mathutils.CrossQuats(pQuat,rQuat)
	gQuat = Mathutils.CrossQuats(parentBB.gQuat,gQuat)
	bb = BoneBase()
	bb.head = head
	bb.tail = tail
	bb.gQuat = gQuat
	boneBases[bone.name] = bb
	if bone.children:
		for b in bone.children:
			calBoneBases2(b,bb,boneBases)

#==============================
#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]])

#==============================
#v1軸上の点をv2軸上の点に移すための最もすなおな回転を求める
def vec_vec2quat(v1,v2):
#==============================
	vec1 = Mathutils.Vector(v1)
	vec2 = Mathutils.Vector(v2)
	vec1.normalize()
	vec2.normalize()
	outer = Mathutils.CrossVecs(vec1,vec2)
	oLen = outer.length
	inner = Mathutils.DotVecs(vec1,vec2)
	if (oLen == 0.0):
		if (inner > 0.0):
			return Mathutils.Quaternion([1.0,0.0,0.0,0.0])
		else:
			print "gaha:warning 180"
#			return Mathutils.Quaternion([0.0,1.0,0.0,0.0]) # <-- *****んー。重要なんだけどビミョウ*****
#			return Mathutils.Quaternion([0.0,0.0,1.0,0.0]) # <-- *****んー。重要なんだけどビミョウ*****
			return Mathutils.Quaternion([0.0,0.0,0.0,1.0]) # <-- *****んー。重要なんだけどビミョウ*****
	outer.normalize()
	theta = math.acos(inner)
	cos = math.cos(theta/2.0)
	sin = math.sin(theta/2.0)
	qTmp = [cos,sin*outer[0],sin*outer[1],sin*outer[2]]
	return Mathutils.Quaternion(qTmp)

#==============================
def needRotChannels(bone):
#==============================
	try:
		ipoName = "Action." + bone.name
		ipo = Blender.Ipo.Get(ipoName)
		ipoCurve = ipo.getCurve('QuatW')
		if (ipoCurve != None):
			return True
		else:
			return False
	except NameError:
		return False

#==============================
def needLocChannels(bone):
#==============================
	try:
		ipoName = "Action." + bone.name
		ipo = Blender.Ipo.Get(ipoName)
		ipoCurve = ipo.getCurve('LocX')
		if (ipoCurve != None):
			return True
		else:
			return False
	except NameError:
		return False

#==============================
def needScaleChannels(bone):
#==============================
	try:
		ipoName = "Action." + bone.name
		ipo = Blender.Ipo.Get(ipoName)
		#ipoCurve = ipo.getCurve('SizeX')
		ipoCurve = ipo.getCurve('ScaleX')
		if (ipoCurve != None):
			return True
		else:
			return False
	except NameError:
		return False

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

#==============================
def printHierarchy(rootBones,boneBases,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,boneBases,1,file)
	file.write('}\n')

#==============================
def printBlock(bone,boneBases,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,boneBases) + '\n')
	file.write(indent + '  ' + channelsString(bone) + '\n')
	bs = bone.children
	if (bs == None):
		printEndSite(bone,boneBases,i+1,file)
	else:
		for b in bs:
			printBlock(b,boneBases,i+1,file)
	file.write(indent + '}\n')

#==============================
def offsetString(bone,boneBases):
#==============================
	global scale
	parentBone = bone.parent
	if (parentBone == None):
		parentBoneHead = [0,0,0]
	else:
		parentBoneHead = boneBases[parentBone.name].head
	boneHead = boneBases[bone.name].head
	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,boneBases,i,file):
#==============================
	global scale
	boneHead = boneBases[bone.name].head
	boneTail = boneBases[bone.name].tail
	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,boneBases,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):
		Blender.Set('curframe',f)
		file.write('0.0 0.0 0.0 ') # <-- for BSRootBone
		for bone in rootBones:
			printMotion2(bone,boneBases,file)
		file.write('\n')
	Blender.Set('curframe',startFrame)

#==============================
def printMotion2(bone,boneBases,file):
#==============================
	global scale
	gQuat = boneBases[bone.name].gQuat
	if (needLocChannels(bone)):
		loc = getLoc(bone)
		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)
		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)
		file.write('%f %f %f ' % (lScale[0],lScale[1],lScale[2]))
	bs = bone.children
	if bs:
		for b in bs:
			printMotion2(b,boneBases,file)

#==============================
#この実装，将来変更の予定あり
def getLoc(bone):
#==============================
	try:
		ipo = Blender.Ipo.Get('Action.'+bone.name)
	except NameError:
		return Mathutils.Vector(0,0,0);
	cx = ipo.getCurve('LocX')
	cy = ipo.getCurve('LocY')
	cz = ipo.getCurve('LocZ')
	f = Blender.Get('curframe')
	return Mathutils.Vector(cx.evaluate(f),cy.evaluate(f),cz.evaluate(f))

#==============================
#この実装，将来変更の予定あり
def getQuat(bone):
#==============================
	try:
		ipo = Blender.Ipo.Get('Action.'+bone.name)
	except NameError:
		print 'Action.'+bone.name
		return Mathutils.Quaternion(1,0,0,0);
	cw = ipo.getCurve('QuatW')
	cx = ipo.getCurve('QuatX')
	cy = ipo.getCurve('QuatY')
	cz = ipo.getCurve('QuatZ')
	f = Blender.Get('curframe')
	w = cw.evaluate(f)
	x = cx.evaluate(f)
	y = cy.evaluate(f)
	z = cz.evaluate(f)
	return Mathutils.Quaternion(w,x,y,z)

#==============================
#この実装，将来変更の予定あり
def getSize(bone):
#==============================
	try:
		ipo = Blender.Ipo.Get('Action.'+bone.name)
	except NameError:
		return Mathutils.Vector(0,0,0);
	sx = ipo.getCurve('ScaleX')
	sy = ipo.getCurve('ScaleY')
	sz = ipo.getCurve('ScaleZ')
	f = Blender.Get('curframe')
	return Mathutils.Vector(sx.evaluate(f),sy.evaluate(f),sz.evaluate(f))

#==============================
#回転の回転。回転の合成とは違うことに注意。
#座標系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
	global guiArmatureMenu, guiScaleNumber, guiStartFrameNumber, guiEndFrameNumber, guiFileNameString
	if (evt == EVENT_EXIT):
		print 'EVENT_EXIT'
		Draw.Exit()
	elif (evt == EVENT_EXPORT):
		print 'EVENT_EXPORT'
		exportBVH()
	elif (evt == EVENT_SELECT_ARMATURE):
		print 'EVENT_SELECT_ARMATURE'
		print(guiArmatureMenu.val)
		targetArmature = armatures[guiArmatureMenu.val-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, guiEndFrameNumber, guiFileNameString
	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 ver0.4")
	BGL.glRasterPos2i(10,195)
	Draw.Text("USAGE: --> http://vesma.sf.net/bs/")
	name = "Root Bone List%t"
	for a in armatures:
		name = name + "|" + a.name
		#name = name + "|" + a.getName()
	guiArmatureMenu = Draw.Menu(name,EVENT_SELECT_ARMATURE,10,160,200,20,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://vesma.sf.net/bs/")
	Draw.Button("Exit",EVENT_EXIT,10,100,200,20,"Quit.")

#============================================================
# MAIN
#============================================================
init()
if (targetArmature == None):
	Draw.Register(errorGUI, event, button_event)
else:
	Draw.Register(gui, event, button_event)
