//
// Copyright (C) Single Scatter Fog
// 
// File: SingleScatterFogNode.cpp
//
// Dependency Graph Node: SingleScatterFog
//
// Author: Maya Plug-in Wizard 2.0
//

#include "SingleScatterFogNode.h"

#include <maya/MPlug.h>
#include <maya/MDataBlock.h>
#include <maya/MDataHandle.h>
#include <maya/MFloatVector.h>
#include <maya/MItDag.h>
#include <maya/MFnPointLight.h>
#include <maya/MDagPath.h>
#include <maya/MMatrix.h>
#include <maya/MColor.h>
#include <maya/MGlobal.h>

#include <math.h>

#define E 2.718281828459f
#define PI 3.1415926535898f

// You MUST change this to a unique value!!!  The id is a 32bit value used
// to identify this type of node in the binary file format.  
//
MTypeId SingleScatterFog::id( 0x00001 );

MObject SingleScatterFog::aColor, SingleScatterFog::aScatteringCoef,
	SingleScatterFog::aFarPointW, SingleScatterFog::aPointW;
MObject SingleScatterFog::aOutColor, SingleScatterFog::aOutTransparency,
	SingleScatterFog::aOutAlpha, SingleScatterFog::aOutScatteringCoef;

TableReader SingleScatterFog::F;

SingleScatterFog::SingleScatterFog() {}
SingleScatterFog::~SingleScatterFog() {}

MStatus SingleScatterFog::compute( const MPlug& plug, MDataBlock& data )
{
	if( plug == aOutColor
		|| plug == aOutTransparency
		|| plug == aOutAlpha )
	{
		MFloatVector &Color = data.inputValue( aColor ).asFloatVector(),
			&pointWorld = data.inputValue( aPointW ).asFloatVector(),
			&farPointWorld = data.inputValue( aFarPointW ).asFloatVector();

		float &scatteringCoef = data.inputValue( aScatteringCoef ).asFloat();

		MFloatVector resultColor( 0.0, 0.0, 0.0 );

		MFloatVector vecVP = pointWorld - farPointWorld;
		float Tvp = scatteringCoef * vecVP.length();

		MStatus stat;
		MItDag dagNodes( MItDag::kBreadthFirst, MFn::kPointLight, &stat );

		if(!stat)
		{
			cerr << "dagNodes error\n";
			return stat;
		}

		for(;!dagNodes.isDone();dagNodes.next())
		{
			MObject lightNode = dagNodes.item();
			MFnPointLight pointLight( lightNode, &stat );

			if( !stat )
			{
				cerr << "error getting point light\n";
				return stat;
			}

			MDagPath lightPath;
			CHECK_MSTATUS_AND_RETURN_IT( pointLight.getPath( lightPath ) );

			MMatrix lightTransform = lightPath.inclusiveMatrix( &stat );

			if( !stat )
			{
				cerr << "error getting inclusive matrix\n";
				return stat;
			}

			MFloatVector lightPos( (float)lightTransform( 3, 0 ),
								   (float)lightTransform( 3, 1 ),
								   (float)lightTransform( 3, 2 ) );

			MFloatVector vecSV = pointWorld - lightPos;

			float gamma = vecSV.angle( vecVP );

			float Dsv = vecSV.length(),
				Tsv = scatteringCoef * vecSV.length();
			
			float cos_gamma = cos(gamma),
				sin_gamma = sin(gamma);
			
			float A0 = (scatteringCoef * pointLight.intensity() * pow( E, -Tsv * cos_gamma ))
				/ (2.0f * PI * Dsv * sin_gamma),
				A1 = Tsv * sin_gamma,
				b = PI/4.0f + (1.0f/2.0f) * atan( (Tvp - Tsv * cos_gamma) / (Tsv * sin_gamma) );

			float f1 = F.LookUp( b, A1 ),
				f2 = F.LookUp( gamma/2.0f, A1 );

			float airlight = A0 * (f1 - f2);

			MColor finalColor = airlight * pointLight.color();
			
			resultColor.x += finalColor.r;
			resultColor.y += finalColor.g;
			resultColor.z += finalColor.b;

//			cerr << A0 << " " << A1 << " " << b << " " << f1 << " " << f2 << " " << airlight << endl;
		}

		float fog_power = min( pow( E, -Tvp ), 1.0f );

		resultColor.x *= Color.x;
		resultColor.y *= Color.y;
		resultColor.z *= Color.z;
		
		resultColor += Color * (1 - fog_power);

		MDataHandle outColorHandle = data.outputValue( aOutColor );
		MFloatVector& outColor = outColorHandle.asFloatVector();
		outColor = resultColor;
		outColorHandle.setClean();

		MDataHandle outTransparencyHandle = data.outputValue( aOutTransparency );
		MFloatVector& outTransparency = outTransparencyHandle.asFloatVector();
		outTransparency.x = outTransparency.y = outTransparency.z = fog_power;
		
		data.setClean(plug);
	}
	else if( plug == aOutScatteringCoef )
	{
		float &scatteringCoef = data.inputValue( aScatteringCoef ).asFloat();
		float &outScatteringCoef = data.outputValue( aOutScatteringCoef ).asFloat();

		outScatteringCoef = scatteringCoef;
		
		data.setClean(plug);
	}
	else
	{
		return MS::kUnknownParameter;
	}

	return MS::kSuccess;
}

void* SingleScatterFog::creator()
{
	return new SingleScatterFog();
}

MStatus SingleScatterFog::initialize()
{
	MFnNumericAttribute nAttr;

	// Inputs
	//
	aColor = nAttr.createColor( "color", "c" );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setKeyable(true) );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setStorable(true) );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setHidden(false) );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setDefault(1.0f, 1.0f, 1.0f) );

	aScatteringCoef = nAttr.create( "scatteringCoef", "scac", MFnNumericData::kFloat );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setKeyable(true) );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setStorable(true) );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setHidden(false) );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setDefault( 0.05f ) );

	aPointW = nAttr.createPoint("pointWorld", "pw" );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setStorable(false) );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setHidden(true) );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setReadable(true) );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setWritable(true) );

	aFarPointW = nAttr.createPoint("farPointWorld", "fw" );
    CHECK_MSTATUS_AND_RETURN_IT(nAttr.setStorable(false)  );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setHidden(true) );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setReadable(true) );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setWritable(true) );


	// Outputs
	//
	aOutColor = nAttr.createColor( "outColor", "oc" );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setStorable(false) );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setHidden(false) );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setReadable(true) );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setWritable(false) );

    aOutTransparency = nAttr.createColor( "outTransparency", "ot" );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setStorable(false) );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setHidden(false) );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setReadable(true) );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setWritable(false) );
	CHECK_MSTATUS_AND_RETURN_IT( nAttr.setDefault( 1.0, 1.0, 1.0 ) );

    aOutAlpha = nAttr.create( "outAlpha", "oa", MFnNumericData::kFloat );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setStorable(false) );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setHidden(false) );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setReadable(true) );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setWritable(false) );
	CHECK_MSTATUS_AND_RETURN_IT( nAttr.setDefault( 1.0 ) );

	aOutScatteringCoef = nAttr.create( "outScatteringCoef", "osc", MFnNumericData::kFloat );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setStorable(false) );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setHidden(false) );
	CHECK_MSTATUS_AND_RETURN_IT( nAttr.setConnectable(true) );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setReadable(true) );
    CHECK_MSTATUS_AND_RETURN_IT( nAttr.setWritable(false) );
	CHECK_MSTATUS_AND_RETURN_IT( nAttr.setDefault( 1.0 ) );


	// Add the attributes
	//
	CHECK_MSTATUS_AND_RETURN_IT( addAttribute(aColor) );
	CHECK_MSTATUS_AND_RETURN_IT( addAttribute(aScatteringCoef) );
	CHECK_MSTATUS_AND_RETURN_IT( addAttribute(aPointW) );
	CHECK_MSTATUS_AND_RETURN_IT( addAttribute(aFarPointW) );
	
	CHECK_MSTATUS_AND_RETURN_IT( addAttribute(aOutColor) );
    CHECK_MSTATUS_AND_RETURN_IT( addAttribute(aOutTransparency) );
    CHECK_MSTATUS_AND_RETURN_IT( addAttribute(aOutAlpha) );
	CHECK_MSTATUS_AND_RETURN_IT( addAttribute(aOutScatteringCoef) );

	// Set dependencies
	//
	CHECK_MSTATUS_AND_RETURN_IT( attributeAffects( aColor, aOutColor ) );
	CHECK_MSTATUS_AND_RETURN_IT( attributeAffects( aScatteringCoef, aOutColor ) );
	CHECK_MSTATUS_AND_RETURN_IT( attributeAffects( aPointW, aOutColor ) );
	CHECK_MSTATUS_AND_RETURN_IT( attributeAffects( aFarPointW, aOutColor ) );

	CHECK_MSTATUS_AND_RETURN_IT( attributeAffects( aScatteringCoef, aOutTransparency ) );
	CHECK_MSTATUS_AND_RETURN_IT( attributeAffects( aPointW, aOutTransparency ) );
	CHECK_MSTATUS_AND_RETURN_IT( attributeAffects( aFarPointW, aOutTransparency ) );

	CHECK_MSTATUS_AND_RETURN_IT( attributeAffects( aScatteringCoef, aOutScatteringCoef ) );

	CHECK_MSTATUS_AND_RETURN_IT( F.Load( "F_512.csv" ) );

	return MS::kSuccess;
}

