/*=========================================================================

  Program:   ParaView
  Module:    vtkBivariateNoiseMapper.cxx

  Copyright (c) Kitware, Inc.
  All rights reserved.
  See Copyright.txt or http://www.paraview.org/HTML/Copyright.html for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notice for more information.

=========================================================================*/
#include "vtkBivariateNoiseMapper.h"
#include "vtkCompositeMapperHelper2.h"
#include "vtkFloatArray.h"
#include "vtkObjectFactory.h"
#include "vtkOpenGLUniforms.h"
#include "vtkOpenGLVertexBufferObjectGroup.h"
#include "vtkPointData.h"
#include "vtkPolyData.h"
#include "vtkRenderer.h"
#include "vtkShader.h"
#include "vtkShaderProgram.h"
#include "vtkShaderProperty.h"

#include <chrono>

//----------------------------------------------------------------------------
struct vtkBivariateNoiseMapper::vtkInternals
{
  double FrequencyModifier = 10.0;
  double AmplitudeModifier = 0.05;
  long StartTime = 0;
};

//----------------------------------------------------------------------------
class vtkBivariateNoiseMapperHelper : public vtkCompositeMapperHelper2
{
public:
  static vtkBivariateNoiseMapperHelper* New();
  vtkTypeMacro(vtkBivariateNoiseMapperHelper, vtkCompositeMapperHelper2);

protected:
  vtkBivariateNoiseMapperHelper() = default;
  ~vtkBivariateNoiseMapperHelper() override = default;

  /**
   * Contain most of the shader replacements specific to this mapper.
   * Define the noise function and use it to generate the output of the fragment shader.
   */
  void ReplaceShaderColor(
    std::map<vtkShader::Type, vtkShader*> shaders, vtkRenderer* ren, vtkActor* act) override;

  /**
   * Pass and interpolate vertex positions (in model coordinates)
   * from the vertex shader to the fragment shader.
   */
  void ReplaceShaderPositionVC(
    std::map<vtkShader::Type, vtkShader*> shaders, vtkRenderer* ren, vtkActor* actor) override;

  /**
   * Define the custom uniforms in the fragment shader (frequency, amplitude, time).
   */
  void ReplaceShaderCustomUniforms(
    std::map<vtkShader::Type, vtkShader*> shaders, vtkActor* actor) override;

  /**
   * Set the custom uniforms values in the fragment shader.
   */
  void SetCustomUniforms(vtkOpenGLHelper& cellBO, vtkActor* actor) override;

  /**
   * Pass the 2nd array to visualize with noise intensity.
   */
  void AppendOneBufferObject(vtkRenderer* ren, vtkActor* act, vtkCompositeMapperHelperData* hdata,
    vtkIdType& flat_index, std::vector<unsigned char>& colors, std::vector<float>& norms) override;

private:
  vtkBivariateNoiseMapperHelper(const vtkBivariateNoiseMapperHelper&) = delete;
  void operator=(const vtkBivariateNoiseMapperHelper&) = delete;
};

//----------------------------------------------------------------------------
vtkStandardNewMacro(vtkBivariateNoiseMapperHelper);

//----------------------------------------------------------------------------
void vtkBivariateNoiseMapperHelper::ReplaceShaderColor(
  std::map<vtkShader::Type, vtkShader*> shaders, vtkRenderer* ren, vtkActor* act)
{
  // InterpolateScalarsBeforeMapping: noise effect is applied during color lookup in the FS.
  // ColorCoordinates: we need the first array values.
  if (this->InterpolateScalarsBeforeMapping && this->ColorCoordinates && !this->DrawingVertices)
  {
    // The bivariate data correspond to the second array to show
    // in the form of Curl noise intensity.
    std::string colorDecVS =
      R"(
in float bivariateData;
out float vertexBivariateDataVSOut;
)";

    // The bivariate data should be interpolated between VS and FS.
    // We interpolate the scalars before mapping the colors in the FS.
    std::string colorImplVS =
      R"(
  vertexBivariateDataVSOut = bivariateData;
  )";

    // Here we define our Curl 3D noise function.
    std::string colorDecFS =
      R"(
// Hash function
vec3 Hash(in vec3 p, in float numCells)
{
  // This is tiling part, adjusts with the scale
  p = mod(p, numCells);

  p = vec3( dot(p,vec3(127.1,311.7, 74.7)),
      dot(p,vec3(269.5,183.3,246.1)),
      dot(p,vec3(113.5,271.9,124.6)));

  return -1.0 + fract(sin(p)*43758.5453123) * 2.0;
}

// Base noise function
float TileableNoise(in vec3 p, in float numCells )
{
  vec3 f, i;

  p *= numCells;

  f = fract(p); // Separate integer from fractional
  i = floor(p);

  vec3 u = f*f*(3.0-2.0*f); // Cosine interpolation approximation

  return mix( mix( mix( dot( Hash( i + vec3(0.0,0.0,0.0), numCells ), f - vec3(0.0,0.0,0.0) ),
                        dot( Hash( i + vec3(1.0,0.0,0.0), numCells ), f - vec3(1.0,0.0,0.0) ), u.x),
                   mix( dot( Hash( i + vec3(0.0,1.0,0.0), numCells ), f - vec3(0.0,1.0,0.0) ),
                        dot( Hash( i + vec3(1.0,1.0,0.0), numCells ), f - vec3(1.0,1.0,0.0) ), u.x), u.y),
              mix( mix( dot( Hash( i + vec3(0.0,0.0,1.0), numCells ), f - vec3(0.0,0.0,1.0) ),
                        dot( Hash( i + vec3(1.0,0.0,1.0), numCells ), f - vec3(1.0,0.0,1.0) ), u.x),
                   mix( dot( Hash( i + vec3(0.0,1.0,1.0), numCells ), f - vec3(0.0,1.0,1.0) ),
                        dot( Hash( i + vec3(1.0,1.0,1.0), numCells ), f - vec3(1.0,1.0,1.0) ), u.x), u.y), u.z );
}

// Fractal Brownian Motion
// Add multiple iteration of noise over different octaves
float TileableNoiseFBM(in vec3 p, float numCells, int octaves)
{
  float f = 0.0;

  // Change starting scale to any integer value...
  p = mod(p, vec3(numCells));
  float amp = 0.5;
  float sum = 0.0;

  for (int i = 0; i < octaves; i++)
  {
    f += TileableNoise(p, numCells) * amp;
    sum += amp;
    amp *= 0.5;

    // numCells must be multiplied by an integer value...
    numCells *= 2.0;
  }

  return f / sum;
}

vec3 snoiseVec3( vec3 x, in float numCells, int octaves )
{
  float s  = TileableNoiseFBM(vec3( x ), numCells, octaves);
  float s1 = TileableNoiseFBM(vec3( x.y - 19.1 , x.z + 33.4 , x.x + 47.2 ), numCells, octaves);
  float s2 = TileableNoiseFBM(vec3( x.z + 74.2 , x.x - 124.5 , x.y + 99.4 ), numCells, octaves);
  vec3 c = vec3( s , s1 , s2 );
  return c;
}

// Final Curl noise
vec3 TileableCurlNoise(in vec3 p, in float numCells, in int octaves)
{
  const float e = .1;
  vec3 dx = vec3( e   , 0.0 , 0.0 );
  vec3 dy = vec3( 0.0 , e   , 0.0 );
  vec3 dz = vec3( 0.0 , 0.0 , e   );

  vec3 p_x0 = snoiseVec3( p - dx, numCells, octaves );
  vec3 p_x1 = snoiseVec3( p + dx, numCells, octaves );
  vec3 p_y0 = snoiseVec3( p - dy, numCells, octaves );
  vec3 p_y1 = snoiseVec3( p + dy, numCells, octaves );
  vec3 p_z0 = snoiseVec3( p - dz, numCells, octaves );
  vec3 p_z1 = snoiseVec3( p + dz, numCells, octaves );

  float x = p_y1.z - p_y0.z - p_z1.y + p_z0.y;
  float y = p_z1.x - p_z0.x - p_x1.z + p_x0.z;
  float z = p_x1.y - p_x0.y - p_y1.x + p_y0.x;

  const float divisor = 1.0 / ( 2.0 * e );
  return normalize( vec3( x , y , z ) * divisor );
}

in float vertexBivariateDataVSOut;
)";

    // Here we apply the value of the curl noise (depending on the 2nd data array)
    // on the value of 2D texture coordinates (tcoordVCVSOutput.s corresponds to the 1st data
    // array with values clamped between 0.0 and 1.0, tcoordVCVSOutput.t is always equal to 0).
    // This ends up oscillating over the color texture (1D).
    // Bigger is the 2nd data array value, bigger is the amplitude.
    std::string colorImplFS =
      R"(
  // Multiply time with displacement vector to "move" in the 3D curl noise.
  // This vector had been determined empirically to obtain a good result on 2D data.
  vec3 curl_offset = currentTime * vec3(0.0, 0.2, 0.8);

  // Compute and apply the noise value to modify the color texture coordinates.
  // The number of cells and octaves are chosen empirically, based on the visual result and performances.
  vec3 curl = TileableCurlNoise((vertexMCVSOutput.xyz * frequencyMod + curl_offset), 10.0, 2) - vec3(-1.);
  vec2 _texCoord = tcoordVCVSOutput.st + vec2(curl.x, 0.0) * amplitudeMod * vertexBivariateDataVSOut;
  )";

    vtkShaderProgram::Substitute(
      shaders[vtkShader::Vertex], "//VTK::Color::Dec", colorDecVS + "\n//VTK::Color::Dec\n");
    vtkShaderProgram::Substitute(
      shaders[vtkShader::Vertex], "//VTK::Color::Impl", colorImplVS + "\n//VTK::Color::Impl\n");
    vtkShaderProgram::Substitute(
      shaders[vtkShader::Fragment], "//VTK::Color::Dec", colorDecFS + "\n//VTK::Color::Dec\n");
    vtkShaderProgram::Substitute(
      shaders[vtkShader::Fragment], "//VTK::Color::Impl", colorImplFS + "\n//VTK::Color::Impl\n");

    this->Superclass::ReplaceShaderColor(shaders, ren, act);

    // Now we need to override the texColor value (after the replacements defined in
    // vtkOpenGLPolyDataMapper) in order to use our custom _texCoord value.
    vtkShaderProgram::Substitute(shaders[vtkShader::Fragment],
      "vec4 texColor = texture(colortexture, tcoordVCVSOutput.st);",
      "vec4 texColor = texture(colortexture, _texCoord);");
  }
  else
  {
    // Just call superclass
    this->Superclass::ReplaceShaderColor(shaders, ren, act);
  }
}

//----------------------------------------------------------------------------
void vtkBivariateNoiseMapperHelper::ReplaceShaderPositionVC(
  std::map<vtkShader::Type, vtkShader*> shaders, vtkRenderer* ren, vtkActor* actor)
{
  // We need to interpolate the vertex positions (in model coordinate)
  // between the vertex shader and the fragment shader, because the Curl
  // noise function take it as a parameter.
  vtkShaderProgram::Substitute(shaders[vtkShader::Vertex], "//VTK::PositionVC::Dec",
    R"(
//VTK::PositionVC::Dec
out vec4 vertexMCVSOutput;
)");

  vtkShaderProgram::Substitute(shaders[vtkShader::Vertex], "//VTK::PositionVC::Impl",
    R"(
  //VTK::PositionVC::Impl
  vertexMCVSOutput = vertexMC;
  )");

  // Fragment shader substitutions
  vtkShaderProgram::Substitute(shaders[vtkShader::Fragment], "//VTK::PositionVC::Dec",
    R"(
//VTK::PositionVC::Dec
in vec4 vertexMCVSOutput;
)");

  this->Superclass::ReplaceShaderPositionVC(shaders, ren, actor);
}

//----------------------------------------------------------------------------
void vtkBivariateNoiseMapperHelper::ReplaceShaderCustomUniforms(
  std::map<vtkShader::Type, vtkShader*> shaders, vtkActor* actor)
{
  vtkShaderProgram::Substitute(shaders[vtkShader::Fragment], "//VTK::CustomUniforms::Dec",
    R"(
//VTK::CustomUniforms::Dec
uniform float frequencyMod = 1.0;
uniform float amplitudeMod = 1.0;
uniform float currentTime = 1.0;
)");

  this->Superclass::ReplaceShaderCustomUniforms(shaders, actor);
}

//----------------------------------------------------------------------------
void vtkBivariateNoiseMapperHelper::SetCustomUniforms(vtkOpenGLHelper& cellBO, vtkActor* actor)
{
  this->Superclass::SetCustomUniforms(cellBO, actor);
  vtkBivariateNoiseMapper* parent = vtkBivariateNoiseMapper::SafeDownCast(this->Parent);
  cellBO.Program->SetUniformf("frequencyMod", parent->Internals->FrequencyModifier);
  cellBO.Program->SetUniformf("amplitudeMod", parent->Internals->AmplitudeModifier);
  auto time =
    (std::chrono::steady_clock::now().time_since_epoch().count() - parent->Internals->StartTime) *
    0.0000000003;
  cellBO.Program->SetUniformf("currentTime", time);
}

//------------------------------------------------------------------------------
void vtkBivariateNoiseMapperHelper::AppendOneBufferObject(vtkRenderer* ren, vtkActor* act,
  vtkCompositeMapperHelperData* hdata, vtkIdType& voffset, std::vector<unsigned char>& newColors,
  std::vector<float>& newNorms)
{
  vtkPolyData* poly = hdata->Data;
  vtkDataArray* array = this->GetInputArrayToProcess(0, poly);
  if (array)
  {
    vtkNew<vtkFloatArray> floatArray;
    floatArray->DeepCopy(array);
    this->VBOs->AppendDataArray("bivariateData", floatArray, VTK_FLOAT);
  }

  this->Superclass::AppendOneBufferObject(ren, act, hdata, voffset, newColors, newNorms);
}

//----------------------------------------------------------------------------
vtkStandardNewMacro(vtkBivariateNoiseMapper);

//----------------------------------------------------------------------------
vtkBivariateNoiseMapper::vtkBivariateNoiseMapper()
  : Internals(new vtkBivariateNoiseMapper::vtkInternals())
{
  this->Internals->StartTime = std::chrono::steady_clock::now().time_since_epoch().count();
}

//----------------------------------------------------------------------------
vtkBivariateNoiseMapper::~vtkBivariateNoiseMapper() = default;

//----------------------------------------------------------------------------
void vtkBivariateNoiseMapper::PrintSelf(ostream& os, vtkIndent indent)
{
  this->Superclass::PrintSelf(os, indent);
}

//----------------------------------------------------------------------------
vtkCompositeMapperHelper2* vtkBivariateNoiseMapper::CreateHelper()
{
  auto* helper = vtkBivariateNoiseMapperHelper::New();
  return helper;
}

//----------------------------------------------------------------------------
void vtkBivariateNoiseMapper::CopyMapperValuesToHelper(vtkCompositeMapperHelper2* helper)
{
  this->Superclass::CopyMapperValuesToHelper(helper);
  helper->SetInputArrayToProcess(0, this->GetInputArrayInformation(0));
}

//----------------------------------------------------------------------------
void vtkBivariateNoiseMapper::SetFrequencyModifier(double fMod)
{
  this->Internals->FrequencyModifier = fMod;
}

//----------------------------------------------------------------------------
void vtkBivariateNoiseMapper::SetAmplitudeModifier(double fMod)
{
  this->Internals->AmplitudeModifier = fMod * 0.01;
}
