import 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( pyExpressionCmd() )

class pyExpressionCmd(OpenMayaMPx.MPxCommand):
	
	def __init__(self):
		
		OpenMayaMPx.MPxCommand.__init__(self)

	def doIt(self, argList):

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

		#sNode =  
		#if not cmds.objExists(sNode):
		#	OpenMaya.MGlobal.displayError('// Node not found: '+sNode)
		#	return

		# get node
		oNode = OpenMaya.MObject()
		sList = OpenMaya.MSelectionList()
		sList.add(sNode)
		try:
			sList.getDependNode(0,oNode)
		except e:
			raise RuntimeError, e
		
		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.asDouble()
		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()
			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

		# MESH INPUT
		PP = []
		PR = []
		PC = []
		PA = []
		PU = []
		PV = []
		myDict['PP'] = []
		myDict['PR'] = []
		myDict['PC'] = []
		myDict['PA'] = []
		myDict['PU'] = []
		myDict['PV'] = []
		plugInMesh = fnDep.findPlug('inMesh')
		if(plugInMesh.isConnected() == True):
			pArrP = OpenMaya.MPointArray()
			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)
#			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()
						fnMesh.getFaceVertexNormal(nArrConnFaceIds[0], i, vN, OpenMaya.MSpace.kWorld)
						fnMesh.getFaceVertexTangent(nArrConnFaceIds[0], i, vT, OpenMaya.MSpace.kWorld)
						fnMesh.getFaceVertexBinormal(nArrConnFaceIds[0], i, vB, OpenMaya.MSpace.kWorld)
						vN.normalize()
						vT.normalize()
						vB.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

	   # 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)
		plugInColor = fnDep.findPlug('inColor')
		if omfx.MDynamicsUtil().hasValidDynamics2dTexture(plugInColor.node(), plugInColor.attribute()) == 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(plugInColor.node(), plugInColor.attribute(), 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 = []
		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)
