#-
# ==========================================================================
# Copyright (C) 1995 - 2006 Autodesk, Inc. and/or its licensors.  All 
# rights reserved.
#
# The coded instructions, statements, computer programs, and/or related 
# material (collectively the "Data") in these files contain unpublished 
# information proprietary to Autodesk, Inc. ("Autodesk") and/or its 
# licensors, which is protected by U.S. and Canadian federal copyright 
# law and by international treaties.
#
# The Data is provided for use exclusively by You. You have the right 
# to use, modify, and incorporate this Data into other products for 
# purposes authorized by the Autodesk software license agreement, 
# without fee.
#
# The copyright notices in the Software and this entire statement, 
# including the above license grant, this restriction and the 
# following disclaimer, must be included in all copies of the 
# Software, in whole or in part, and all derivative works of 
# the Software, unless such copies or derivative works are solely 
# in the form of machine-executable object code generated by a 
# source language processor.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND. 
# AUTODESK DOES NOT MAKE AND HEREBY DISCLAIMS ANY EXPRESS OR IMPLIED 
# WARRANTIES INCLUDING, BUT NOT LIMITED TO, THE WARRANTIES OF 
# NON-INFRINGEMENT, MERCHANTABILITY OR FITNESS FOR A PARTICULAR 
# PURPOSE, OR ARISING FROM A COURSE OF DEALING, USAGE, OR 
# TRADE PRACTICE. IN NO EVENT WILL AUTODESK AND/OR ITS LICENSORS 
# BE LIABLE FOR ANY LOST REVENUES, DATA, OR PROFITS, OR SPECIAL, 
# DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES, EVEN IF AUTODESK 
# AND/OR ITS LICENSORS HAS BEEN ADVISED OF THE POSSIBILITY 
# OR PROBABILITY OF SUCH DAMAGES.
#
# ==========================================================================
#+


import maya.OpenMaya as OpenMaya
import maya.OpenMayaMPx as OpenMayaMPx
import maya.OpenMayaRender as OpenMayaRender
import math
import random
import sys

kSimpleEmitterNodeName = "spSimpleEmitter"
kSimpleEmitterNodeID = OpenMaya.MTypeId(0x8700F)

SEGMENTS = 20
EDGES = 30
TORUS_PI = 3.14159265
TORUS_2PI = 2.0 * TORUS_PI

glRenderer = OpenMayaRender.MHardwareRenderer.theRenderer()
glFT = glRenderer.glFunctionTable()

class simpleEmitter(OpenMayaMPx.MPxEmitterNode):
	def __init__(self):
		OpenMayaMPx.MPxEmitterNode.__init__(self)

	def isFullValue( self, plugIndex, block ):
		value = 1
		mIsFull = OpenMayaMPx.cvar.MPxEmitterNode_mIsFull

		try:
			mhValue = block.inputArrayValue( mIsFull )
			mhValue.jumpToElement( plugIndex )
			hValue = mhValue.inputValue( )
			value = hValue.asBool()
		except:
			sys.stderr.write("Error getting the input array value\n")
			raise
		return value
					
	
	def getWorldPosition( self, point ):
		try:
			thisNode = simpleEmitter.thisMObject( self )
			fnThisNode = OpenMaya.MFnDependencyNode(thisNode)

			worldMatrixAttr = fnThisNode.attribute( "worldMatrix" )

			matrixPlug = OpenMaya.MPlug( thisNode, worldMatrixAttr )
			matrixPlug = matrixPlug.elementByLogicalIndex( 0 )

			matrixObject = OpenMaya.MObject()
			matrixObject = matrixPlug.asMObject(  )

			worldMatrixData = OpenMaya.MFnMatrixData( matrixObject )
			worldMatrix = worldMatrixData.matrix( )

			point.x = worldMatrix( 3, 0 )
			point.y = worldMatrix( 3, 1 )
			point.z = worldMatrix( 3, 2 )
		except:
			sys.stderr.write( "Error in getWorldPosition\n" )
			raise

	def currentTimeValue( self, block ):
		try:
			mCurrentTime = OpenMayaMPx.cvar.MPxEmitterNode_mCurrentTime
			hValue = block.inputValue( mCurrentTime )
			value = hValue.asTime()
		except:
			sys.stderr.write("Error getting current time value, returning 0")
			value = OpenMaya.MTime(0.0)
		return value

	def startTimeValue( self, plugIndex, block ):
		try:
			mStartTime = OpenMayaMPx.cvar.MPxEmitterNode_mStartTime
			mhValue = block.inputArrayValue( mStartTime )
			mhValue.jumpToElement( plugIndex )
			hValue = mhValue.inputValue( )
			value = hValue.asTime( )		
		except:
			sys.stderr.write("Error getting start time value, setting to 0")
			value = OpenMaya.MTime(0.0)
		return value

	def deltaTimeValue( self, plugIndex, block ):
		try:
			mDeltaTime = OpenMayaMPx.cvar.MPxEmitterNode_mDeltaTime

			mhValue = block.inputArrayValue( mDeltaTime )
			mhValue.jumpToElement( plugIndex )

			hValue = mhValue.inputValue()
			value = hValue.asTime()
		except:
			sys.stderr.write("Error getting delta time value, setting to 0\n")
			value = OpenMaya.MTime(0.0)
		return value

	def rateValue( self, block ):
		try:
			mRate = OpenMayaMPx.cvar.MPxEmitterNode_mRate
			hValue = block.inputValue( mRate )
			value = hValue.asDouble()		
		except:
			sys.stderr.write("Error getting rate value, setting to 0\n")
			value = 0.0
		return value

	def directionValue( self, block ):
		try:
			mDirection = OpenMayaMPx.cvar.MPxEmitterNode_mDirection
			hValue = block.inputValue( mDirection )
			value = hValue.asDouble3()
			dirV = OpenMaya.MVector(value[0], value[1], value[2])
		except:
			sys.stderr.write("Error getting direction value, setting to 0,0,0\n")
			dirV = OpenMaya.MVector(0.0, 0.0, 0.0)
		return dirV

	def speedValue( self, block ):
		try:
			mSpeed = OpenMayaMPx.cvar.MPxEmitterNode_mSpeed
			hValue = block.inputValue( mSpeed )
			value = hValue.asDouble()		
		except:
			sys.stderr.write("Error getting speed value, setting to 0\n")
			value = 0.0
		return value

	def inheritFactorValue( self, plugIndex, block ):
		try:
			mInheritFactor = OpenMayaMPx.cvar.MPxEmitterNode_mInheritFactor
			mhValue = block.inputArrayValue( mInheritFactor )
			mhValue.jumpToElement( plugIndex )
			hValue = mhValue.inputValue( )
			value = hValue.asDouble()		
		except:
			sys.stderr.write("Error getting inherit factor value, setting to 0\n")
			value = 0.0
		return value

	def useRotation( self, direction ):
		try:
			thisNode = simpleEmitter.thisMObject(self)
			fnThisNode = OpenMaya.MFnDependencyNode( thisNode )

			worldMatrixAttr = fnThisNode.attribute( "worldMatrix" )

			matrixPlug = OpenMaya.MPlug( thisNode, worldMatrixAttr )
			matrixPlug = matrixPlug.elementByLogicalIndex( 0 )

			matrixObject = OpenMaya.MObject()
			matrixObject = matrixPlug.asMObject( )

			worldMatrixData = OpenMaya.MFnMatrixData( matrixObject )
			worldMatrix = worldMatrixData.matrix( )

			rotatedVector = OpenMaya.MVector()
			rotatedVector = direction * worldMatrix
		except:
			sys.stderr.write("Error getting rotation value, setting to 0,0,0\n")
			rotatedVector = OpenMaya.MVector(0.0, 0.0, 0.0)
		return rotatedVector


	def compute(self, plug, block):
		mOutput = OpenMayaMPx.cvar.MPxEmitterNode_mOutput

		# Determine if we are requesting the output plug for this emitter node.
		#
		if plug == mOutput:
			# Get the logical index of the element this plug refers to,
			# because the node can be emitting particles into more than
			# one particle shape.
			#
			try:
				multiIndex = plug.logicalIndex( )

				# Get output data arrays (position, velocity, or parentId)
				# that the particle shape is holding from the previous frame.
				#
				hOutArray = block.outputArrayValue ( mOutput )

				# Create a builder to aid in the array construction efficiently.
				#
				bOutArray = hOutArray.builder( )

				# Get the appropriate data array that is being currently evaluated.
				#
				hOut = bOutArray.addElement( multiIndex )

				# Create the data and apply the function set,
				# particle array initialized to length zero, 
				# fnOutput.clear()
				#
				fnOutput = OpenMaya.MFnArrayAttrsData()
				dOutput = fnOutput.create( )

				# Check if the particle object has reached it's maximum,
				# hence is full. If it is full then just return with zero particles.
				#
				beenFull = simpleEmitter.isFullValue( self, multiIndex, block )
				if beenFull == 1:
					return

				# Get input position and velocity arrays where new particles are from,
				# also known as the owner. An owner is determined if connections exist
				# to the emitter node from a shape such as nurbs, polymesh, curve, 
				# or a lattice shape.
				#
				
				# Get a single position from world transform
				#
				inPosAry = OpenMaya.MVectorArray()
				inPosAry.clear()
			
				worldPos = OpenMaya.MPoint(0.0, 0.0, 0.0)
				simpleEmitter.getWorldPosition( self, worldPos )

				worldV = OpenMaya.MVector(worldPos[0], worldPos[1], worldPos[2])
				inPosAry.append( worldV )

				# Create a single velocity			
				inVelAry = OpenMaya.MVectorArray()
				inVelAry.clear()
				velocity = OpenMaya.MVector(0, 0, 0)
				inVelAry.append( velocity )

				# Get deltaTime, currentTime and startTime.
				# If deltaTime <= 0.0, or currentTime <= startTime,
				# do not emit new pariticles and return.
				#
				cT = simpleEmitter.currentTimeValue( self, block )
				sT = simpleEmitter.startTimeValue( self, multiIndex, block )
				dT = simpleEmitter.deltaTimeValue( self, multiIndex, block )

				dTValue = dT.value()
			
				if cT <= sT or dTValue <= 0.0:
					# We do not emit particles before the start time, 
					# and do not emit particles when moving backwards in time.
					# 

					# This code is necessary primarily the first time to 
					# establish the new data arrays allocated, and since we have 
					# already set the data array to length zero it does 
					# not generate any new particles.
					# 
					hOut.setMObject( dOutput )
					block.setClean( plug )
					return

				# Compute and store an emission rate
				#
				emitCountPP = OpenMaya.MIntArray()
				emitCountPP.clear()

				plugIndex = plug.logicalIndex( )

				# Get rate and delta time.
				#
				rate = simpleEmitter.rateValue( self, block )
				dtRate = simpleEmitter.deltaTimeValue( self, plugIndex, block )
				dtRateDbl = dtRate.as( OpenMaya.MTime.kSeconds )
				dblCount = rate * dtRateDbl
				intCount = int(dblCount)
				emitCountPP.append( intCount )

				# Get speed, direction vector, and inheritFactor attributes.
				#			
				speed = simpleEmitter.speedValue( self, block )
				dirV = simpleEmitter.directionValue( self, block )
				inheritFactor = simpleEmitter.inheritFactorValue( self, multiIndex, block )

				# Get the position, velocity, and normalized time arrays to append new particle data.
				#
				fnOutPos = fnOutput.vectorArray( "position" )
				fnOutVel = fnOutput.vectorArray( "velocity" )
				fnOutTime = fnOutput.doubleArray( "timeInStep" )

				# Convert deltaTime into seconds.
				#			
				dt = dT.as( OpenMaya.MTime.kSeconds )

				# Rotate the direction attribute by world transform
				rotatedV = simpleEmitter.useRotation( self, dirV )

				# Start emitting particles.
				#
				simpleEmitter.emit( self, inPosAry, inVelAry, emitCountPP, dt, speed, inheritFactor,\
						rotatedV, fnOutPos, fnOutVel, fnOutTime)

				# Update the data block with new dOutput and set plug clean.
				#	
				# sys.__stdout__.write( " handle: " + str(hOut) + "\n" )		
				hOut.setMObject( dOutput )
				block.setClean( plug )
			except:
				sys.stderr.write("simpleEmitter compute error\n")
				raise
		else:
			return OpenMaya.kUnknownParameter

	def emit( self, inPosAry, inVelAry, emitCountPP, dt, speed, inheritFactor, dirV, outPosAry, outVelAry, outTimeAry):
		try:
			posLength = inPosAry.length()
			velLength = inVelAry.length()
			countLength = emitCountPP.length()

			if not posLength == velLength or not posLength == countLength:
				return

			totalCount = 0
			for index in range(countLength):
				totalCount += emitCountPP[index]
			if totalCount <= 0:
				return

			dirV.normalize()
			
			for index in range(posLength):
				emitCount = emitCountPP[index]
				if emitCount <= 0:
					continue
			
				sPos = inPosAry[index]
				sVel = inVelAry[index]
				prePos = sPos - sVel * dt

				for i in range(emitCount):
					alpha = ( float(i) + random.random() ) / float(emitCount)
					newPos = prePos * (1.0 - alpha) + sPos * alpha
					newVel = dirV * speed

					newPos += newVel * ( dt * (1.0 - alpha) )
					newVel += sVel * inheritFactor

					# Add new data into output arrays.
					#
					outPosAry.append( newPos )
					outVelAry.append( newVel )
					outTimeAry.append( alpha )
		except Exception, e:
			sys.stderr.write( "Error in simpleEmitter.emit\n" )
			raise

	def draw( self, view, path, style, status):
		view.beginGL()

		for j in range(0, SEGMENTS):
			glFT.glPushMatrix()
			glFT.glRotatef(float(360.0 * j/SEGMENTS), 0.0, 1.0, 0.0)
			glFT.glTranslatef( 1.5, 0.0, 0.0)
			
			for i in range(0, EDGES):
				glFT.glBegin(OpenMayaRender.MGL_LINE_STRIP)

				p0 = float(TORUS_2PI * i/EDGES)
				p1 = float(TORUS_2PI * (i+1)/EDGES)
				glFT.glVertex2f(math.cos(p0), math.sin(p0))
				glFT.glVertex2f(math.cos(p1), math.sin(p1))
				
				glFT.glEnd()

			glFT.glPopMatrix()

		view.endGL()


def nodeCreator():
	return OpenMayaMPx.asMPxPtr( simpleEmitter() )

def nodeInitializer():
	return

# initialize the script plug-in
def initializePlugin(mobject):
	mplugin = OpenMayaMPx.MFnPlugin(mobject)

	try:
		mplugin.registerNode( kSimpleEmitterNodeName, kSimpleEmitterNodeID, \
								nodeCreator, nodeInitializer, OpenMayaMPx.MPxNode.kEmitterNode )
	except:
		sys.stderr.write( "Failed to register node: %s\n" % kSimpleEmitterNodeName )
		raise

# uninitialize the script plug-in
def uninitializePlugin(mobject):
	mplugin = OpenMayaMPx.MFnPlugin(mobject)

	try:
		mplugin.deregisterNode( kSimpleEmitterNodeID )
	except:
		sys.stderr.write( "Failed to unregister node: %s\n" % kSimpleEmitterNodeName )
		raise