#include <maya/MIOStream.h>
#include <math.h>
#include <stdlib.h>
#include <sweptEmitter.h>
#include <maya/MDataHandle.h>
#include <maya/MFnDynSweptGeometryData.h>
#include <maya/MDynSweptLine.h>
#include <maya/MDynSweptTriangle.h>
#include <maya/MVectorArray.h>
#include <maya/MDoubleArray.h>
#include <maya/MIntArray.h>
#include <maya/MMatrix.h>
#include <maya/MFnDependencyNode.h>
#include <maya/MFnNumericAttribute.h>
#include <maya/MFnUnitAttribute.h>
#include <maya/MFnVectorArrayData.h>
#include <maya/MFnDoubleArrayData.h>
#include <maya/MFnArrayAttrsData.h>
#include <maya/MFnMatrixData.h>
MTypeId sweptEmitter::id( 0x80016 );
sweptEmitter::sweptEmitter()
:       lastWorldPoint(0, 0, 0, 1)
{
}
sweptEmitter::~sweptEmitter()
{
}
void *sweptEmitter::creator()
{
    return new sweptEmitter;
}
MStatus sweptEmitter::initialize()
{
        return( MS::kSuccess );
}
MStatus sweptEmitter::compute(const MPlug& plug, MDataBlock& block)
{
        MStatus status;
        
        
        if( !(plug == mOutput) )
        return( MS::kUnknownParameter );
        
        
    
        
        int multiIndex = plug.logicalIndex( &status );
        McheckErr(status, "ERROR in plug.logicalIndex.\n");
        
        
        
        MArrayDataHandle hOutArray = block.outputArrayValue(mOutput, &status);
        McheckErr(status, "ERROR in hOutArray = block.outputArrayValue.\n");
        
        
        MArrayDataBuilder bOutArray = hOutArray.builder( &status );
        McheckErr(status, "ERROR in bOutArray = hOutArray.builder.\n");
        
        
        MDataHandle hOut = bOutArray.addElement(multiIndex, &status);
        McheckErr(status, "ERROR in hOut = bOutArray.addElement.\n");
    
    
    MFnArrayAttrsData fnOutput;
    MObject dOutput = fnOutput.create ( &status );
    McheckErr(status, "ERROR in fnOutput.create.\n");
        
        
        
        bool beenFull = isFullValue( multiIndex, block );
        if( beenFull )
        {
                return( MS::kSuccess );
        }
        
        
        
        
        MTime cT = currentTimeValue( block );
        MTime sT = startTimeValue( multiIndex, block );
        MTime dT = deltaTimeValue( multiIndex, block );
        if( (cT <= sT) || (dT <= 0.0) )
        {
                
                
                
                
                
                
                
                
                hOut.set( dOutput );
                block.setClean( plug );
                return( MS::kSuccess );
        }
        
        
        double speed = speedValue( block );
        MVector dirV = directionVector( block );
        double inheritFactor = inheritFactorValue( multiIndex, block );
        
        
        MVectorArray fnOutPos = fnOutput.vectorArray("position", &status);
        MVectorArray fnOutVel = fnOutput.vectorArray("velocity", &status);
        
        
        double dt = dT.as( MTime::kSeconds );
        
        
        MVector rotatedV = useRotation ( dirV );
        
        MVectorArray inPosAry;
        
        MVectorArray inVelAry;
        
        MIntArray emitCountPP;
        
        
        MObject thisObj = this->thisMObject();
        MPlug sweptPlug( thisObj, mSweptGeometry );
        if ( sweptPlug.isConnected() ) 
        {
                MDataHandle sweptHandle = block.inputValue( mSweptGeometry );
                
                MObject sweptData = sweptHandle.data();
                MFnDynSweptGeometryData fnSweptData( sweptData );
                
                
                if (fnSweptData.lineCount() > 0) {
                        int numLines = fnSweptData.lineCount();
                
                        for ( int i=0; i<numLines; i++ )
                        {
                                inPosAry.clear();
                                inVelAry.clear();
                                emitCountPP.clear();
                                MDynSweptLine line = fnSweptData.sweptLine( i );
                                
                                MVector p1 = line.vertex( 0 );
                                MVector p2 = line.vertex( 1 );
                                inPosAry.append( p1 );
                                inPosAry.append( p2 );
                                inVelAry.append( MVector( 0,0,0 ) );
                                inVelAry.append( MVector( 0,0,0 ) );
                                
                                emitCountPP.clear();
                                status = emitCountPerPoint( plug, block, 2, emitCountPP );
                                emit( inPosAry, inVelAry, emitCountPP,
                                        dt, speed, inheritFactor, rotatedV, fnOutPos, fnOutVel );
                        }
                }
                
                
                if (fnSweptData.triangleCount() > 0) {
                        int numTriangles = fnSweptData.triangleCount();
                
                        for ( int i=0; i<numTriangles; i++ )
                        {
                                inPosAry.clear();
                                inVelAry.clear();
                                emitCountPP.clear();
                                MDynSweptTriangle tri = fnSweptData.sweptTriangle( i );
                                
                                MVector p1 = tri.vertex( 0 );
                                MVector p2 = tri.vertex( 1 );
                                MVector p3 = tri.vertex( 2 );
                                MVector center = p1 + p2 + p3;
                                center /= 3.0;
                                inPosAry.append( center );
                                inVelAry.append( MVector( 0,0,0 ) );
                                
                                emitCountPP.clear();
                                status = emitCountPerPoint( plug, block, 1, emitCountPP );
                                emit( inPosAry, inVelAry, emitCountPP,
                                        dt, speed, inheritFactor, rotatedV, fnOutPos, fnOutVel );
                        }
                }
        }
        
        
        hOut.set( dOutput );
        block.setClean( plug );
        return( MS::kSuccess );
}
void sweptEmitter::emit
        (
                const MVectorArray &inPosAry,   
                const MVectorArray &inVelAry,   
                const MIntArray &emitCountPP,   
                double dt,                                              
                double speed,                                   
                double inheritFactor,                   
                MVector dirV,                                   
                MVectorArray &outPosAry,                
                MVectorArray &outVelAry                 
        )
{
        
        
        int posLength = inPosAry.length();
        int velLength = inVelAry.length();
        int countLength = emitCountPP.length();
        if( (posLength != velLength) || (posLength != countLength) )
                return;
        
        
        int index;
        int totalCount = 0;
        for( index = 0; index < countLength; index ++ )
                totalCount += emitCountPP[index];
        if( totalCount <= 0 )
                return;
        
        
        dirV.normalize();
        
        
        int emitCount;
        MVector newPos, newVel;
        MVector prePos, sPos, sVel;
        for( index = 0; index < posLength; index++ )
        {
                emitCount = emitCountPP[index];
                if( emitCount <= 0 )
                        continue;
                sPos = inPosAry[index];
                sVel = inVelAry[index];
                prePos = sPos - sVel * dt;
                for( int i = 0; i < emitCount; i++ )
                {
                        double alpha = ( (double)i + drand48() ) / (double)emitCount;
                        newPos = (1 - alpha) * prePos + alpha * sPos;
                        newVel = dirV * speed;
                        newPos += newVel * ( dt * (1 - alpha) );
                        newVel += sVel * inheritFactor;
                        
                        
                        outPosAry.append( newPos );
                        outVelAry.append( newVel );
                }
        }
}
MVector sweptEmitter::useRotation ( MVector &direction )
{
        MStatus status;
        MVector rotatedVector;
        MObject thisNode = thisMObject();
        MFnDependencyNode fnThisNode( thisNode );
        
        
        MObject worldMatrixAttr = fnThisNode.attribute( "worldMatrix" );
        
        
        
        MPlug matrixPlug( thisNode, worldMatrixAttr );
        matrixPlug = matrixPlug.elementByLogicalIndex( 0 );
        
        
        MObject matrixObject;
        status = matrixPlug.getValue( matrixObject );
        if( !status )
        {
                status.perror("sweptEmitter::useRotation: get matrixObject");
                return ( direction );
        }
        MFnMatrixData worldMatrixData( matrixObject, &status );
        if( !status )
        {
                status.perror("sweptEmitter::useRotation: get worldMatrixData");
                return( direction );
        }
        MMatrix worldMatrix = worldMatrixData.matrix( &status );
        if( !status )
        {
                status.perror("sweptEmitter::useRotation: get worldMatrixData.matrix");
                return( direction );
        }
        rotatedVector = direction * worldMatrix;
    return( rotatedVector );
}
MStatus sweptEmitter::emitCountPerPoint
        (
                const MPlug &plug,
                MDataBlock &block,
                int length,                                     
                MIntArray &emitCountPP          
        )
{
        MStatus status;
        int plugIndex = plug.logicalIndex( &status );
        McheckErr(status, "ERROR in emitCountPerPoint: when plug.logicalIndex.\n");
        
        
        double rate = rateValue( block );
        MTime dt = deltaTimeValue( plugIndex, block );
        
        
        double dblCount = rate * dt.as( MTime::kSeconds );
        int intCount = (int)dblCount;
        for( int i = 0; i < length; i++ )
        {
                emitCountPP.append( intCount );
        }
        return( MS::kSuccess );
}
MStatus initializePlugin(MObject obj)
{
        MStatus status;
        MFnPlugin plugin(obj, PLUGIN_COMPANY, "3.0", "Any");
        status = plugin.registerNode( "sweptEmitter", sweptEmitter::id,
                                                        &sweptEmitter::creator, &sweptEmitter::initialize,
                                                        MPxNode::kEmitterNode );
        if (!status) {
                status.perror("registerNode");
                return status;
        }
        return status;
}
MStatus uninitializePlugin(MObject obj)
{
        MStatus status;
        MFnPlugin plugin(obj);
        status = plugin.deregisterNode( sweptEmitter::id );
        if (!status) {
                status.perror("deregisterNode");
                return status;
        }
        return status;
}