// Shave and a Haircut
// (c) 2019 Epic Games
// US Patent 6720962

#include "ShaveAndHaircut.h"

#include <maya/MPlugArray.h>
#include <maya/MPlug.h>

#include <ai_nodes.h>

#include "utils/time.h"
#include "utils/Version.h"

class CExtensionAttrHelper;

MStatus CShaveTranslator::UpdateHairInfo()
{
   MStatus status;
   // shaveItHair init only reads MObjectArray, so we append our node into an array
   //
   MObjectArray hairNodes(1, m_dagPath.node());

   // Export shaveAndHaircut info into a variable
   //
   status = shaveAPI::exportHair(hairNodes, &m_hairInfo);
   return status;
}

AtNode*  CShaveTranslator::CreateArnoldNodes()
{
   return AddArnoldNode("curves");
}


AtNode* CShaveTranslator::CreateShaveShader(AtNode* curve)
{   
   char nodeName[MAX_NAME_SIZE];
   // check if a shaderHair node hasn't been already created in a previous export
#if (MTOA_ARCH_VERSION_NUM < 2)
   AtNode* shader = AiNode("ShaveHair");
#else
   AtNode* shader = GetArnoldNode("ShaveHair");
#endif
   if (shader == NULL)
   {
#if (MTOA_ARCH_VERSION_NUM >= 2)
      shader = AddArnoldNode("ShaveHair", "ShaveHair");
#endif
      AiNodeSetStr(shader, "name", NodeUniqueName(shader, nodeName));
   }

   // Fade the hairstrand towards the tip.
   MPlug plug = FindMayaPlug("tipFade");
   if (plug.asBool())
   {
      // If tip fade is on then the hairs are not opaque no
      // matter the attribute setting.
      AiNodeSetBool(curve, "opaque", false);

#if (MTOA_ARCH_VERSION_NUM < 2)
      AtNode* ramp = AiNode("MayaRamp");
#else
      AtNode* ramp = GetArnoldNode("MayaRamp");
#endif
      if (ramp == NULL)
      {
#if (MTOA_ARCH_VERSION_NUM >= 2)
         ramp = AddArnoldNode("MayaRamp", "MayaRamp");
#endif
         AiNodeSetStr(ramp, "name", NodeUniqueName(ramp, nodeName));
      }
#if (MTOA_ARCH_VERSION_NUM < 2)
      AtNode* placementNode = AiNode("MayaPlace2DTexture");
#else
      AtNode* placementNode = GetArnoldNode("MayaPlace2DTexture");
#endif
      if (placementNode == NULL)
      {
#if (MTOA_ARCH_VERSION_NUM >= 2)
         placementNode = AddArnoldNode("MayaPlace2DTexture", "MayaPlace2DTexture");
#endif
         AiNodeSetStr(placementNode, "name", NodeUniqueName(placementNode, nodeName));
      }
      AiNodeSetStr(ramp, "type", "v");

      AtArray* positions  = AiArrayAllocate(2, 1, AI_TYPE_FLOAT);
      AtArray* colors  = AiArrayAllocate(2, 1, AI_TYPE_RGB);
      
      AiArraySetFlt(positions, 0, 0.55f);
      AiArraySetFlt(positions, 1, 1.0f);
      
      AtRGB color;
      
      color.r = 1.0f; color.g = 1.0f; color.b = 1.0f;
      AiArraySetRGB(colors, 0, color);
      
      color.r = 0.0f; color.g = 0.0f; color.b = 0.0f;
      AiArraySetRGB(colors, 1, color);
      
      AiNodeSetArray(ramp, "position", positions);
      AiNodeSetArray(ramp, "color", colors);
   
      AiNodeLink (placementNode, "uvCoord", ramp);
      AiNodeLink (ramp, "strand_opacity", shader);
   }

#if (MTOA_ARCH_VERSION_NUM < 2)
   // Add shader uparam and vparam names.
   AiNodeSetStr(shader, "uparam",   "uparamcoord");
   AiNodeSetStr(shader, "vparam",   "vparamcoord");
#endif

   // Add root and tip color.
   AiNodeSetStr(shader, "rootcolor","rootcolorparam");
   AiNodeSetStr(shader, "tipcolor", "tipcolorparam");

   // Set specular and gloss.
   ProcessParameter(shader, "spec", AI_TYPE_FLOAT, "specular");
   ProcessParameter(shader, "gloss", AI_TYPE_FLOAT, "gloss");
   ProcessParameter(shader, "spec_color", AI_TYPE_RGB, "specularTint");
   ProcessParameter(shader, "ambdiff", AI_TYPE_FLOAT, "amb/diff");
   ProcessParameter(shader, "kd_ind", AI_TYPE_FLOAT, "aiIndirect");
#if (MTOA_ARCH_VERSION_NUM < 2)
   ProcessParameter(shader, "diffuse_cache", AI_TYPE_BOOLEAN, "aiDiffuseCache");   ProcessParameter(shader, "direct_diffuse", AI_TYPE_FLOAT, "aiDirectDiffuse");
#endif
   ProcessParameter(shader, "aov_direct_diffuse", AI_TYPE_STRING, "aiAovDirectDiffuse");
   ProcessParameter(shader, "aov_indirect_diffuse", AI_TYPE_STRING, "aiAovIndirectDiffuse");
   ProcessParameter(shader, "aov_direct_specular", AI_TYPE_STRING, "aiAovDirectSpecular");

   return shader;
}

void CShaveTranslator::Export(AtNode* curve)
{
   // Only translate the shave node if its marked as a active
   if (!FindMayaPlug("active").asBool())
      return;

   Update(curve);
}


void CShaveTranslator::Update(AtNode* curve)
{
   // Export shaveAndHaircut info into a variable
   if (UpdateHairInfo() != MS::kSuccess)
      return;

   // The shader nodes
   // TODO: Kill these and export it properly.
   AtNode* shader       = NULL;

   // Export the transform matrix
#if (MTOA_ARCH_VERSION_NUM < 2)
   ExportMatrix(curve, 0);
#else
   ExportMatrix(curve);
#endif

   // Get the visibiliy and render flags set.
   ProcessRenderFlags(curve);

   // Curves shader
   MPlug plug;
   plug = FindMayaPlug("aiOverrideHair");
   if (!plug.isNull() && plug.asBool())
   {
      MPlugArray curveShaderPlug;
      plug = FindMayaPlug("aiHairShader");
      if (!plug.isNull())
      {
         plug.connectedTo(curveShaderPlug, true, false);
         if (curveShaderPlug.length() > 0)
         {
#if (MTOA_ARCH_VERSION_NUM < 2)
            shader = ExportRootShader(curveShaderPlug[0]);
#else
            shader = ExportConnectedNode(curveShaderPlug[0]);
#endif
         }
      }
   }

   // Default to the ShaveHair shader if nothing else has been set.
   if (shader == NULL)
   {
#if (MTOA_ARCH_VERSION_NUM < 2)
      shader = ExportRootShader(CreateShaveShader(curve));
#else
      shader = CreateShaveShader(curve);
#endif
   }

#if (MTOA_ARCH_VERSION_NUM < 2)
   if (shader != NULL)
   {
      AiNodeSetPtr(curve, "shader", shader);
   }
#else
   SetRootShader(shader);
#endif
   
   // Should we export the hair root and tip colour? Default to true.
   // Turning it off gives us a slimmer ass.
   plug = FindMayaPlug("aiExportHairColors");
   bool export_curve_color = true;
   if (!plug.isNull())
   {
      export_curve_color = plug.asBool();
   }
   
   // The numPoints array (int array the size of numLines, no motionsteps)
   AtArray* curveNumPoints             = AiArrayAllocate(m_hairInfo.numHairs, 1, AI_TYPE_INT);

   // The root and tip color array
   AtArray* rootColor = NULL;
   AtArray* tipColor = NULL;
   if(export_curve_color)
   {
      rootColor = AiArrayAllocate(m_hairInfo.numHairs, 1, AI_TYPE_RGB);
      tipColor  = AiArrayAllocate(m_hairInfo.numHairs, 1, AI_TYPE_RGB);
   }

   plug = FindMayaPlug("aiExportHairIDs");
   bool export_curve_id = true;
   if (!plug.isNull())
   {
      export_curve_id = plug.asBool();
   }

   AtArray * curveID = NULL;
   if (export_curve_id)
   {
      curveID = AiArrayAllocate(m_hairInfo.numHairs, 1, AI_TYPE_UINT);
   }
   

   // The U and V paramcoords array
#if (MTOA_ARCH_VERSION_NUM < 2)
   AtArray* curveUParamCoord          = AiArrayAllocate(m_hairInfo.numHairs, 1, AI_TYPE_FLOAT);
   AtArray* curveVParamCoord          = AiArrayAllocate(m_hairInfo.numHairs, 1, AI_TYPE_FLOAT);
#else
   AtArray* curveParamCoord          = AiArrayAllocate(m_hairInfo.numHairs, 1, AI_TYPE_VECTOR2);
#endif
   

   for (int i = 0; i < m_hairInfo.numHairs; ++i)
   {
      int hairRootIndex = m_hairInfo.hairStartIndices[i];

      // Set numPoints
      int numRenderLineCVs = m_hairInfo.hairEndIndices[i] - hairRootIndex;
      AiArraySetInt(curveNumPoints, i, (numRenderLineCVs+2));

      // Set UV
#if (MTOA_ARCH_VERSION_NUM < 2)
      AiArraySetFlt(curveUParamCoord, i, m_hairInfo.uvws[hairRootIndex].x);
      AiArraySetFlt(curveVParamCoord, i, m_hairInfo.uvws[hairRootIndex].y);
#else
      AtVector2 hairUVs(m_hairInfo.uvws[hairRootIndex].x, m_hairInfo.uvws[hairRootIndex].y);
      AiArraySetVec2(curveParamCoord, i, hairUVs);
#endif
      
      // Root and tip colours for the ShaveHair shader.
      // TODO: Make exporting all the info for the ShaveHair shader an option.
      if (export_curve_color)
      {
         AtRGB shaveRootColors;
         shaveRootColors.r = m_hairInfo.rootColors[i].r;
         shaveRootColors.g = m_hairInfo.rootColors[i].g;
         shaveRootColors.b = m_hairInfo.rootColors[i].b;
   
         AtRGB shaveTipColors;
         shaveTipColors.r = m_hairInfo.tipColors[i].r;
         shaveTipColors.g = m_hairInfo.tipColors[i].g;
         shaveTipColors.b = m_hairInfo.tipColors[i].b;

         AiArraySetRGB(rootColor, i, shaveRootColors);
         AiArraySetRGB(tipColor,  i, shaveTipColors);
      }

      if (export_curve_id)
      {
         AiArraySetUInt(curveID, i, (i));
      }
   }

   // Allocate memory for all curve points and widths (aka radii)

   // The required size of the radius array varies by curve basis.  The equation is:
   //
   //  (np - 4) / vstep + 2
   //
   // where np is the number of points in the strand, and vstep is the number of points per segment
   // (which varies depending on which curve basis you choose).
   //
   // table of vstep values:
   //   3 - bezier
   //   1 - b-spline
   //   1 - catmull-rom
   //   2 - hermite
   //   4 - power
   // for shave we use catmull-rom, which means np -2.
   // It seems (though I have yet to confirm with SA) that the discrepancy is due to duplicate knots at the
   // beginning and end of each strand in arnold curves, which are not considered in the "radius" array.
   // Shave does not include these extra knots in its HairInfo.vertices array.
   AtArray* curveWidths = AiArrayAllocate(m_hairInfo.numHairVertices, 1, AI_TYPE_FLOAT);
   int numPointsInterpolation = m_hairInfo.numHairVertices + (m_hairInfo.numHairs * 2);

#if (MTOA_ARCH_VERSION_NUM < 2)
   const int kPointType = AI_TYPE_POINT;
#else
   const int kPointType = AI_TYPE_VECTOR;
#endif
   AtArray* curvePoints = NULL;
   // TODO: Change this to use RequiresMotionDeformData()
   if (RequiresMotionData() && IsMotionBlurEnabled(MTOA_MBLUR_DEFORM))
   {
      curvePoints = AiArrayAllocate(numPointsInterpolation, GetNumMotionSteps(), kPointType);
   }
   else
   {
      curvePoints = AiArrayAllocate(numPointsInterpolation, 1, kPointType);
   }

   // Set the required arrays
   AiNodeSetArray(curve, "num_points",             curveNumPoints);
   AiNodeSetArray(curve, "points",                 curvePoints);
   AiNodeSetArray(curve, "radius",                 curveWidths);

#if (MTOA_ARCH_VERSION_NUM < 2)
   AiNodeDeclare(curve, "uparamcoord",             "uniform FLOAT");
   AiNodeDeclare(curve, "vparamcoord",             "uniform FLOAT");
   AiNodeSetArray(curve, "uparamcoord",            curveUParamCoord);
   AiNodeSetArray(curve, "vparamcoord",            curveVParamCoord);
#else
   AiNodeSetArray(curve, "uvs",            curveParamCoord);
#endif
   
   // Finally add and set the hair color arrays as needed.
   if(export_curve_color)
   {
      AiNodeDeclare(curve, "rootcolorparam",       "uniform RGB");
      AiNodeDeclare(curve, "tipcolorparam",        "uniform RGB");
      AiNodeSetArray(curve, "rootcolorparam",      rootColor);
      AiNodeSetArray(curve, "tipcolorparam",       tipColor);
   }

   if(export_curve_id)
   {
      AiNodeDeclare(curve, "curve_id",        "uniform UINT");
      AiNodeSetArray(curve, "curve_id",       curveID);
   }
   
   // Set tesselation method
   AiNodeSetStr(curve, "basis", "catmull-rom");

   // Hair specific Arnold render settings.
   plug = FindMayaPlug("aiMinPixelWidth");
   if (!plug.isNull())
   {
      AiNodeSetFlt(curve, "min_pixel_width", plug.asFloat());
   }

   // Mode is an enum, 0 == ribbon, 1 == tubes.
   plug = FindMayaPlug("aiMode");
   if (!plug.isNull())
   {
      AiNodeSetInt(curve, "mode", plug.asInt());
   }

   ProcessHairLines(0,
                    curvePoints,
                    curveNumPoints,
                    curveWidths);

   m_hairInfo.clear();
}

#if (MTOA_ARCH_VERSION_NUM < 2)
void CShaveTranslator::ExportMotion(AtNode* curve, unsigned int step)
#else
void CShaveTranslator::ExportMotion(AtNode* curve)
#endif
{
   // Check if motionblur is enabled and early out if it's not.
   if (!IsMotionBlurEnabled()) return;

   // Set transform matrix
#if (MTOA_ARCH_VERSION_NUM < 2)
   ExportMatrix(curve, step);
#else
   ExportMatrix(curve);
   int step = GetMotionStep();
#endif

   if (IsMotionBlurEnabled(MTOA_MBLUR_DEFORM))
   {
      // Bail early if we've trouble getting data from Shave.
      //
      if (UpdateHairInfo() != MS::kSuccess) return;

      ProcessHairLines(step,
                       AiNodeGetArray(curve, "points"),
                       AiNodeGetArray(curve, "num_points"),
                       AiNodeGetArray(curve, "radius"));

      m_hairInfo.clear();
   }
}

void CShaveTranslator::ProcessHairLines(const unsigned int step,
                                        AtArray* curvePoints,
                                        AtArray* curveNumPoints,
                                        AtArray* curveWidths)
{
   // Get number of CVs per hair/line/strand.
   const int numPointsPerStep = m_hairInfo.numHairVertices + (m_hairInfo.numHairs * 2);

   // Tells us where the nextline starts.
   int curveLineInterpStartsIdx = 0;
   int curveLineStartsIdx       = 0;

   int nanCount = 0;
   // Process all hair lines
   for (int strand = 0; strand < m_hairInfo.numHairs; ++strand)
   {
      int numHairPoints = m_hairInfo.hairEndIndices[strand] - m_hairInfo.hairStartIndices[strand];

      // Curve radius
      const float rootRadius = m_hairInfo.rootRadii[strand];
      const float tipRadius  = m_hairInfo.tipRadii[strand];
      const float radiusStepSize = (rootRadius - tipRadius) / (numHairPoints-2);

      float curveSize = rootRadius;
      int index = 0;

      // This is a pointer to the first vertex in the hair strand.
      // It's incremement in the for loop below.
      shaveAPI::Vertex * vertex = &(m_hairInfo.vertices[m_hairInfo.hairStartIndices[strand]]);

      if(!AiIsFinite(vertex->x) || !AiIsFinite(vertex->y) || !AiIsFinite(vertex->z))
      {
         nanCount++;
         // the root point is NaN
         // not good...
#if (MTOA_ARCH_VERSION_NUM < 2)
         AtPoint previousPoint;
         AiV3Create(previousPoint, 0.f, 0.f, 0.f);
#else
         AtVector previousPoint (0.f, 0.f, 0.f);
#endif
         bool searchPoint = true;

         // search for the first valid point we can find in our full list
         // and repeat this position as the amount of points in this strand
         for (int tmpStrand = 0; (tmpStrand < m_hairInfo.numHairs) && searchPoint; ++tmpStrand)
         {
            shaveAPI::Vertex * tmpVertex = &(m_hairInfo.vertices[m_hairInfo.hairStartIndices[tmpStrand]]);
            for (int j = m_hairInfo.hairStartIndices[tmpStrand];
                 j < m_hairInfo.hairEndIndices[tmpStrand];
                 ++j, ++tmpVertex)
            {
               if (AiIsFinite(tmpVertex->x) && AiIsFinite(tmpVertex->y) && AiIsFinite(tmpVertex->z))
               {
                  // found a valid point in this hair system
                  searchPoint = false;
#if (MTOA_ARCH_VERSION_NUM < 2)
                  AiV3Create(previousPoint, tmpVertex->x, tmpVertex->y, tmpVertex->z);
#else
                  previousPoint = AtVector(tmpVertex->x, tmpVertex->y, tmpVertex->z);
#endif
                  break;
               }
            }
         }

         // copy this position for all points in this strand
         // yes, we'll loose this strand.... 
         // if we wanted to completely remove this strand, it would require 
         // a deeper rewriting of this code
         for (int j = 0; j < m_hairInfo.hairEndIndices[strand] - m_hairInfo.hairStartIndices[strand] + 2; ++j)
         {
#if (MTOA_ARCH_VERSION_NUM < 2)
            AiArraySetPnt(curvePoints,
                    (curveLineInterpStartsIdx + j  + (step * numPointsPerStep)),
                    previousPoint);
#else
            AiArraySetVec(curvePoints,
                    (curveLineInterpStartsIdx + j  + (step * numPointsPerStep)),
                    previousPoint);
#endif
         }
         if (step == 0)
         {
            // fill the strand width (as NULL)
            for (int j = 0; j < m_hairInfo.hairEndIndices[strand] - m_hairInfo.hairStartIndices[strand]; ++j)
            {         
               AiArraySetFlt(curveWidths, (j + curveLineStartsIdx) , AI_EPSILON);
            }
         }
         curveLineInterpStartsIdx += numHairPoints +2;
         curveLineStartsIdx += numHairPoints;
         continue;
      }
      

#if (MTOA_ARCH_VERSION_NUM < 2)
      AtPoint arnoldCurvePoint;
      AiV3Create(arnoldCurvePoint, vertex->x, vertex->y, vertex->z);
#else
      AtVector arnoldCurvePoint(vertex->x, vertex->y, vertex->z);
#endif
     

      // Create a first point on the curve. Start and end are duplicated vertices.
#if (MTOA_ARCH_VERSION_NUM < 2)
      AiArraySetPnt(curvePoints,
                    (curveLineInterpStartsIdx + (step * numPointsPerStep)),
                    arnoldCurvePoint);
#else
      AiArraySetVec(curvePoints,
                    (curveLineInterpStartsIdx + (step * numPointsPerStep)),
                    arnoldCurvePoint);
#endif
      
      bool nanFound = false; // says if we found an invalid value on this strand

      // Loop through all the hair strand indices.
      for (int j = m_hairInfo.hairStartIndices[strand];
           j < m_hairInfo.hairEndIndices[strand];
           ++j, ++index, ++vertex)
      {

         // Add the point.
#if (MTOA_ARCH_VERSION_NUM < 2)
         AiV3Create(arnoldCurvePoint, vertex->x, vertex->y, vertex->z);
#else
         arnoldCurvePoint = AtVector(vertex->x, vertex->y, vertex->z);
#endif

         shaveAPI::Vertex * tmpVertex = vertex;
         int tmpIndex = index;
         while(!AiIsFinite(arnoldCurvePoint.x) || !AiIsFinite(arnoldCurvePoint.y) || !AiIsFinite(arnoldCurvePoint.z))
         {
            if (!nanFound) nanCount++;

            nanFound = true;
            tmpVertex--;
            tmpIndex--;
            if (tmpIndex < 0) 
            {
               // this should never happen since the root point is valid
#if (MTOA_ARCH_VERSION_NUM < 2)
               AiV3Create(arnoldCurvePoint, 0.f, 0.f, 0.f);
#else
               arnoldCurvePoint = AtVector( 0.f, 0.f, 0.f);
#endif
               break;
            }
#if (MTOA_ARCH_VERSION_NUM < 2)
            AiV3Create(arnoldCurvePoint, tmpVertex->x, tmpVertex->y, tmpVertex->z);
#else
            arnoldCurvePoint = AtVector(tmpVertex->x, tmpVertex->y, tmpVertex->z);
#endif
         }

#if (MTOA_ARCH_VERSION_NUM < 2)
         AiArraySetPnt(curvePoints,
                       ((index+1) + curveLineInterpStartsIdx + (step * numPointsPerStep)),
                       arnoldCurvePoint);
#else
         AiArraySetVec(curvePoints,
                       ((index+1) + curveLineInterpStartsIdx + (step * numPointsPerStep)),
                       arnoldCurvePoint);
#endif

         // Animated widths are not supported, so just on STEP 0
         if (step == 0)
            AiArraySetFlt(curveWidths, (index + curveLineStartsIdx) , curveSize);

         // guard against minus values
         if (curveSize > 0.0f) curveSize -= radiusStepSize;
         if (curveSize < 0.0f) curveSize = 0.0f;
      }

      // Last point (duplicate of the previous point).
#if (MTOA_ARCH_VERSION_NUM < 2)
      AiArraySetPnt(curvePoints,
                    ((index+1) + curveLineInterpStartsIdx + (step * numPointsPerStep)),
                    arnoldCurvePoint);
#else
      AiArraySetVec(curvePoints,
                    ((index+1) + curveLineInterpStartsIdx + (step * numPointsPerStep)),
                    arnoldCurvePoint);
#endif

      curveLineInterpStartsIdx += numHairPoints +2;
      curveLineStartsIdx += numHairPoints;
   }
   if (nanCount > 0)
   {
      AiMsgError("[SHAVE]: %d invalid strands were found", nanCount);
   }
}

void CShaveTranslator::NodeInitializer(CAbTranslator context)
{
   CExtensionAttrHelper helper(context.maya, "curves");
   CShapeTranslator::MakeCommonAttributes(helper);
   helper.MakeInput("min_pixel_width");
   helper.MakeInput("mode");

   CAttrData data;

#if (MTOA_ARCH_VERSION_NUM < 2)
   data.defaultValue.BOOL = true;
#else
   data.defaultValue.BOOL() = true;
#endif
   data.name = "aiExportHairColors";
   data.shortName = "ai_export_hair_colours";
   helper.MakeInputBoolean(data);

#if (MTOA_ARCH_VERSION_NUM < 2)
   data.defaultValue.BOOL = true;
#else
   data.defaultValue.BOOL() = true;
#endif
   data.name = "aiExportHairIDs";
   data.shortName = "ai_export_hair_ids";
   helper.MakeInputBoolean(data);

#if (MTOA_ARCH_VERSION_NUM < 2)
   data.defaultValue.BOOL = false;
#else
   data.defaultValue.BOOL() = false;
#endif
   data.name = "aiOverrideHair";
   data.shortName = "ai_override_hair";
   helper.MakeInputBoolean(data);

   data.name = "aiHairShader";
   data.shortName = "ai_hair_shader";
   helper.MakeInputNode(data);

   CExtensionAttrHelper helper2(context.maya, "ShaveHair");
#if (MTOA_ARCH_VERSION_NUM < 2)
   helper2.MakeInput("diffuse_cache");
#endif

   helper2.MakeInput("direct_diffuse");

   helper2.GetAttrData("kd_ind", data);
   data.name = "aiIndirect";
   helper2.MakeInputFloat(data);

   helper2.MakeInput("aov_direct_diffuse");
   helper2.MakeInput("aov_direct_specular");
   helper2.MakeInput("aov_indirect_diffuse");
}