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

#pragma kernel surface_nets_move_point

#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"
#include "../../Shader/VoxelFuncs.cginc"
#include "../../Shader/VoxelCacheFuncs.cginc"

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

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

[numthreads(kThreadGroupSize, 1, 1)]
void surface_nets_move_point(int3 id : SV_DispatchThreadID)
{
#if defined(MUDBUN_DISABLE_SURFACE_NETS)
  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), 
  };

  int aEdgeVertIndex[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 }, 
  };

  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;
  float3 aEdgeOffset[12];
  for (int iEdge = 0; iEdge < 12; ++iEdge)
  {
    int iVert0 = aEdgeVertIndex[iEdge][0];
    int iVert1 = aEdgeVertIndex[iEdge][1];
    float res0 = aCornerRes[iVert0];
    float res1 = aCornerRes[iVert1];
    if (res0 * res1 > 0)
      continue;

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

    int iEdgeOutput = numEdges++;

    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;
    aEdgeOffset[iEdgeOutput] = edgeOffset;
  }

  if (numEdges <= 0)
    return;

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

  // gradient descent
  {
    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;
}

