#include <maya/MIOStream.h>
#include <math.h>
#include <stdlib.h>
#include <simpleFluidEmitter.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>
#include <maya/MDagPath.h>
#include <maya/MMatrix.h>
#include <maya/MTransformationMatrix.h>
#include <maya/MFnDynSweptGeometryData.h>
#include <maya/MDynSweptTriangle.h>
MTypeId simpleFluidEmitter::id( 0x81020 );
simpleFluidEmitter::simpleFluidEmitter()
{
}
simpleFluidEmitter::~simpleFluidEmitter()
{
}
void *simpleFluidEmitter::creator()
{
    return new simpleFluidEmitter;
}
MStatus simpleFluidEmitter::initialize()
{
        return( MS::kSuccess );
}
MStatus simpleFluidEmitter::compute(const MPlug& plug, MDataBlock& block)
{
        if( plug.attribute() == mEmissionFunction )
        {
                
                
                return MS::kUnknownParameter;
        }
        else
        {
                
                
                return MS::kUnknownParameter;
        }
}
MStatus 
simpleFluidEmitter::fluidEmitter( 
        const MObject& fluidObject, 
        const MMatrix& worldMatrix, 
        int plugIndex 
)
{
        
        
        
        
        
        MFnFluid fluid( fluidObject );
        if( fluid.object() != MObject::kNullObj )
        {
                return MS::kSuccess;
        }
        
        
        MDataBlock block = forceCache();
        
        
        double dTime = getDeltaTime( plugIndex, block ).as(MTime::kSeconds);
        if( dTime == 0.0 )
        {
                
                
                return MS::kSuccess;
        }
        
        
        
        MTime cTime = getCurrentTime( block );
        MTime sTime = getStartTime( plugIndex, block );
        
        
        
        if( cTime < sTime )
        {
                resetRandomState( plugIndex, block );
                return MS::kSuccess;
        }
        
        
        
        
        
        
        double density = fluidDensityEmission( block ); 
        double heat = fluidHeatEmission( block );       
        double fuel = fluidFuelEmission( block );       
        bool doColor = fluidEmitColor( block ); 
        
        MFnFluid::FluidMethod densityMode, tempMode, fuelMode;
        MFnFluid::ColorMethod colorMode;
        MFnFluid::FluidGradient grad;
        MFnFluid::FalloffMethod falloffMode;
        fluid.getDensityMode( densityMode, grad );
        fluid.getTemperatureMode( tempMode, grad );
        fluid.getFuelMode( fuelMode, grad );
        fluid.getColorMode( colorMode );
        fluid.getFalloffMode( falloffMode );
        
        bool densityToEmit = (density != 0.0) && ((densityMode == MFnFluid::kDynamicGrid)||(densityMode == MFnFluid::kStaticGrid));
        bool heatToEmit = (heat != 0.0) && ((tempMode == MFnFluid::kDynamicGrid)||(tempMode == MFnFluid::kStaticGrid));
        bool fuelToEmit = (fuel != 0.0) && ((fuelMode == MFnFluid::kDynamicGrid)||(fuelMode == MFnFluid::kStaticGrid));
        bool colorToEmit = doColor && ((colorMode == MFnFluid::kDynamicColorGrid)||(colorMode == MFnFluid::kStaticColorGrid));
        bool falloffEmit = (falloffMode == MFnFluid::kStaticFalloffGrid);
        
        
        
        if( !densityToEmit && !heatToEmit && !fuelToEmit && !colorToEmit && !falloffEmit )
        {
                return MS::kSuccess;
        }
        
        
        double dropoff = fluidDropoff( block );
        
        
        
        
        MTransformationMatrix xform( worldMatrix );
        double xformScale[3];
        xform.getScale( xformScale, MSpace::kWorld );
        double dropoffScale = sqrt( xformScale[0]*xformScale[0] + 
                                                                xformScale[1]*xformScale[1] + 
                                                                xformScale[2]*xformScale[2] );
        if( dropoffScale > 0.1 )
        {
                dropoff /= dropoffScale;
        }
        
        
        
        
        
        
        getRandomState( plugIndex, block );
        
        
        
        double conversion = 0.01;
        MEmitterType emitterType = getEmitterType( block );
        switch( emitterType )
        {
                case kOmni:
                        omniFluidEmitter( fluid, worldMatrix, plugIndex, block, dTime,
                                                          conversion, dropoff );
                        break;
                case kVolume:
                        volumeFluidEmitter( fluid, worldMatrix, plugIndex, block, dTime,
                                                            conversion, dropoff );
                        break;
                        
                case kSurface:
                        surfaceFluidEmitter( fluid, worldMatrix, plugIndex, block, dTime,
                                                             conversion, dropoff );
                        break;
                default:
                        break;
        }
        
        
        setRandomState( plugIndex, block );
        return MS::kSuccess;
}
#define MIN(x,y) ((x)<(y)?(x):(y))
#define MAX(x,y) ((x)>(y)?(x):(y))
void 
simpleFluidEmitter::omniFluidEmitter(
        MFnFluid&               fluid,
        const MMatrix&  fluidWorldMatrix,
        int                     plugIndex,
        MDataBlock&     block,
        double                  dt,
        double                  conversion,
        double                  dropoff
)
{
        
        
        MVectorArray emitterPositions;
        
        
        
        
        
        bool gotOwnerPositions = false;
        MObject ownerShape = getOwnerShape();
        if( ownerShape != MObject::kNullObj )
        {
                MStatus status;
                MDataHandle hOwnerPos = block.inputValue( mOwnerPosData, &status );
                if( status == MS::kSuccess )
                {
                        MObject dOwnerPos = hOwnerPos.data();
                        MFnVectorArrayData fnOwnerPos( dOwnerPos );
                        MVectorArray posArray = fnOwnerPos.array( &status );
                        if( status == MS::kSuccess )
                        {
                                
                                
                                for( unsigned int i = 0; i < posArray.length(); i ++ )
                                {
                                        emitterPositions.append( posArray[i] );
                                }
                                
                                gotOwnerPositions = true;
                        }
                }
        }
        
        
        
        
        if( !gotOwnerPositions )
        {
                MPoint emitterPos = getWorldPosition();
                emitterPositions.append( emitterPos );
        }
        
        
        double densityEmit = fluidDensityEmission( block );
        double fuelEmit = fluidFuelEmission( block );
        double heatEmit = fluidHeatEmission( block );
        bool doEmitColor = fluidEmitColor( block );
        MColor emitColor = fluidColor( block );
        
        
        
        
        double theRate = getRate(block) * dt * conversion;
        
        
        double size[3];
        unsigned int res[3];
        fluid.getDimensions( size[0], size[1], size[2] );
        fluid.getResolution( res[0], res[1], res[2] );
        
        
        double dx = size[0] / res[0];
        double dy = size[1] / res[1];
        double dz = size[2] / res[2];
        
        
        double Ox = -size[0]/2;
        double Oy = -size[1]/2;
        double Oz = -size[2]/2; 
        
        
        
        double minDist = getMinDistance( block );
        double maxDist = getMaxDistance( block );
        
        
        
        
        
        MTransformationMatrix fluidXform( fluidWorldMatrix );
        double fluidScale[3];
        fluidXform.getScale( fluidScale, MSpace::kWorld );
        
        
        double wsX =  fabs(fluidScale[0]*dx);
        double wsY = fabs(fluidScale[1]*dy);
        double wsZ = fabs(fluidScale[2]*dz);
        double wsMin = MIN( MIN( wsX, wsY), wsZ );
        double wsMax = MAX( MAX( wsX, wsY), wsZ );
        double wsDiag  = wsMin * sqrt(3.0);
        
        if ( maxDist <= minDist || maxDist <= (wsDiag/2.0) ) {
                if ( minDist < 0 ) minDist = 0;
                maxDist = minDist + wsDiag/2.0;
                dropoff = 0;
        }
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        
        int numSamples = 1;
        
        if(wsMin >.00001) 
        {
                numSamples = (int)(wsMax/wsMin + .5);
                if(numSamples > 8) 
                        numSamples = 8;
                if(numSamples < 1)
                        numSamples = 1;
        }
        
        bool jitter =  fluidJitter(block);
        if( !jitter )
        {
                
                
                
                
                numSamples = 1;
        }
        for( unsigned int p = 0; p < emitterPositions.length(); p++ )
        {
                MPoint emitterWorldPos = emitterPositions[p];
                
                
                
                for( unsigned int i = 0; i < res[0]; i++ )
                {
                        double x = Ox + i*dx;
                        
                        for( unsigned int j = 0; j < res[1]; j++ )
                        {
                                double y = Oy + j*dy;
                                
                                for( unsigned int k = 0; k < res[2]; k++ )
                                {
                                        double z = Oz + k*dz;
        
                                        int si;
                                        for( si = 0; si < numSamples; si++ )
                                        {
                                                
                                                
                                                double rx, ry, rz;
                                                if( jitter )
                                                {
                                                        rx = x + randgen()*dx;
                                                        ry = y + randgen()*dy;
                                                        rz = z + randgen()*dz;
                                                }
                                                else
                                                {
                                                        rx = x + 0.5*dx;
                                                        ry = y + 0.5*dy;
                                                        rz = z + 0.5*dz;
                                                }
                                                
                                                
                                                MPoint point( rx, ry, rz );
                                                point *= fluidWorldMatrix;
                                                MVector diff = point - emitterWorldPos;
                                                double distSquared = diff * diff;
                                                double dist = diff.length();
                                        
                                                
                                                
                                                if( (dist < minDist) || (dist > maxDist) )
                                                {
                                                        continue;
                                                }
                                                
                                                
                                                
                                                
                                                
                                                double distDrop = dropoff * distSquared;
                                                double newVal = theRate * exp( -distDrop ) / (double)numSamples;
                                                
                                                
                                                if( newVal != 0 )
                                                {
                                                        fluid.emitIntoArrays( (float) newVal, i, j, k, (float)densityEmit, (float)heatEmit, (float)fuelEmit, doEmitColor, emitColor );
                                                }
                                                float *fArray = fluid.falloff();
                                                if( fArray != NULL )
                                                {
                                                        MPoint midPoint( x+0.5*dx, y+0.5*dy, z+0.5*dz );
                                                        midPoint.x *= 0.2;
                                                        midPoint.y *= 0.2;
                                                        midPoint.z *= 0.2;
                                                        float fdist = (float) sqrt( midPoint.x*midPoint.x + midPoint.y*midPoint.y + midPoint.z*midPoint.z );
                                                        fdist /= sqrtf(3.0f);
                                                        fArray[fluid.index(i,j,k)] = 1.0f-fdist;
                                                }
                                        }
                                }
                        }
                }
        }
}
void 
simpleFluidEmitter::volumeFluidEmitter(
        MFnFluid&               fluid,
        const MMatrix&  fluidWorldMatrix,
        int                     plugIndex,
        MDataBlock&     block,
        double                  dt,
        double                  conversion,
        double                  dropoff
)
{
        
        
        MPoint emitterPos = getWorldPosition();
        MMatrix emitterWorldMatrix = getWorldMatrix();
        MMatrix fluidInverseWorldMatrix = fluidWorldMatrix.inverse();
        
        
        
        double densityEmit = fluidDensityEmission( block );
        double fuelEmit = fluidFuelEmission( block );
        double heatEmit = fluidHeatEmission( block );
        bool doEmitColor = fluidEmitColor( block );
        MColor emitColor = fluidColor( block );
        
        
        
        
        
        double theRate = getRate(block) * dt * conversion;
        
        
        
        double size[3];
        unsigned int res[3];
        fluid.getDimensions( size[0], size[1], size[2] );
        fluid.getResolution( res[0], res[1], res[2] );
        
        double dx = size[0] / res[0];
        double dy = size[1] / res[1];
        double dz = size[2] / res[2];
        
        
        double Ox = -size[0]/2;
        double Oy = -size[1]/2;
        double Oz = -size[2]/2; 
        
        
        
        MBoundingBox bbox;
        if( !volumePrimitiveBoundingBox( bbox ) )
        {
                
                
                return;
        }
        
        
        
        bbox.transformUsing( emitterWorldMatrix );
        bbox.transformUsing( fluidInverseWorldMatrix );
        MPoint lowCorner = bbox.min();
        MPoint highCorner = bbox.max();
        
        
        int3 lowCoords;
        int3 highCoords;
        fluid.toGridIndex( lowCorner, lowCoords );
        fluid.toGridIndex( highCorner, highCoords );
        
        int i;
        for ( i = 0; i < 3; i++ )
        {
                if ( lowCoords[i] < 0 ) {
                        lowCoords[i] = 0;
                } else if ( lowCoords[i] > ((int)res[i])-1 ) {
                        lowCoords[i] = ((int)res[i])-1;
                }
                if ( highCoords[i] < 0 ) {
                        highCoords[i] = 0;
                } else if ( highCoords[i] > ((int)res[i])-1 ) {
                        highCoords[i] = ((int)res[i])-1;
                }
                
        }
        
        
        
        
        
        double emitterVoxelSize[3];
        emitterVoxelSize[0] = (highCorner[0]-lowCorner[0])/dx;
        emitterVoxelSize[1] = (highCorner[1]-lowCorner[1])/dy;
        emitterVoxelSize[2] = (highCorner[2]-lowCorner[2])/dz;
                
        double minVoxelSize = MIN(emitterVoxelSize[0],MIN(emitterVoxelSize[1],emitterVoxelSize[2]));
        if( minVoxelSize < 1.0 )
        {
                minVoxelSize = 1.0;
        }
        int maxSamples = 8;
        int numSamples = (int)(8.0/(minVoxelSize*minVoxelSize*minVoxelSize) + 0.5);
        if( numSamples < 1 ) numSamples = 1;
        if( numSamples > maxSamples ) numSamples = maxSamples;
        
        
        
        
        bool jitter = fluidJitter(block);
        if( !jitter )
        {
                numSamples = 1;
        }
        
        
        
        
        
        
        for( i = lowCoords[0]; i <= highCoords[0]; i++ )
        {
                double x = Ox + (i+0.5)*dx;
                        
                for( int j = lowCoords[1]; j < highCoords[1]; j++ )
                {
                        double y = Oy + (j+0.5)*dy;
                        for( int k = lowCoords[2]; k < highCoords[2]; k++ )
                        {
                                double z = Oz + (k+0.5)*dz;
                                
                                for ( int si = 0; si < numSamples; si++) {
                                        
                                        
                                        
                                        double rx, ry, rz;
                                        if(jitter) {
                                                rx = x + dx*(randgen() - 0.5);
                                                ry = y + dy*(randgen() - 0.5);
                                                rz = z + dz*(randgen() - 0.5);
                                        } else {
                                                rx = x;
                                                ry = y;
                                                rz = z;
                                        }
                                        
                                        
                                        MPoint pt( rx, ry, rz );
                                        pt *= fluidWorldMatrix;
                                        
                                        
                                        if( volumePrimitivePointInside( pt, emitterWorldMatrix ) )
                                        {
                                                
                                                
                                                double dist = pt.distanceTo( emitterPos );
                                                double distDrop = dropoff * (dist*dist);
                                                double newVal = (theRate * exp( -distDrop )) / (double)numSamples;
                                                
                                                
                                                
                                                if( newVal != 0.0 )
                                                {
                                                        fluid.emitIntoArrays( (float) newVal, i, j, k, (float)densityEmit, (float)heatEmit, (float)fuelEmit, doEmitColor, emitColor );
                                                }
                                        }
                                }
                        }
                }
        }
}
void 
simpleFluidEmitter::surfaceFluidEmitter(
        MFnFluid&               fluid,
        const MMatrix&  fluidWorldMatrix,
        int                     plugIndex,
        MDataBlock&     block,
        double                  dt,
        double                  conversion,
        double                  dropoff
)
{
        
        
        MMatrix fluidInverseWorldMatrix = fluidWorldMatrix.inverse();
        
        
        
        double densityEmit = fluidDensityEmission( block );
        double fuelEmit = fluidFuelEmission( block );
        double heatEmit = fluidHeatEmission( block );
        bool doEmitColor = fluidEmitColor( block );
        MColor emitColor = fluidColor( block );
        
        
        
        
        
        double theRate = getRate(block) * dt * conversion;
        
        
        double size[3];
        unsigned int res[3];
        fluid.getDimensions( size[0], size[1], size[2] );
        fluid.getResolution( res[0], res[1], res[2] );
                
        
        double dx = size[0] / res[0];
        double dy = size[1] / res[1];
        double dz = size[2] / res[2];
        
        
        double Ox = -size[0]/2;
        double Oy = -size[1]/2;
        double Oz = -size[2]/2; 
        
        
        
        
        
        MDataHandle sweptHandle = block.inputValue( mSweptGeometry );
        MObject sweptData = sweptHandle.data();
        MFnDynSweptGeometryData fnSweptData( sweptData );
        
        
        
        
        
        bool jitter = fluidJitter(block);
        if( !jitter )
        {
                resetRandomState( plugIndex, block );
        }
        if( fnSweptData.triangleCount() > 0 )
        {
                
                
                
                
                double vfArea = pow(dx*dy*dz, 2.0/3.0);
                
                
                
                
                
                
                MFnDependencyNode fnNode( thisMObject() );
                MObject rateTextureAttr = fnNode.attribute( "textureRate" );
                MObject colorTextureAttr = fnNode.attribute( "particleColor" );
                bool texturedRate = hasValidEmission2dTexture( rateTextureAttr );
                bool texturedColor = hasValidEmission2dTexture( colorTextureAttr );
                
                
                
                MDoubleArray uCoords, vCoords;
                if( texturedRate || texturedColor )
                {
                        uCoords.setLength( fnSweptData.triangleCount() );
                        vCoords.setLength( fnSweptData.triangleCount() );
                        
                        int t;
                        for( t = 0; t < fnSweptData.triangleCount(); t++ )
                        {
                                MDynSweptTriangle tri = fnSweptData.sweptTriangle( t );
                                MVector uv0 = tri.uvPoint(0);
                                MVector uv1 = tri.uvPoint(1);
                                MVector uv2 = tri.uvPoint(2);
                                
                                MVector uvMid = (uv0+uv1+uv2)/3.0;
                                uCoords[t] = uvMid[0];
                                vCoords[t] = uvMid[1];
                        }
                }
                
                
                MDoubleArray texturedRateValues;
                if( texturedRate )
                {
                        texturedRateValues.setLength( uCoords.length() );
                        evalEmission2dTexture( rateTextureAttr, uCoords, vCoords, NULL, &texturedRateValues );
                }
                
                MVectorArray texturedColorValues;
                if( texturedColor )
                {
                        texturedColorValues.setLength( uCoords.length() );
                        evalEmission2dTexture( colorTextureAttr, uCoords, vCoords, &texturedColorValues, NULL );
                }
                
                for( int t = 0; t < fnSweptData.triangleCount(); t++ )
                {
                        
                        
                        double curTexturedRate = texturedRate ? texturedRateValues[t] : 1.0;
                        MColor curTexturedColor;
                        if( texturedColor )
                        {
                                MVector& curVec = texturedColorValues[t];
                                curTexturedColor.r = (float)curVec[0];
                                curTexturedColor.g = (float)curVec[1];
                                curTexturedColor.b = (float)curVec[2];
                                curTexturedColor.a = 1.0;
                        }
                        else
                        {
                                curTexturedColor = emitColor;
                        }
                        MDynSweptTriangle tri = fnSweptData.sweptTriangle( t );
                        MVector v0 = tri.vertex(0);
                        MVector v1 = tri.vertex(1);
                        MVector v2 = tri.vertex(2);
                        
                        
                        
                        
                        double triArea = tri.area();
                        int numSamples = (int)(triArea / vfArea);
                        if( numSamples < 1 ) numSamples = 1;
                        
                        
                        
                        
                        
                        double triRate = (theRate*(triArea/vfArea))/numSamples;
                        
                        triRate *= curTexturedRate;
                        
                        for( int j = 0; j < numSamples; j++ )
                        {
                                
                                
                                
                                double r1 = randgen();
                                double r2 = randgen();
                                
                                if( r1 + r2 > 1 )
                                {
                                        r1 = 1-r1;
                                        r2 = 1-r2;
                                }
                                double r3 = 1 - (r1+r2);
                                MPoint randPoint = r1*v0 + r2*v1 + r3*v2;
                                randPoint *= fluidInverseWorldMatrix;
                                
                                
                                
                                int3 coord;
                                fluid.toGridIndex( randPoint, coord );
                                
                                if( (coord[0]<0) || (coord[1]<0) || (coord[2]<0) ||
                                        (coord[0]>=(int)res[0]) || (coord[1]>=(int)res[1]) || (coord[2]>=(int)res[2]) )
                                {
                                        continue;
                                }
                                
                                
                                
                                
                                MPoint gridPoint;
                                gridPoint.x = Ox + (coord[0]+0.5)*dx;
                                gridPoint.y = Oy + (coord[1]+0.5)*dy;
                                gridPoint.z = Oz + (coord[2]+0.5)*dz;
                                
                                MVector diff = gridPoint - randPoint;
                                double distSquared = diff * diff;
                                double distDrop = dropoff * distSquared;
                                
                                double newVal = triRate * exp( -distDrop );
                
                                
                                
                                if( newVal != 0 )
                                {
                                        fluid.emitIntoArrays( (float) newVal, coord[0], coord[1], coord[2], (float)densityEmit, (float)heatEmit, (float)fuelEmit, doEmitColor, curTexturedColor );              
                                }
                        }
                }
        }
}
MStatus initializePlugin(MObject obj)
{
        MStatus status;
        MFnPlugin plugin(obj, PLUGIN_COMPANY, "3.0", "Any");
        status = plugin.registerNode( "simpleFluidEmitter", simpleFluidEmitter::id,
                                                        &simpleFluidEmitter::creator, &simpleFluidEmitter::initialize,
                                                        MPxNode::kFluidEmitterNode );
        if (!status) {
                status.perror("registerNode");
                return status;
        }
        return status;
}
MStatus uninitializePlugin(MObject obj)
{
        MStatus status;
        MFnPlugin plugin(obj);
        status = plugin.deregisterNode( simpleFluidEmitter::id );
        if (!status) {
                status.perror("deregisterNode");
                return status;
        }
        return status;
}