import sys, math, traceback
import maya.OpenMaya as OpenMaya
import maya.OpenMayaMPx as OpenMayaMPx
import maya.OpenMayaFX as omfx

from pyExpressionUtils import Vector
from pyExpressionUtils import Matrix

def creator():
	return OpenMayaMPx.asMPxPtr( pyExpressionNode() )

# ----------------------------------------------------------------------------

class pyExpressionNode( 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):
		'''Initialize class
		'''
		OpenMayaMPx.MPxNode.__init__(self)
		
	def compute(self, plug, dataBlock):
		'''Compute node
		'''
		if  plug != pyExpressionNode.OF and plug != pyExpressionNode.OV and plug != pyExpressionNode.OM:
			return

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

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

		# GET CURRENT NODE
		oNode = self.thisMObject()
		fnDep = OpenMaya.MFnDependencyNode(oNode)

		#
		# GET INPUTS
		#

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

		# TIME INPUT
		t = dataBlock.inputValue(pyExpressionNode.time).asDouble()
		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(pyExpressionNode.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(pyExpressionNode.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(pyExpressionNode.M)
		for i in range(hArrData.elementCount()):
			hArrData.jumpToArrayElement(i)
			m = Matrix(hArrData.inputValue().asMatrix())
			M.append(m)
		myDict['M'] = M

		# MESH INPUT
		myDict['PP'] = []
		myDict['PR'] = []
		myDict['PC'] = []
		myDict['PA'] = []
		myDict['PU'] = []
		myDict['PV'] = []
		plugInMesh = fnDep.findPlug('inMesh')
		
		# > inMesh connected 
		if plugInMesh.isConnected() == True:
			PP = []
			PR = []
			PC = []
			PA = []
			PU = []
			PV = []
			
			pArrP = OpenMaya.MPointArray()
			vArrN = OpenMaya.MVectorArray()
			vArrB = OpenMaya.MVectorArray()
			vArrT = OpenMaya.MVectorArray()
			cArrC = OpenMaya.MColorArray()
			fArrA = OpenMaya.MFloatArray()
			
			util = OpenMaya.MScriptUtil()
			util.createFromList([0.0, 0.0], 2)
			ptr = util.asFloat2Ptr()
			dagPath = OpenMaya.MDagPath();
			plugArr = OpenMaya.MPlugArray()
			plugInMesh.connectedTo(plugArr, True, False)
			OpenMaya.MFnDagNode(plugArr[0].node()).getPath(dagPath) # < --- replace by mobject
			#oInMesh = plugInMesh.asMObject()
			fnMesh = OpenMaya.MFnMesh(dagPath)
			itMesh = OpenMaya.MItMeshVertex(dagPath)
			# point positions
			fnMesh.getPoints(pArrP, OpenMaya.MSpace.kWorld)
			for i in range(0, pArrP.length()):
				PP.append(Vector(pArrP[i].x, pArrP[i].y, pArrP[i].z))
				try:
					itMesh.next()
					itMesh.getUV(ptr)
					PU.append(util.getFloat2ArrayItem(ptr, 0, 0))
					PV.append(util.getFloat2ArrayItem(ptr, 0, 1))
				except:
					PU.append(0.0)
					PV.append(0.0)
			# point rotations
			nArrConnFaceIds = OpenMaya.MIntArray()
			itMesh.reset()
			for i in range(0, pArrP.length()):
				itMesh.getConnectedFaces(nArrConnFaceIds)
				if nArrConnFaceIds.length() == 0:
					PR.append(Vector(0.0, 0.0, 0.0))
				else:
					try:
						vN = OpenMaya.MVector()
						vT = OpenMaya.MVector()
						vB = OpenMaya.MVector()
						itMesh.getNormal(vN, OpenMaya.MSpace.kWorld)
						if itMesh.onBoundary() == True:
							fnMesh.getFaceVertexTangent(nArrConnFaceIds[0], i, vT, OpenMaya.MSpace.kWorld)
						else:
							vT2 = OpenMaya.MVector()
							for nFaceId in nArrConnFaceIds:
								fnMesh.getFaceVertexTangent(nFaceId, i, vT2, OpenMaya.MSpace.kWorld)
								vT+=vT2
						if vT.length() == 0:
							fnMesh.getFaceVertexTangent(nArrConnFaceIds[0], i, vT, OpenMaya.MSpace.kWorld)
						vT.normalize()
						a = numpy.cross([vN[0],vN[1],vN[2]], [vT[0],vT[1],vT[2]])
						vB.x = a[0]; vB.y = a[1]; vB.z = a[2]
						vB.normalize()
						a = numpy.cross([vN.x,vN.y,vN.z], [vB.x,vB.y,vB.z])
						vT.x = a[0]; vT.y = a[1]; vT.z = a[2]
						vT.normalize()
						m = OpenMaya.MMatrix()
						listM = (vN.x, vN.y, vN.z, 0.0, vT.x, vT.y, vT.z, 0.0, vB.x, vB.y, vB.z, 0.0, 0.0, 0.0, 0.0, 1.0)
						OpenMaya.MScriptUtil().createMatrixFromList(listM, m)
						eR = OpenMaya.MTransformationMatrix(m).eulerRotation()
						PR.append(Vector(eR[0] * 57.2958, eR[1] * 57.2958, eR[2] * 57.2958))
					except:
						PR.append(Vector(0.0, 0.0, 0.0))
				itMesh.next()
			# point colors
			if fnMesh.numColors() > 0:
				fnMesh.getVertexColors(cArrC)
				for i in range(0, cArrC.length()):
					PC.append(Vector(cArrC[i].r, cArrC[i].g, cArrC[i].b))
					PA.append(cArrC[i].a)
			myDict["PP"] = PP
			myDict["PR"] = PR
			myDict["PC"] = PC
			myDict["PA"] = PA
			myDict["PU"] = PU
			myDict["PV"] = PV
		# < end inMesh connected

		# TEXTURE INPUT
		TC = []
		TA = []
		myDict['TC'] = []
		myDict['TA'] = []
		meshIsect = OpenMaya.MMeshIntersector()
		pom = OpenMaya.MPointOnMesh()
		p = OpenMaya.MPoint()
		plugInMesh = fnDep.findPlug('inMesh')
		bIsConn = plugInMesh.isConnected()
		if bIsConn == True:
			oInMesh = plugInMesh.asMObject()
			meshIsect.create(oInMesh)
		if omfx.MDynamicsUtil().hasValidDynamics2dTexture(oNode, pyExpressionNode.inColor) == True:
			dArrU = OpenMaya.MDoubleArray()
			dArrV = OpenMaya.MDoubleArray()
			for v in V:
				# if nothing is connected to the inMesh attribute, use the world XZ plane as coordinate system
				if bIsConn == False:
					dArrU.append(v.x)
					dArrV.append(v.z)
				else: # otherwise find the closest point to the mesh and get its UV coordinates
					p.x = v.x; p.y = v.y; p.z = v.z
					meshIsect.getClosestPoint(p, pom, 999999999.9)
					p = OpenMaya.MPoint(pom.getPoint())
					fnMesh.getUVAtPoint(p, ptr, OpenMaya.MSpace.kWorld)
					dArrU.append(util.getFloat2ArrayItem(ptr, 0, 0))
					dArrV.append(util.getFloat2ArrayItem(ptr, 0, 1))
			vArrC = OpenMaya.MVectorArray()
			dArrA = OpenMaya.MDoubleArray()
			omfx.MDynamicsUtil().evalDynamics2dTexture(oNode, pyExpressionNode.inColor, dArrU, dArrV, vArrC, dArrA)
			for i in range(0, vArrC.length()):
				TC.append(Vector(vArrC[i].x, vArrC[i].y, vArrC[i].z))
				TA.append(dArrA[i])
			myDict['TC'] = TC
			myDict['TA'] = TA

		#
		# GET OUTPUTS
		#

		# FLOAT OUTPUT
		OF = []
		hArrData = dataBlock.outputArrayValue(pyExpressionNode.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(pyExpressionNode.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(pyExpressionNode.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

		#
		# SET OUTPUTS
		#

		# FLOAT OUTPUT
		if  plug == pyExpressionNode.OF:

			hArrData = dataBlock.outputArrayValue(pyExpressionNode.OF)
			bData = OpenMaya.MArrayDataBuilder(pyExpressionNode.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 == pyExpressionNode.OV:

			hArrData = dataBlock.outputArrayValue(pyExpressionNode.OV)
			bData = OpenMaya.MArrayDataBuilder(pyExpressionNode.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])
				else: continue

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

		# MATRIX OUTPUT
		if  plug == pyExpressionNode.OM:

			hArrData = dataBlock.outputArrayValue(pyExpressionNode.OM)
			bData = OpenMaya.MArrayDataBuilder(pyExpressionNode.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):
					hData = bData.addElement(i)
					hData.setMMatrix(super(Matrix,om))
				else: continue
			hArrData.set(bData)
			hArrData.setAllClean()

# initializer
def nodeInitializer():
	'''
	# local variables:
	================
	
	OS - this node name
	
	INPUT ATTRIBUTES:
	T (float) - input time attribute value
	N (list of strings) - names of nodes connected to the input node attributes
	F (list of floats) - input float attribute values
	V (list of vectors) - input vector attribute values
	M (list of matrices) - input matrix attribute values
	
	INPUT MESH DATA:
	PP (list of vectors) - point positions
	PR (list of vectors) - point rotations
	PC (list of vectors) - point colors - red, green, blue
	PA (list of floats) - point colors - alpha
	PU (list of floats) - point U coordinates
	PU (list of floats) - point V coordinates
	
	INPUT TEXTURE DATA:
	TC (list of vectors) - red, green, blue color values per input vector attribute
	TA (list of floats) - alpha values per input vector attribute
	
	OUTPUT ATTRIBUTES:
	OF (list of floats) - output float attribute values
	OV (list of vectors) - output vector attribute values
	OM (list of matrices) - output matrix attribute values
	
	
	built-in classes:
	=================
	
	In order to simplify the syntax and streamline
	the workflow we provide two built-in classes -
	'Vector' and 'Matrix' based on the C++ API ones
	MVector and MTransformationMatrix.
	
	For example, instead of doing this:
	import maya.OpenMaya as OpenMaya
	vec = om.MVector(1.0, 2.0, 3.0)
	print vec.length()
	we can do:
	vec = Vector(1.0, 2.0, 3.0)
	print vec.length()
	
	Similar when working with matrices.
	For example, instead of doing this:
	import maya.OpenMaya as OpenMaya
	listMat = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0]
	mat = OpenMaya.MMatrix()
	OpenMaya.MScriptUtil().createMatrixFromList(listMat, mat)
	tMat = OpenMaya.MTransformationMatrix(mat)
	matScl = tMat.asScaleMatrix()()
	print matScl(0,0), matScl(1,1), matScl(2,2)
	we can do:
	mat = Matrix(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0)
	print mat.getScale()
	
	Both classes inherit all methods of their
	\"parents\" (MVector or MTransformationMatrix).
	The Matrix class has 3 additional methods:
	getTranslate(), getRotate() and getScale().
	All they return Vectors.
	getRotate() returns euler values that can be
	directly plugged into rotate attributes of
	transform/joint nodes.
	
	'''
	
	gAttr = OpenMaya.MFnGenericAttribute()
	cAttr = OpenMaya.MFnCompoundAttribute()
	nAttr = OpenMaya.MFnNumericAttribute()
	mAttr = OpenMaya.MFnMatrixAttribute()
	tAttr = OpenMaya.MFnTypedAttribute()
	meAttr = OpenMaya.MFnMessageAttribute()

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

	# NOTES
	oString = OpenMaya.MFnStringData().create( nodeInitializer.__doc__ )
	pyExpressionNode.notes = tAttr.create("notes", "notes", OpenMaya.MFnData.kString, oString)
	tAttr.setStorable(True)
	tAttr.setKeyable(False)
	pyExpressionNode.addAttribute( pyExpressionNode.notes )
	
	# TIME
	pyExpressionNode.time = nAttr.create("time", "time", OpenMaya.MFnNumericData.kDouble, 0.0)
	nAttr.setKeyable(True)
	pyExpressionNode.addAttribute(pyExpressionNode.time)

	# INPUT MESH
	pyExpressionNode.inMesh = gAttr.create("inMesh", "inMesh")
	gAttr.addDataAccept(OpenMaya.MFnData.kMesh)
	pyExpressionNode.addAttribute(pyExpressionNode.inMesh)

	# INPUT TEXTURE
	pyExpressionNode.inColor = nAttr.createColor("inColor", "inColor")
	pyExpressionNode.addAttribute(pyExpressionNode.inColor)

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

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

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

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

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

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

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

	# RELATIONSHIPS
	pyExpressionNode.attributeAffects(pyExpressionNode.time, pyExpressionNode.OF)
	pyExpressionNode.attributeAffects(pyExpressionNode.time, pyExpressionNode.OV)
	pyExpressionNode.attributeAffects(pyExpressionNode.time, pyExpressionNode.OM)
	pyExpressionNode.attributeAffects(pyExpressionNode.inMesh, pyExpressionNode.OF)
	pyExpressionNode.attributeAffects(pyExpressionNode.inMesh, pyExpressionNode.OV)
	pyExpressionNode.attributeAffects(pyExpressionNode.inMesh, pyExpressionNode.OM)
	pyExpressionNode.attributeAffects(pyExpressionNode.inColor, pyExpressionNode.OF)
	pyExpressionNode.attributeAffects(pyExpressionNode.inColor, pyExpressionNode.OV)
	pyExpressionNode.attributeAffects(pyExpressionNode.inColor, pyExpressionNode.OM)
	pyExpressionNode.attributeAffects(pyExpressionNode.N, pyExpressionNode.OF)
	pyExpressionNode.attributeAffects(pyExpressionNode.N, pyExpressionNode.OV)
	pyExpressionNode.attributeAffects(pyExpressionNode.N, pyExpressionNode.OM)
	pyExpressionNode.attributeAffects(pyExpressionNode.F, pyExpressionNode.OF)
	pyExpressionNode.attributeAffects(pyExpressionNode.F, pyExpressionNode.OV)
	pyExpressionNode.attributeAffects(pyExpressionNode.F, pyExpressionNode.OM)
	pyExpressionNode.attributeAffects(pyExpressionNode.V, pyExpressionNode.OF)
	pyExpressionNode.attributeAffects(pyExpressionNode.V, pyExpressionNode.OV)
	pyExpressionNode.attributeAffects(pyExpressionNode.V, pyExpressionNode.OM)
	pyExpressionNode.attributeAffects(pyExpressionNode.M, pyExpressionNode.OF)
	pyExpressionNode.attributeAffects(pyExpressionNode.M, pyExpressionNode.OV)
	pyExpressionNode.attributeAffects(pyExpressionNode.M, pyExpressionNode.OM)
