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

//
// Simple hair shader, roughly based on Kay/Kajiya's shading model
//

#include <ai.h>

#include <cstring>
#include <string>

#define MKSTR_( x) #x
#define MKSTR(x) MKSTR_(x)

#ifndef MIN_ARNOLD_VERSION
# error MIN_ARNOLD_VERSION not defined.
#endif

#ifndef MAX_ARNOLD_VERSION
# error MAX_ARNOLD_VERSION not defined.
#endif

static void decodeVersion(const char* versionStr, int &arch, int &major, int &minor, int &fix)
{
    char c;
    int curPart = 0;
    int parts[4] = { 0, 0, 0, 0 };

    while ((curPart < 4) && (c = *versionStr++))
    {
        if (c == '.')
        {
            ++curPart;
        }
        else
        {
            parts[curPart] = parts[curPart] * 10 + (c - '0');
        }
    }

    arch = parts[0];
    major = parts[1];
    minor = parts[2];
    fix = parts[3];
}


static int compareVersions(
        int arch1, int major1, int minor1, int fix1,
        int arch2, int major2, int minor2, int fix2
)
{
    if (arch1 < arch2) return -1;
    if (arch1 > arch2) return 1;

    if (major1 < major2) return -1;
    if (major1 > major2) return 1;

    if (minor1 < minor2) return -1;
    if (minor1 > minor2) return 1;

    if (fix1 < fix2) return -1;
    if (fix1 > fix2) return 1;

    return 0;
}
    

AI_SHADER_NODE_EXPORT_METHODS(ShaveHairMtd);

enum ShaveHairParams
{
   p_rootcolor,
   p_tipcolor,
   p_strand_opacity,
   p_ambdiff,
   p_ambient,
   p_gloss,
   p_spec_color,
   p_spec,
   p_kd_ind,
#if (AI_VERSION_ARCH_NUM < 5)
   p_uparam,
   p_vparam,
#endif
   p_direct_diffuse,
   p_indirect_diffuse,
#if (AI_VERSION_ARCH_NUM < 5)
   p_diffuse_cache,
#endif
   p_aov_direct_diffuse,
   p_aov_direct_specular,
   p_aov_indirect_diffuse
};

#if (AI_VERSION_ARCH_NUM < 5)
#define nentry mds
#endif

node_parameters
{
   AiParameterStr(       "rootcolor"        , NULL);
   AiParameterStr(       "tipcolor"         , NULL);
   AiParameterRGB(       "strand_opacity"   , 1.0f, 1.0f, 1.0f);
   AiParameterFlt(       "ambdiff"          , 1.0f);
   AiParameterRGB(       "ambient"          , 1.0f, 1.0f, 1.0f);
   AiParameterFlt(       "gloss"            , 10.0f);
   AiParameterRGB(       "spec_color"       , 1.0f, 1.0f, 1.0f);
   AiParameterFlt(       "spec"             , 1.0f);
   AiParameterFlt(       "kd_ind"           , 1.0f);

   AiMetaDataSetFlt(nentry, "kd_ind"           , "softmax", 10.0f);
   AiMetaDataSetFlt(nentry, "kd_ind"           , "min",     0.0f);
#if (AI_VERSION_ARCH_NUM < 5)
   AiParameterSTR(       "uparam"           , NULL);
   AiParameterSTR(       "vparam"           , NULL);
#endif
   AiParameterFlt(       "direct_diffuse"   , 1.0f);
   AiMetaDataSetFlt(nentry, "direct_diffuse"   , "softmax", 1.0f);
   AiMetaDataSetFlt(nentry, "direct_diffuse"   , "min",     0.0f);
   AiParameterFlt(       "indirect_diffuse" , 1.0f);
#if (AI_VERSION_ARCH_NUM < 5)
   AiParameterBOOL(      "diffuse_cache"    , true);
#endif
 
   AiParameterStr ( "aov_direct_diffuse"             , "direct_diffuse"    );
   AiMetaDataSetInt(nentry, "aov_direct_diffuse"        , "aov.type", AI_TYPE_RGB);
   AiParameterStr ( "aov_direct_specular"            , "direct_specular"   );
   AiMetaDataSetInt(nentry, "aov_direct_specular"       , "aov.type", AI_TYPE_RGB);
   AiParameterStr ( "aov_indirect_diffuse"           , "indirect_diffuse"  );
   AiMetaDataSetInt(nentry, "aov_indirect_diffuse"      , "aov.type", AI_TYPE_RGB);

   AiMetaDataSetStr(nentry, NULL, "maya.name", "shaveHair");

    //  mtoa will attempt to create an instance of the 'shaveHair' node
    //  type. If the Shave plugin is not yet loaded then mtoa will register
    //  the node type itself which will cause Shave to fail.
    //
    //  We can prevent mtoa from creating the instance by setting it to
    //  hidden.
    //
    AiMetaDataSetBool(nentry, NULL, "maya.hide", true);
}

#if (AI_VERSION_ARCH_NUM < 5)
#undef nentry
#endif

typedef struct
{
#if (AI_VERSION_ARCH_NUM < 5)
   float              gamma;
#endif
   int                max_diffuse_depth;
}
ShaderData;


node_initialize
{
   ShaderData *data = (ShaderData*) AiMalloc(sizeof(ShaderData));
   data->max_diffuse_depth = -1;
   AiNodeSetLocalData(node, data);
}

node_update
{
#if (AI_VERSION_ARCH_NUM < 5) || ((AI_VERSION_ARCH_NUM == 5) && (AI_VERSION_MAJOR_NUM < 1))
    AtNode *options = AiUniverseGetOptions();
#else
    AtUniverse *universe = AiNodeGetUniverse(node);
    AtNode *options = AiUniverseGetOptions(universe);
#endif
    ShaderData *data = (ShaderData*)AiNodeGetLocalData(node);
    int depth = AiNodeGetInt(options, "GI_diffuse_depth");
    if (depth >= 0)
        data->max_diffuse_depth = depth;
    else
        data->max_diffuse_depth = 0;
#if (AI_VERSION_ARCH_NUM < 5)
    data->gamma = 1.0f / AiNodeGetFlt(options, "shader_gamma");
#endif
}

node_finish
{
   ShaderData *data = (ShaderData*)AiNodeGetLocalData(node);
   AiFree(data);
}

shader_evaluate
{
   AtRGB opacity = AiShaderEvalParamRGB(p_strand_opacity);

   // This piece of user-data is automatically set by the curves node when
   // using auto-enlargement (min_pixel_width > 0)
   float geo_opacity;
   static AtString geo_opacityStr("geo_opacity");
#if (AI_VERSION_ARCH_NUM < 5)
   if (AiUDataGetFlt(geo_opacityStr, &geo_opacity))
#else
   if (AiUDataGetFlt(geo_opacityStr, geo_opacity))
#endif
   {
      opacity *= geo_opacity;
   }  
      
#if (AI_VERSION_ARCH_NUM < 5)
   if (AiShaderGlobalsApplyOpacity(sg, opacity))
      return;
      
   // early out for shadow rays and totally transparent objects
   if ((sg->Rt & AI_RAY_SHADOW) || AiColorIsZero(sg->out_opacity))
      return;
#else
   // early out for shadow rays and totally transparent objects
   if (sg->Rt & AI_RAY_SHADOW)
      return;
#endif

   AtVector V = -sg->Rd;
   AtVector T = AiV3Normalize(sg->dPdv);

   AtRGB Cdiff = AI_RGB_BLACK;
   AtRGB Cspec = AI_RGB_BLACK;

#if (AI_VERSION_ARCH_NUM < 5)
   // change the current (u,v) position according to the specified userdata channels
   float oldU = sg->u;
   float oldV = sg->v;

   AtParamValue *params = AiNodeGetParams(node);

   AiUDataGetFlt(params[p_uparam].STR, &(sg->u));
   AiUDataGetFlt(params[p_vparam].STR, &(sg->v));
#endif

   float ambdiff    = AiShaderEvalParamFlt(p_ambdiff);
   float gloss      = AiShaderEvalParamFlt(p_gloss) * 2000;
   float spec       = AiShaderEvalParamFlt(p_spec);

   float direct_c   = AiShaderEvalParamFlt(p_direct_diffuse);
   float indirect_c = AiShaderEvalParamFlt(p_indirect_diffuse);

   AtRGB spec_color = AiShaderEvalParamRGB(p_spec_color);
   AtRGB root_color;
   AtRGB tip_color;

#if (AI_VERSION_ARCH_NUM < 5)
   // FIXME: we need to gamma correct according to global settings
   AiUDataGetRGB(params[p_rootcolor].STR, &root_color);
   AiUDataGetRGB(params[p_tipcolor].STR, &tip_color);

   ShaderData *data = (ShaderData*)AiNodeGetLocalData(node);
   AiColorGamma(&root_color, data->gamma);
   AiColorGamma(&tip_color, data->gamma);

   // restore original (u,v)
   sg->u = oldU;
   sg->v = oldV;

   // mix root and tip colors
   AtColor diff_color;
   AiColorLerp(diff_color, sg->v, root_color, tip_color);

   if (AiShaderEvalParamBool(p_diffuse_cache) && (sg->Rt & AI_RAY_DIFFUSE))
   {
      // quick viz
      Cdiff = AiHairDirectDiffuseCache(sg);
      sg->out.RGB = Cdiff * diff_color;
      return;
   }
#else
   AiUDataGetRGB(AiShaderEvalParamStr(p_rootcolor), root_color);
   AiUDataGetRGB(AiShaderEvalParamStr(p_tipcolor), tip_color);
   
   // mix root and tip colors
   AtRGB diff_color = (sg->v * tip_color) + ((1.f - sg->v) * root_color);
#endif
   
   // Since the curves are represented as ray-facing ribbons
   // their normal is usually pointing roughly towards the incoming ray
   // - but not always. It can have discontinuities - which then causes
   // some lights to be ignored if they happen to be on the wrong side of this
   // normal. Setting sg->fhemi forces the lights to be gathered regardless of
   // where the normal is pointing
   sg->fhemi = false;

   // direct lighting
   AiLightsPrepare(sg);
#if (AI_VERSION_ARCH_NUM < 5)
   while (AiLightsGetSample(sg))
   {
      if (AiLightGetAffectDiffuse(sg->Lp))
      {
         float TdotL = (float)AiV3Dot(T, sg->Ld);
         float d = 1 - TdotL * TdotL;
         d = d > 0 ? sqrtf(d) : 0;
         float diffterm = (1 - ambdiff) + d * ambdiff;  // limits gamut of diffuse term
         // diffuse x illumination
         Cdiff += (sg->Li * diffterm * sg->we) * direct_c;
      }

      if (spec > 0 && AiLightGetAffectSpecular(sg->Lp))
      {
         AtVector H = sg->Ld + V;
         AiV3Normalize(H, H);
         float HdotT = (float)AiV3Dot(H, T);
         float s = 1 - HdotT * HdotT;
         if (s > 0)
         {
            // note: s holds sin^2 of the angle between H and T
            //       we don't need to take the sqrt, because we
            //       compensate by halving the gloss factor
            s = powf(s, gloss * 0.5f);
            // specular exponent x illumination
            Cspec += sg->Li * s * sg->we;
         }
      }
   }
#else
   AtLightSample light_sample;
   while (AiLightsGetSample(sg, light_sample))
   {
      float lDiff = AiLightGetDiffuse(light_sample.Lp);
      if (lDiff > 0)
      {
         float TdotL = (float)AiV3Dot(T, light_sample.Ld);
         float d = 1 - TdotL * TdotL;
         d = d > 0 ? sqrtf(d) : 0;
         float diffterm = (1 - ambdiff) + d * ambdiff;  // limits gamut of diffuse term
         // diffuse x illumination
         Cdiff += (light_sample.Li * diffterm) * direct_c * lDiff;
      }

      float lSpec = AiLightGetSpecular(light_sample.Lp);
      if (spec > 0 && lSpec > 0)
      {
         AtVector H = light_sample.Ld + V;
         H = AiV3Normalize(H);
         float HdotT = (float)AiV3Dot(H, T);
         float s = 1 - HdotT * HdotT;
         if (s > 0)
         {
            // note: s holds sin^2 of the angle between H and T
            //       we don't need to take the sqrt, because we
            //       compensate by halving the gloss factor
            s = powf(s, gloss * 0.5f);
            // specular exponent x illumination
            Cspec += light_sample.Li * s * lSpec;
         }
      }
   }
#endif
   Cdiff *= diff_color;
   AiAOVSetRGB(sg, AiShaderEvalParamStr(p_aov_direct_diffuse), Cdiff);

   Cspec *= spec * spec_color;
   AiAOVSetRGB(sg, AiShaderEvalParamStr(p_aov_direct_specular), Cspec);

   // indirect diffuse or ambient
   // if kd_ind = 0 then use ambient
   // FIXME: we should be checking for the arnold diffuse depth before we use Indirect gi
   //
   float kd_ind = AiShaderEvalParamFlt(p_kd_ind);
   AtRGB ind_diff;
   if (kd_ind > 0)
   {
#if (AI_VERSION_ARCH_NUM < 5)
      ind_diff = (kd_ind * AiIndirectDiffuse(&V,sg)) * indirect_c;
#else
      ind_diff = (kd_ind * AiIndirectDiffuse(V,sg, AI_RGB_WHITE)) * indirect_c;
#endif
   }
   else
   {
      ind_diff = AiShaderEvalParamRGB(p_ambient);
   }

   ind_diff *= diff_color;
   AiAOVSetRGB(sg, AiShaderEvalParamStr(p_aov_indirect_diffuse), ind_diff);
   Cdiff += ind_diff;

#if (AI_VERSION_ARCH_NUM < 5)
   sg->out.RGB = Cdiff + Cspec;
#else
   sg->out.RGB() = Cdiff + Cspec;
#endif
}

node_loader
{
   if (i > 0)
      return false;

    //  Make sure that we support this version of Arnold.
    //
    const char *minVersionStr = MKSTR( MIN_ARNOLD_VERSION);
    int minArch, minMajor, minMinor, minFix;
    decodeVersion(minVersionStr, minArch, minMajor, minMinor, minFix);

    const char *maxVersionStr = MKSTR( MAX_ARNOLD_VERSION);
    int maxArch, maxMajor, maxMinor, maxFix;
    decodeVersion(maxVersionStr, maxArch, maxMajor, maxMinor, maxFix);

    char archStr[4];
    char majorStr[4];
    char minorStr[4];
    char fixStr[4];

    AiGetVersion(archStr, majorStr, minorStr, fixStr);

    int curArch = atoi(archStr);
    int curMajor = atoi(majorStr);
    int curMinor = atoi(minorStr);
    int curFix = atoi(fixStr);

    if ((compareVersions(curArch, curMajor, curMinor, curFix, minArch, minMajor, minMinor, minFix) < 0)
    ||  (compareVersions(curArch, curMajor, curMinor, curFix, maxArch, maxMajor, maxMinor, maxFix) > 0))
    {
        AiMsgWarning("Shave: Ignore the compatibility warning for shave_shaders-%s below. It's harmless.", AI_VERSION);
        return false;
    }

   node->methods      = ShaveHairMtd;
   node->output_type  = AI_TYPE_RGB;
   node->name         = "ShaveHair";
   node->node_type    = AI_NODE_SHADER;
   strcpy(node->version, AI_VERSION);
   return true;
}