import sys
import math
import traceback
import maya.cmds as cmds
import maya.OpenMaya as OpenMaya
import maya.OpenMayaMPx as OpenMayaMPx

class Utils(): pass

class Vector(OpenMaya.MVector):
	def __init__(self, *args ):

		if len(args)==0:
			super(Vector,self).__init__( )

		elif len(args)==1:
			if isinstance(args[0],list) and len(args[0])==3:
				super(Vector,self).__init__(args[0][0],args[0][1],args[0][2])
			elif isinstance(args[0],OpenMaya.MVector):
				super(Vector,self).__init__(args[0])
			elif isinstance(args[0],Vector):
				super(Vector,self).__init__(args[0].x,args[0].y,args[0].z)

		elif len(args)==3:
			super(Vector,self).__init__( args[0],args[1],args[2] )
				
	def __str__(self):
		return '['+str(self.x)+', '+str(self.y)+', '+str(self.z)+']'

	def __sub__(self,value):
		if isinstance(value,OpenMaya.MVector):
			return Vector( self.x-value.x,self.y-value.y,self.z-value.z )

	def __mul__(self, value):
		if isinstance(value, OpenMaya.MMatrix ):
			#return Vector( OpenMaya.MVector(self.x,self.y,self.z)*value )
			return Vector( super(Vector,self).__mul__(value) )

	
class Matrix(OpenMaya.MMatrix):

	pi = 3.14159265359

	def __init__(self, *args):
		# constructor
		if not args:
			super(Matrix,self).__init__( )
		# initialize
		else:
			if len(args)==1 and isinstance(args[0],OpenMaya.MMatrix):
				super(Matrix,self).__init__(args[0])

			if len(args)==16:
				mMatrix = OpenMaya.MMatrix()
				OpenMaya.MScriptUtil().createMatrixFromList(args,mMatrix)
				super(Matrix,self).__init__(mMatrix)

			if len(args)==4:

				v0 = Vector(args[0])
				v1 = Vector(args[1])
				v2 = Vector(args[2])
				v3 = Vector(args[3])
				
				m = [v0.x,v0.y,v0.z,0,v1.x,v1.y,v1.z,0,v2.x,v2.y,v2.z,0,v3.x,v3.y,v3.z,0]
				mMatrix = OpenMaya.MMatrix()
				OpenMaya.MScriptUtil().createMatrixFromList(m,mMatrix)
				super(Matrix,self).__init__(mMatrix)
		
	def __str__(self):
		return '['+','.join( [str(f) for f in self.get()] )+']'
		#return '['+str(self(0,0))+', '+str(self(0,1))+', '+str(self(0,2))+', '+str(self(0,3))+', '+str(self(1,0))+', '+str(self(1,1))+', '+str(self(1,2))+', '+str(self(1,3))+', '+str(self(2,0))+', '+str(self(2,1))+', '+str(self(2,2))+', '+str(self(2,3))+', '+str(self(3,0))+', '+str(self(3,1))+', '+str(self(3,2))+', '+str(self(3,3))+']'

	def __mul__(self, value):
		if isinstance(value, OpenMaya.MMatrix ):
			#m = [self(0,0),self(0,1),self(0,2),self(0,3),self(1,0),self(1,1),self(1,2),self(1,3),self(2,0),self(2,1),self(2,2),self(2,3),self(3,0),self(3,1),self(3,2),self(3,3)]
			#mMatrix = OpenMaya.MMatrix()
			#OpenMaya.MScriptUtil().createMatrixFromList(m,mMatrix)	
			#return Matrix( mMatrix*value  )

			return Matrix( super(Matrix,self).__mul__(value) )

	def __getitem__(self,index):

		if index>=0 and index<4:
			return self(0,index)
		if index>=4 and index<8:
			return self(1,index+4)
		if index>=8 and index<12:
			return self(2,index+8)
		if index>=12 and index<16:
			return self(3,index+12)

	def get(self):
		"""
		Returns the Matrix as a list
		"""
		return [self(0,0),self(0,1),self(0,2),self(0,3),self(1,0),self(1,1),self(1,2),self(1,3),self(2,0),self(2,1),self(2,2),self(2,3),self(3,0),self(3,1),self(3,2),self(3,3)]

	def inverse(self):
		
		#m = [self(0,0),self(0,1),self(0,2),self(0,3),self(1,0),self(1,1),self(1,2),self(1,3),self(2,0),self(2,1),self(2,2),self(2,3),self(3,0),self(3,1),self(3,2),self(3,3)]
		#mMatrix = OpenMaya.MMatrix()
		#OpenMaya.MScriptUtil().createMatrixFromList(m,mMatrix)
		#return Matrix( mMatrix.inverse() )

		return Matrix( super(Matrix,self).inverse() )
		
	def getTranslate(self):
		
		return Vector( self(3,0),self(3,1),self(3,2) )

	def getRotate(self):
		eRotate = OpenMaya.MTransformationMatrix(self).rotation().asEulerRotation()
		#return Vector(OpenMaya.MVector(eRotate.x*180/math.pi,eRotate.y*180/math.pi,eRotate.z*180/math.pi))
		return Vector( eRotate.x*180.0/self.pi,eRotate.y*180.0/self.pi,eRotate.z*180.0/self.pi )

	def getScale(self):
		mMatrix = OpenMaya.MTransformationMatrix(self).asScaleMatrix()
		return Vector( mMatrix(0.0),mMatrix(1,1),mMatrix(2,2) )

# command
class pyExpressionCmd( OpenMayaMPx.MPxCommand ):
	def __init__(self):
		OpenMayaMPx.MPxCommand.__init__( self )
 
	def doIt( self, argList ):

		# get node
		sNode = None
		if argList.length()!=1:
			OpenMaya.MGlobal.displayError('// You have to specify the node name at the end.')
			return

		sNode = argList.asString( argList.length() -1 )
		if not cmds.objExists(sNode):
			OpenMaya.MGlobal.displayError('// Node not found: '+sNode )
			return

		# get node
		oNode = OpenMaya.MObject()
		sList = OpenMaya.MSelectionList()
		sList.add(sNode)
		sList.getDependNode(0,oNode)
		fnDep = OpenMaya.MFnDependencyNode(oNode)

		# GET EXPRESSION CODE
		pythonCode = ''
		try:
			pythonCode = fnDep.findPlug('expression').asString()
		except:
			return
		if not pythonCode or len(pythonCode)<2:
			return

		# GET INPUTS
		myDict = dict(locals())
		myDict['Vector'] = Vector
		myDict['Matrix'] = Matrix

		#
		# GET INPUTS
		#

		# OBJ
		myDict['OS'] = fnDep.name()

		# TIME INPUT
		plug = fnDep.findPlug('time')
		t = plug.asInt()
		myDict['T']=t

		# NODE INPUT
		N = []
		plug = fnDep.findPlug('N')
		plugArr = OpenMaya.MPlugArray()
		for i in range( plug.numElements() ):
			plug.elementByLogicalIndex(i).connectedTo(plugArr,True,False)
			if plugArr.length()>0:
				N.append( OpenMaya.MFnDependencyNode( plugArr[0].node() ).name() )
		myDict['N'] = N

		# FLOAT INPUT
		F = []
		plug = fnDep.findPlug('F')
		for i in range(plug.numElements()):
			if plug.elementByLogicalIndex(i).isConnected() == False: continue
			f =  plug.elementByLogicalIndex(i).asFloat()
			#sAlias = plug.elementByLogicalIndex(i).partialName(False,False,False,True,False,False)
			#if sAlias!='F':
			#	myDict[sAlias] = fValue
			F.append(f)
		myDict['F'] = F

		# VECTOR INPUT
		V = []
		plug = fnDep.findPlug('V')
		for i in range(plug.numElements()):
			if plug.elementByLogicalIndex(i).isConnected() == False: continue
			vx = plug.elementByLogicalIndex(i).child(0).asFloat( )
			vy = plug.elementByLogicalIndex(i).child(1).asFloat( )
			vz = plug.elementByLogicalIndex(i).child(2).asFloat( )
			V.append( Vector( vx,vy,vz ) )
		myDict['V'] = V

		# MATRIX INPUT
		M = []
		plug = fnDep.findPlug('M')
		for i in range(plug.numElements()):
			if plug.elementByLogicalIndex(i).isConnected() == False: continue
			o = plug.elementByLogicalIndex(i).asMObject()
			m = OpenMaya.MFnMatrixData(o).matrix()
			M.append( Matrix(m) )
		myDict['M'] = M

		#
		# GET OUTPUTS
		#

		# FLOAT OUTPUT
		OF = []
		plug = fnDep.findPlug('OF')
		for i in range(plug.numElements()):
			of =  plug.elementByLogicalIndex(i).asFloat()
			OF.append( of )
		myDict['OF'] = OF

		# VECTOR OUTPUT
		OV = []
		plug = fnDep.findPlug('OV')
		for i in range(plug.numElements()):
			ofx = plug.elementByLogicalIndex(i).child(0).asFloat( )
			ofy = plug.elementByLogicalIndex(i).child(1).asFloat( )
			ofz = plug.elementByLogicalIndex(i).child(2).asFloat( )
			OV.append( Vector( ofx,ofy,ofz ) )
		myDict['OV'] = OV

		# MATRIX OUTPUT
		OM = []
		plug = fnDep.findPlug('OM')
		for i in range(plug.numElements()):
			oo = plug.elementByLogicalIndex(i).asMObject()
			om = OpenMaya.MFnMatrixData(oo).matrix()
			OM.append( Matrix(om) )
		myDict['OM'] = OM

		# COMPUTE
		try:
			exec pythonCode in myDict
		except Exception,e:
			OpenMaya.MGlobal.displayError(traceback.format_exc())
			return
			#print "ERROR!! The code does't work: ",str(e)

		#
		# SET VALUES
		#

		# OUT FLOAT
		plug = fnDep.findPlug('OF')
		for i in range( plug.numElements() ):
			plug.elementByLogicalIndex(i).setFloat( myDict['OF'][i] )

		# OUT VECTOR
		plug = fnDep.findPlug('OV')
		for i in range( plug.numElements() ):
			ov = myDict['OV'][i]
			sType = type(ov)
			fX,fY,fZ = 0,0,0

			if sType==OpenMaya.MVector or sType==Vector:
				fX = ov.x
				fY = ov.y
				fZ = ov.z

			elif (sType==list or sType==tuple) and len(ov)==3:
				fX =  ov[0] 
				fY =  ov[1]
				fZ =  ov[2] 

			plug.elementByLogicalIndex(i).child(0).setFloat( fX )
			plug.elementByLogicalIndex(i).child(1).setFloat( fY )
			plug.elementByLogicalIndex(i).child(2).setFloat( fZ )

		# OUT MATRIX
		plug = fnDep.findPlug('OM')
		for i in range( plug.numElements() ):

			sType = type(myDict['OM'][i])
			mMatrix = OpenMaya.MMatrix()

			if sType==OpenMaya.MMatrix:
				mMatrix = myDict['OM'][i]

			elif sType==Matrix:
				OpenMaya.MScriptUtil().createMatrixFromList(myDict['OM'][i].get(),mMatrix)
				#mMatrix = super(Matrix, myDict['OM'][i] )

			elif (sType==list or sType==tuple) and len(myDict['OM'][i])==16:
				OpenMaya.MScriptUtil().createMatrixFromList(myDict['OM'][i],mMatrix)

			fnMatrixData = OpenMaya.MFnMatrixData()
			oMatrixData = fnMatrixData.create(mMatrix)
			plug.elementByLogicalIndex(i).setMObject( oMatrixData )

#
# NODE DEFINITION
#

class pyExpression(OpenMayaMPx.MPxNode):

	expression = OpenMaya.MObject()
	time = OpenMaya.MObject()
	N = OpenMaya.MObject()
	F = OpenMaya.MObject()
	V = OpenMaya.MObject()
	M = OpenMaya.MObject()
	OF = OpenMaya.MObject()
	OV = OpenMaya.MObject()
	OM = OpenMaya.MObject()

	def __init__(self):
		OpenMayaMPx.MPxNode.__init__(self)

	def compute(self,plug,dataBlock):

		if  plug != pyExpression.OF and plug != pyExpression.OV and plug != pyExpression.OM:
			return

		oNode = self.thisMObject() 
		fnDep = OpenMaya.MFnDependencyNode(oNode)

		myDict = dict(locals())
		myDict['Vector'] = Vector
		myDict['Matrix'] = Matrix

		# GET EXPRESSION CODE
		pythonCode = ''
		try:
			pythonCode = fnDep.findPlug('expression').asString()
			#pythonCode = dataBlock.inputValue(pyExpression.expression).asString()
		except:
			return
		if not pythonCode or len(pythonCode)<2:
			return

		#
		# GET INPUTS
		#

		# OBJ
		myDict['OS'] = fnDep.name()

		# TIME INPUT
		t = dataBlock.inputValue( pyExpression.time ).asInt()
		myDict['T'] = t 

		# NODE INPUT
		N = []
		plugN = fnDep.findPlug('N')
		plugArr = OpenMaya.MPlugArray()
		for i in range( plugN.numElements() ):
			plugN.elementByLogicalIndex(i).connectedTo(plugArr,True,False)
			if plugArr.length()>0:
				N.append( OpenMaya.MFnDependencyNode( plugArr[0].node() ).name() )
		myDict['N'] = N

		# FLOAT INPUT
		F = []
		hArrData = dataBlock.inputArrayValue( pyExpression.F )
		for i in range(hArrData.elementCount()):
			hArrData.jumpToArrayElement(i)
			f = hArrData.inputValue().asFloat()
			F.append( f )
		myDict['F'] = F

		# VECTOR INPUT
		V = []
		hArrData = dataBlock.inputArrayValue( pyExpression.V )
		for i in range(hArrData.elementCount()):
			hArrData.jumpToArrayElement(i)
			v = Vector( hArrData.inputValue().asVector() )
			V.append( v )
		myDict['V'] = V

		# MATRIX INPUT
		M = []
		hArrData = dataBlock.inputArrayValue( pyExpression.M )
		for i in range(hArrData.elementCount()):
			hArrData.jumpToArrayElement(i)
			m = Matrix( hArrData.inputValue().asMatrix() )
			M.append( m )
		myDict['M'] = M

		#
		# GET OUTPUTS
		#

		# FLOAT OUTPUT
		OF = []
		hArrData = dataBlock.outputArrayValue( pyExpression.OF )
		for i in range(hArrData.elementCount()):
			hArrData.jumpToArrayElement(i)
			of = hArrData.outputValue().asFloat()
			OF.append( of )
		myDict['OF'] = OF

		# VECTOR OUTPUT
		OV = []
		hArrData = dataBlock.outputArrayValue( pyExpression.OV )
		for i in range(hArrData.elementCount()):
			hArrData.jumpToArrayElement(i)
			ov = Vector( hArrData.outputValue().asVector() )
			OV.append( ov )
		myDict['OV'] = OV

		# MATRIX OUTPUT
		OM = []
		hArrData = dataBlock.outputArrayValue( pyExpression.OM )
		for i in range(hArrData.elementCount()):
			hArrData.jumpToArrayElement(i)
			om = Matrix( hArrData.outputValue().asMatrix() )
			OM.append( om )
		myDict['OM'] = OM

		# COMPUTE
		try:
			exec pythonCode in myDict
		except Exception,e:
			OpenMaya.MGlobal.displayError(traceback.format_exc())
			return
			#print "ERROR!! The code does't work: ",str(e)

		#
		# SET OUTPUTS
		#

		# FLOAT OUTPUT
		if  plug == pyExpression.OF:

			hArrData = dataBlock.outputArrayValue( pyExpression.OF )
			bData = OpenMaya.MArrayDataBuilder(pyExpression.OF, hArrData.elementCount() )

			for i in range(hArrData.elementCount()):
				of = myDict['OF'][i]
				hData = bData.addElement(i)
				hData.setFloat(of)
			hArrData.set(bData)
			hArrData.setAllClean()

		# VECTOR OUTPUT
		if  plug == pyExpression.OV:

			hArrData = dataBlock.outputArrayValue( pyExpression.OV )
			bData = OpenMaya.MArrayDataBuilder(pyExpression.OV, hArrData.elementCount() )

			for i in range(hArrData.elementCount()):

				ov = myDict['OV'][i]
				if isinstance(ov, OpenMaya.MVector ):
					pass
				elif isinstance(ov, Vector ):
					ov = OpenMaya.MVector(Vector.x,Vector.y,Vector.z)
				elif isinstance(ov, list ) and len(ov)==3:
					ov = OpenMaya.MVector(ov[0],ov[1],ov[2])

				hData = bData.addElement(i)
				hData.setMVector(ov)
			hArrData.set(bData)
			hArrData.setAllClean()

		# MATRIX OUTPUT
		if  plug == pyExpression.OM:

			hArrData = dataBlock.outputArrayValue( pyExpression.OM )
			bData = OpenMaya.MArrayDataBuilder(pyExpression.OM, hArrData.elementCount() )
					
			for i in range(hArrData.elementCount()):
				om = myDict['OM'][i]
				if isinstance(om, OpenMaya.MMatrix ):
					hData = bData.addElement(i)
					hData.setMMatrix(om)
				elif isinstance(om, Matrix ):
					#om = 
					hData = bData.addElement(i)
					hData.setMMatrix(super(Matrix,om))
			hArrData.set(bData)
			hArrData.setAllClean()

# initializer
def nodeInitializer():

	cAttr = OpenMaya.MFnCompoundAttribute()
	nAttr = OpenMaya.MFnNumericAttribute()
	mAttr = OpenMaya.MFnMatrixAttribute()
	tAttr = OpenMaya.MFnTypedAttribute()
	meAttr = OpenMaya.MFnMessageAttribute()

	# EXPRESSION
	oString = OpenMaya.MFnStringData().create('# ojete')
	pyExpression.expression = tAttr.create( "expression", "expression", OpenMaya.MFnData.kString, oString )
	tAttr.setStorable(True)
	tAttr.setKeyable(False)
	pyExpression.addAttribute( pyExpression.expression )

	# TIME
	pyExpression.time = nAttr.create( "time", "time", OpenMaya.MFnNumericData.kInt, 0 )
	nAttr.setKeyable(True)
	pyExpression.addAttribute( pyExpression.time )

	# TEXTURE
	pyExpression.color = nAttr.createColor("color", "color" )
	pyExpression.alpha = nAttr.create( "alpha", "alpha", OpenMaya.MFnNumericData.kFloat, 1.0 )
	pyExpression.texture = cAttr.create('texture','texture')
	cAttr.addChild(pyExpression.color)
	cAttr.addChild(pyExpression.alpha)
	pyExpression.addAttribute( pyExpression.texture )

	# INPUT NODE
	pyExpression.N = meAttr.create( "N", "N")
	meAttr.setArray(True)
	meAttr.setKeyable(True)
	pyExpression.addAttribute( pyExpression.N )

	# INPUT FLOAT
	pyExpression.F = nAttr.create( "F", "F", OpenMaya.MFnNumericData.kFloat, 0.0 )
	nAttr.setArray(True)
	nAttr.setKeyable(True)
	pyExpression.addAttribute( pyExpression.F )

	# INPUT VECTOR
	pyExpression.VX = nAttr.create( "VX", "VX", OpenMaya.MFnNumericData.kDouble, 0.0 )
	pyExpression.VY = nAttr.create( "VY", "VY", OpenMaya.MFnNumericData.kDouble, 0.0 )
	pyExpression.VZ = nAttr.create( "VZ", "VZ", OpenMaya.MFnNumericData.kDouble, 0.0 )
	pyExpression.V = nAttr.create( "V", "V", pyExpression.VX,pyExpression.VY,pyExpression.VZ)
	nAttr.setArray(True)
	nAttr.setKeyable(True)
	pyExpression.addAttribute( pyExpression.V )

	# INPUT MATRIX 
	#pyExpression.m = mAttr.create( "m", "m", OpenMaya.MFnMatrixAttribute.kDouble )
	pyExpression.M = mAttr.create( "M", "M")
	mAttr.setArray(True)
	mAttr.setKeyable(True)
	pyExpression.addAttribute( pyExpression.M )

	# OUTPUT FLOAT
	pyExpression.OF = nAttr.create( "OF", "OF", OpenMaya.MFnNumericData.kFloat, 0.0 )
	nAttr.setArray(True)
	nAttr.setReadable(True)
	nAttr.setWritable(False)
	nAttr.setUsesArrayDataBuilder(True)
	pyExpression.addAttribute( pyExpression.OF )

	# OUTPUT VECTOR
	pyExpression.OVX = nAttr.create( "OVX", "OVX", OpenMaya.MFnNumericData.kDouble, 0.0 )
	pyExpression.OVY = nAttr.create( "OVY", "OVY", OpenMaya.MFnNumericData.kDouble, 0.0 )
	pyExpression.OVZ = nAttr.create( "OVZ", "OVZ", OpenMaya.MFnNumericData.kDouble, 0.0 )
	pyExpression.OV = nAttr.create( "OV", "OV", pyExpression.OVX,pyExpression.OVY,pyExpression.OVZ)
	nAttr.setArray(True)
	nAttr.setReadable(True)
	nAttr.setWritable(False)
	nAttr.setUsesArrayDataBuilder(True)
	pyExpression.addAttribute( pyExpression.OV )

	# OUTPUT MATRIX
	pyExpression.OM = mAttr.create( "OM", "OM")
	mAttr.setArray(True)
	mAttr.setReadable(True)
	mAttr.setWritable(False)
	mAttr.setUsesArrayDataBuilder(True)
	pyExpression.addAttribute( pyExpression.OM )

	# RELATIONSHIPS
	pyExpression.attributeAffects( pyExpression.time, pyExpression.OF )
	pyExpression.attributeAffects( pyExpression.time, pyExpression.OV )
	pyExpression.attributeAffects( pyExpression.time, pyExpression.OM )
	pyExpression.attributeAffects( pyExpression.N, pyExpression.OF )
	pyExpression.attributeAffects( pyExpression.N, pyExpression.OV )
	pyExpression.attributeAffects( pyExpression.N, pyExpression.OM )
	pyExpression.attributeAffects( pyExpression.F, pyExpression.OF )
	pyExpression.attributeAffects( pyExpression.F, pyExpression.OV )
	pyExpression.attributeAffects( pyExpression.F, pyExpression.OM )
	pyExpression.attributeAffects( pyExpression.V, pyExpression.OF )
	pyExpression.attributeAffects( pyExpression.V, pyExpression.OV )
	pyExpression.attributeAffects( pyExpression.V, pyExpression.OM )
	pyExpression.attributeAffects( pyExpression.M, pyExpression.OF )
	pyExpression.attributeAffects( pyExpression.M, pyExpression.OV )
	pyExpression.attributeAffects( pyExpression.M, pyExpression.OM )
