//
// Copyright (C) 
// Author: Trevor Sommer

#include <AmbientOcclusionShader.h>

//Node ID
//Warning this ID is just a random number, I cannot garuntee it will be unique
MTypeId     AmbientOcclusionShader::id(0x0066670);

//Node Attributes
MObject  AmbientOcclusionShader::aOutColor;
MObject  AmbientOcclusionShader::aNormalCamera;
MObject  AmbientOcclusionShader::aPointCamera;
MObject  AmbientOcclusionShader::aRaySampler;
MObject  AmbientOcclusionShader::aRayCount;


AmbientOcclusionShader::AmbientOcclusionShader() 
{
}

AmbientOcclusionShader::~AmbientOcclusionShader() 
{
}

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

//  @brief: Node Init
//
//  @return [MStatus]
MStatus AmbientOcclusionShader::initialize()
{
	MStatus status;

    MFnNumericAttribute nAttr;

	//inputs
    aRayCount = nAttr.create("rayCount", "rc", MFnNumericData::kInt);
    status = nAttr.setDefault(25);
    CHECK_MSTATUS_AND_RETURN_IT(status);
    status = nAttr.setMin(1);
    CHECK_MSTATUS_AND_RETURN_IT(status);
    status = nAttr.setSoftMax(500);
    CHECK_MSTATUS_AND_RETURN_IT(status);
    status = addAttribute(aRayCount);
    CHECK_MSTATUS_AND_RETURN_IT(status);

    aNormalCamera = nAttr.createPoint("normalCamera", "n");
    status = nAttr.setHidden(true);
    CHECK_MSTATUS_AND_RETURN_IT(status);
    status = addAttribute(aNormalCamera);
    CHECK_MSTATUS_AND_RETURN_IT(status);

    aPointCamera =  nAttr.createPoint("pointCamera", "p");
    status = nAttr.setHidden(true);
    CHECK_MSTATUS_AND_RETURN_IT(status);
    status = addAttribute(aPointCamera);
    CHECK_MSTATUS_AND_RETURN_IT(status);

    aRaySampler =  nAttr.createAddr("raySampler", "rtr");
    status = nAttr.setHidden(true);
    CHECK_MSTATUS_AND_RETURN_IT(status);
    status = addAttribute(aRaySampler);
    CHECK_MSTATUS_AND_RETURN_IT(status);

	//outputs
    aOutColor = nAttr.createColor("outColor", "oc");
    status = nAttr.setWritable(false);
    CHECK_MSTATUS_AND_RETURN_IT(status);
    status = nAttr.setHidden(true);
    CHECK_MSTATUS_AND_RETURN_IT(status);
    status = addAttribute(aOutColor);
    CHECK_MSTATUS_AND_RETURN_IT(status);

	//attribute Affects
    status = attributeAffects(aNormalCamera, aOutColor);
    CHECK_MSTATUS_AND_RETURN_IT(status);
    status = attributeAffects(aPointCamera, aOutColor);
    CHECK_MSTATUS_AND_RETURN_IT(status);

	return MS::kSuccess;
}

//  @brief: Calculates color value between black and white.
//  Color is based on the number of hits from random raycast vectors from the render pixel.
//
//  @return [MStatus]
MStatus AmbientOcclusionShader::compute(const MPlug& plug, MDataBlock& data)
{
	MStatus status;
    
    MFloatVector outColor;
    MFloatVector normCam;
    MFloatVector pointCam;

    //NOTE: get camera info
    MDataHandle camNormDHL = data.inputValue(aNormalCamera, &status);
    CHECK_MSTATUS_AND_RETURN_IT(status);
    normCam = camNormDHL.asFloatVector();

    MDataHandle camPointDHL = data.inputValue(aPointCamera, &status);
    CHECK_MSTATUS_AND_RETURN_IT(status);
    pointCam = camPointDHL.asFloatVector();

    MDataHandle raySamplerDHL = data.inputValue(aRaySampler, &status);
    CHECK_MSTATUS_AND_RETURN_IT(status);
    void* raySampler = raySamplerDHL.asAddr();

    int raycount;
    MDataHandle rayCountDHL = data.inputValue(aRayCount, &status);
    CHECK_MSTATUS_AND_RETURN_IT(status);
    raycount = rayCountDHL.asInt();

    MDataHandle outColorDHL= data.outputValue(aOutColor, &status);
    CHECK_MSTATUS_AND_RETURN_IT(status);

    int numbHits = 0;
    double hitFraction;
    double hitColor;
    
    MFloatVectorArray origins(raycount);
    MFloatVectorArray directions(raycount);
    MFloatVectorArray intersections(raycount);
    MIntArray didHits(raycount);

    for(int idx=0; idx < raycount; idx++)
    {
        origins[idx] = pointCam;

        MFloatVector randDirect = getRayBounceVector(idx);

            if(randDirect * normCam < 0.0)
            {
                randDirect = -1.0 * randDirect;
            }

        directions[idx] = randDirect;
        
    }

    status = MRenderUtil::raytraceFirstGeometryIntersections(origins,
                                                            directions,
                                                            NULL,
                                                            raySampler,
                                                            intersections,
                                                            didHits);
    CHECK_MSTATUS_AND_RETURN_IT(status);

    for(int idx=0; idx < raycount; idx++)
    {
        if(didHits[idx])
        {
            numbHits++;
        }
        
    }

    hitFraction = (double)numbHits/raycount;
    hitColor = 1.0 - hitFraction;
    outColor = MFloatVector(hitColor, hitColor, hitColor);

    outColorDHL.set(outColor);
    
    return MS::kSuccess;

}

//  @brief: Uses the given random seed value to return a randomly generated vector 
//
//  @return [MFloatVector]
MFloatVector AmbientOcclusionShader::getRayBounceVector(unsigned int seed)
{
    double max = double(RAND_MAX);
	srand(seed);
	
	double x = ((double)(rand() % RAND_MAX) / max) + 2.0 - 1.0;
	double y = ((double)(rand() % RAND_MAX) / max) + 2.0 - 1.0;
	double z = ((double)(rand() % RAND_MAX) / max) + 2.0 - 1.0;

	return MFloatVector(x, y, z);
}