/******************************************************************************/
/*
  Project   - MudBun
  Publisher - Long Bunny Labs
              http://LongBunnyLabs.com
  Author    - Ming-Lun "Allen" Chou
              http://AllenChou.net
*/
/******************************************************************************/

#pragma kernel surface_nets_move_point
#pragma kernel surface_nets_move_point_2d

#include "../../Shader/ComputeCommon.cginc"

#include "../../Shader/AutoSmoothFuncs.cginc"
#include "../../Shader/BrushFuncs.cginc"
#include "../../Shader/DualMeshingFuncs.cginc"
#include "../../Shader/GenPointDefs.cginc"
#include "../../Shader/IndirectArgsDefs.cginc"
#include "../../Shader/Math/MathConst.cginc"
#include "../../Shader/MeshingModeDefs.cginc"
#include "../../Shader/NormalFuncs.cginc"
#include "../../Shader/RenderModeDefs.cginc"

// https://0fps.net/2012/07/12/smooth-voxel-terrain-part-2/

float surfaceNetsDualQuadsBlend;
int surfaceNetsBinarySearchIterations;
int surfaceNetsGradientDescentIterations;
float surfaceNetsGradientDescentFactor;

static int aEdgeVertIndex3d[12][2] = 
{
  { 0, 1 }, 
  { 1, 5 }, 
  { 5, 4 }, 
  { 4, 0 }, 
  { 2, 3 }, 
  { 3, 7 }, 
  { 7, 6 }, 
  { 6, 2 }, 
  { 0, 2 }, 
  { 1, 3 }, 
  { 5, 7 }, 
  { 4, 6 }, 
}; 

static int aEdgeVertIndex2d[4][2] = 
{
  { 0, 1 }, 
  { 1, 2 }, 
  { 2, 3 }, 
  { 3, 0 }, 
};

[numthreads(kThreadGroupSize, 1, 1)]
void surface_nets_move_point(int3 id : SV_DispatchThreadID)
{
#if defined(MUDBUN_DISABLE_SURFACE_NETS) || defined(MUDBUN_FAST_ITERATION)
  return;
#endif

  if (id.x >= indirectDrawArgs[0])
    return;

  int iGenPoint = id.x;

  if (surfaceNetsDualQuadsBlend >= 1.0f)
  {
    if (renderMode == kRenderModeQuadSplats)
      aGenPoint[iGenPoint].material.size *= 0.70711f;
    return;
  }

  float h = 0.5f * voxelSize;
  float3 center = aGenPoint[iGenPoint].posNorm.xyz;
  float3 minCornerOffset = -h;
  int iBrushMask = aGenPoint[iGenPoint].iBrushMask;
  
  float3 aCornerOffset[8] = 
  {
    float3(-h, -h, -h), 
    float3( h, -h, -h), 
    float3(-h,  h, -h), 
    float3( h,  h, -h), 
    float3(-h, -h,  h), 
    float3( h, -h,  h), 
    float3(-h,  h,  h), 
    float3( h,  h,  h), 
  };

  float aCornerRes[8];
  SdfBrushMaterial mat;
  [loop] for (int iCorner = 0; iCorner < 8; ++iCorner)
  {
    aCornerRes[iCorner] = sdf_masked_brushes(center + aCornerOffset[iCorner], iBrushMask, mat);
  }

  float3 avgEdgeOffset = 0.0f;
  int numEdges = 0;
  [loop] for (int iEdge = 0; iEdge < 12; ++iEdge)
  {
    int iVert0 = aEdgeVertIndex3d[iEdge][0];
    int iVert1 = aEdgeVertIndex3d[iEdge][1];
    float res0 = aCornerRes[iVert0];
    float res1 = aCornerRes[iVert1];
    if (res0 * res1 > 0)
      continue;

    ++numEdges;

    float3 offset0 = aCornerOffset[iVert0];
    float3 offset1 = aCornerOffset[iVert1];

    float3 edgeOffset;
    if (res0 == 0.0f && res1 == 0.0f)
    {
      edgeOffset = max(offset0, offset1);
    }
    else if (res0 == 0.0f)
    {
      edgeOffset = offset0;
    }
    else if (res1 == 0.0f)
    {
      edgeOffset = offset1;
    }
    else if (surfaceNetsBinarySearchIterations <= 0)
    {
      // lerp approximation
      float t = -res0 / (res1 - res0);
      edgeOffset = lerp(offset0, offset1, t);
    }
    /*
    else
    {
      // binary search
      edgeOffset = 0.5f * (offset0 + offset1);
      [loop] for (int iSearch = 0; iSearch < surfaceNetsBinarySearchIterations; ++iSearch)
      {
        float resT = sdf_masked_brushes(center + edgeOffset, iBrushMask, mat);
        if (res0 * resT < 0.0f)
        {
          res1 = resT;
          offset1 = edgeOffset;
        }
        else if (resT * res1 < 0.0f)
        {
          res0 = resT;
          offset0 = edgeOffset;
        }
        edgeOffset = 0.5f * (offset0 + offset1);
      }
    }
    */

    avgEdgeOffset += edgeOffset;
  }

  if (numEdges <= 0)
    return;

  avgEdgeOffset /= numEdges;
  float3 avgEdgePos = center + avgEdgeOffset;

  // gradient descent
  /*
  if (surfaceNetsGradientDescentIterations > 0)
  {
    float3 n;
    SDF_NORMAL_FULL(n, avgEdgePos, sdf_masked_brushes, iBrushMask, 1e-2f * voxelSize);
    [loop] for (int iDescent = 0; iDescent < surfaceNetsGradientDescentIterations; ++iDescent)
    {
      float d = sdf_masked_brushes(avgEdgePos, iBrushMask, mat);
      avgEdgeOffset -= surfaceNetsGradientDescentFactor * n * d;
      avgEdgePos = center + avgEdgeOffset;
    }
  }
  */

  aGenPoint[iGenPoint].posNorm.xyz = lerp(avgEdgePos, aGenPoint[iGenPoint].posNorm.xyz, surfaceNetsDualQuadsBlend);

  if (renderMode == kRenderModeQuadSplats)
    aGenPoint[iGenPoint].material.size *= 1.0f - (1.f - 0.70711f) * surfaceNetsDualQuadsBlend;
}

[numthreads(kThreadGroupSize, 1, 1)]
void surface_nets_move_point_2d(int3 id : SV_DispatchThreadID)
{
#if defined(MUDBUN_DISABLE_SURFACE_NETS) || defined(MUDBUN_FAST_ITERATION)
  return;
#endif

  if (id.x >= indirectDrawArgs[0])
    return;

  int iGenPoint = id.x;

  if (surfaceNetsDualQuadsBlend >= 1.0f)
  {
    if (renderMode == kRenderModeQuadSplats)
      aGenPoint[iGenPoint].material.size *= 0.70711f;
    return;
  }

  float h = 0.5f * voxelSize;
  float3 center = aGenPoint[iGenPoint].posNorm.xyz;
  float3 minCornerOffset = -h;
  int iBrushMask = aGenPoint[iGenPoint].iBrushMask;
  
  float3 aCornerOffset[4] = 
  {
    float3(-h, -h, 0.0f), 
    float3(-h,  h, 0.0f),
    float3( h,  h, 0.0f),
    float3( h, -h, 0.0f),
  };

  bool anyOutside = false;
  float aCornerRes[4];
  SdfBrushMaterial mat;
  [loop] for (int iCorner = 0; iCorner < 4; ++iCorner)
  {
    float cornerRes = sdf_masked_brushes(center + aCornerOffset[iCorner], iBrushMask, mat);
    aCornerRes[iCorner] = cornerRes;
    if (cornerRes >= 0.0f)
      anyOutside = true;
  }

  if (!anyOutside)
    return;

  float3 avgEdgeOffset = 0.0f;
  int numEdges = 0;
  [loop] for (int iEdge = 0; iEdge < 4; ++iEdge)
  {
    int iVert0 = aEdgeVertIndex2d[iEdge][0];
    int iVert1 = aEdgeVertIndex2d[iEdge][1];
    float res0 = aCornerRes[iVert0];
    float res1 = aCornerRes[iVert1];
    if (res0 * res1 > 0)
      continue;

    ++numEdges;

    float3 offset0 = aCornerOffset[iVert0];
    float3 offset1 = aCornerOffset[iVert1];

    float3 edgeOffset;
    if (res0 == 0.0f && res1 == 0.0f)
    {
      edgeOffset = max(offset0, offset1);
    }
    else if (res0 == 0.0f)
    {
      edgeOffset = offset0;
    }
    else if (res1 == 0.0f)
    {
      edgeOffset = offset1;
    }
    else if (surfaceNetsBinarySearchIterations <= 0)
    {
      // lerp approximation
      float t = -res0 / (res1 - res0);
      edgeOffset = lerp(offset0, offset1, t);
    }
    /*
    else
    {
      // binary search
      edgeOffset = 0.5f * (offset0 + offset1);
      [loop] for (int iSearch = 0; iSearch < surfaceNetsBinarySearchIterations; ++iSearch)
      {
        float resT = sdf_masked_brushes(center + edgeOffset, iBrushMask, mat);
        if (res0 * resT < 0.0f)
        {
          res1 = resT;
          offset1 = edgeOffset;
        }
        else if (resT * res1 < 0.0f)
        {
          res0 = resT;
          offset0 = edgeOffset;
        }
        edgeOffset = 0.5f * (offset0 + offset1);
      }
    }
    */

    avgEdgeOffset += edgeOffset;
  }

  if (numEdges <= 0)
    return;

  avgEdgeOffset /= numEdges;
  float3 avgEdgePos = center + avgEdgeOffset;

  // gradient descent
  /*
  if (surfaceNetsGradientDescentIterations > 0)
  {
    float3 n;
    SDF_NORMAL_2D(n, avgEdgePos, sdf_masked_brushes, iBrushMask, 1e-2f * voxelSize);
    [loop] for (int iDescent = 0; iDescent < surfaceNetsGradientDescentIterations; ++iDescent)
    {
      float d = sdf_masked_brushes(avgEdgePos, iBrushMask, mat);
      avgEdgeOffset -= surfaceNetsGradientDescentFactor * n * d;
      avgEdgePos = center + avgEdgeOffset;
    }
  }
  */

  aGenPoint[iGenPoint].posNorm.xyz = lerp(avgEdgePos, aGenPoint[iGenPoint].posNorm.xyz, surfaceNetsDualQuadsBlend);
  aGenPoint[iGenPoint].sdfValue = sdf_masked_brushes(avgEdgePos, iBrushMask, mat) + surfaceShift;

  if (renderMode == kRenderModeQuadSplats)
    aGenPoint[iGenPoint].material.size *= 1.0f - (1.f - 0.70711f) * surfaceNetsDualQuadsBlend;
}