// SPDX-FileCopyrightText: Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
// SPDX-License-Identifier: BSD-3-Clause

#include "vtkWebGPUComputeFrustumCuller.h"
#include "FrustumCulling.h"   // For the frustum culling shader source code
#include "vtkActor.h"         // for safe downcasting in the bounds recomputed callback
#include "vtkCamera.h"        // for manipulating the camera data used by the shader
#include "vtkCommand.h"       // for handling observers on bounds recomputed
#include "vtkMatrix4x4.h"     // for manipulating the view-projection matrix
#include "vtkObjectFactory.h" // for vtk standard new macro
#include "vtkProp.h"          // for manipulating props
#include "vtkRenderer.h"      // for manipulating the renderer
#include "vtkWebGPUComputePass.h"

VTK_ABI_NAMESPACE_BEGIN

vtkStandardNewMacro(vtkWebGPUComputeFrustumCuller);

//------------------------------------------------------------------------------
vtkWebGPUComputeFrustumCuller::vtkWebGPUComputeFrustumCuller()
{
  this->Pipeline = vtkSmartPointer<vtkWebGPUComputePipeline>::New();
  this->FrustumCullingPass = this->Pipeline->CreateComputePass();
  this->FrustumCullingPass->SetShaderSource(FrustumCulling);
  this->FrustumCullingPass->SetShaderEntryPoint("frustumCullingEntryPoint");

  // Buffer that will contain the indices of the objects that were not culled.
  // This buffer will need to be initialized when the cull method is called because we don't know
  // how many actor we're going to cull yet (and thus we cannot initialize the size of this buffer)
  this->OutputIndicesBufferIndex = -1;

  // Buffer that will contain the number of objects that have not been culled
  vtkSmartPointer<vtkWebGPUComputeBuffer> outputCountBuffer =
    vtkSmartPointer<vtkWebGPUComputeBuffer>::New();
  outputCountBuffer->SetGroup(0);
  outputCountBuffer->SetBinding(2);
  outputCountBuffer->SetMode(vtkWebGPUComputeBuffer::BufferMode::READ_WRITE_MAP_COMPUTE_STORAGE);
  outputCountBuffer->SetByteSize(sizeof(int));
  outputCountBuffer->SetLabel("Frustum culler output count buffer");
  this->OutputObjectCountBufferIndex = this->FrustumCullingPass->AddBuffer(outputCountBuffer);

  // Actor bounds callback
  this->ActorBoundsRecomputedCallback = vtkSmartPointer<vtkCallbackCommand>::New();
  this->ActorBoundsRecomputedCallback->SetCallback(
    vtkWebGPUComputeFrustumCuller::BoundsRecomputedCallbackFunction);
  this->ActorBoundsRecomputedCallback->SetClientData(this);

  // Camera modified callback
  this->CameraModifiedCallback = vtkSmartPointer<vtkCallbackCommand>::New();
  this->CameraModifiedCallback->SetCallback(
    vtkWebGPUComputeFrustumCuller::CameraModifiedCallbackFunction);
  this->CameraModifiedCallback->SetClientData(this);
}

//------------------------------------------------------------------------------
double vtkWebGPUComputeFrustumCuller::Cull(
  vtkRenderer* ren, vtkProp** propList, int& listLength, int& initialized)
{
  if (listLength != this->PreviousPropsCount)
  {
    // Resizing the culler buffers
    ResizeCuller(propList, listLength);
    // Updates the map that maps the actor to the position of their bounds in the GPU bounds buffer
    UpdateActorToIndex(propList, listLength);
    // Recreating the observers
    AddBoundsObservers(propList, listLength);
    // Adding the observer to the camera if not already added
    AddCameraObserver(ren);
  }

  // Makes sure the bounds of the props are up-to-date
  TriggerBoundsRecomputation(propList, listLength);
  // Reuploads the bounds if some bounds actually need reupload
  UpdateBoundsBuffer(propList, listLength);
  // Reuploads the camera data if the camera was modified since last time
  UpdateCamera(ren);

  this->PreviousPropsCount = listLength;

  // Zeroing the counter of objects that passed the culling test
  std::vector<int> zero = { 0 };
  this->FrustumCullingPass->UpdateBufferData(OutputObjectCountBufferIndex, zero);
  this->FrustumCullingPass->Dispatch();
  this->FrustumCullingPass->ReadBufferFromGPU(this->OutputObjectCountBufferIndex,
    vtkWebGPUComputeFrustumCuller::OutputObjectCountMapCallback, &listLength);

  vtkWebGPUComputeFrustumCuller::OutputIndicesCallbackData callbackData;
  callbackData.indicesCount = &listLength;
  callbackData.propList = propList;
  callbackData.scratchList = &this->CallbackScratchList;
  this->FrustumCullingPass->ReadBufferFromGPU(this->OutputIndicesBufferIndex,
    vtkWebGPUComputeFrustumCuller::OutputObjectIndicesMapCallback, &callbackData);
  this->Pipeline->Update();

  // The allocated time of the props has now been initialized
  initialized = true;

  // We need to return the total time allocated to all the props.
  // Because we're either completely culling a prop or not, each not culled prop has an allocated
  // time of 1 Thus, the total time allocated is just the number of props that were not culled
  return listLength;
}

//------------------------------------------------------------------------------
void vtkWebGPUComputeFrustumCuller::ResizeCuller(vtkProp** propList, int newPropsCount)
{
  // Recomputing the number of workgroups needed to cover the new number of props in the compute
  // shader
  int groupsX = std::ceil(newPropsCount / 32.0f);
  this->FrustumCullingPass->SetWorkgroups(groupsX, 1, 1);

  ResizeBoundsBuffer(propList, newPropsCount);
  ResizeOutputIndicesBuffer(propList, newPropsCount);
  ResizeScratchList(propList, newPropsCount);
}

//------------------------------------------------------------------------------
void vtkWebGPUComputeFrustumCuller::CreateInputBoundsBuffer(vtkProp** propList, int nbProps)
{
  // Gathering the bounds of the actors to upload to the GPU buffer
  std::vector<float> propsBounds(6 * nbProps);

  for (int i = 0; i < nbProps; i++)
  {
    double* bounds = propList[i]->GetBounds();

    for (int j = 0; j < 6; j++)
    {
      propsBounds[i * 6 + j] = static_cast<float>(bounds[j]);
    }
  }

  vtkNew<vtkWebGPUComputeBuffer> inputBoundsBuffer;
  inputBoundsBuffer->SetGroup(0);
  inputBoundsBuffer->SetBinding(0);
  inputBoundsBuffer->SetLabel("Input bounds buffer");
  inputBoundsBuffer->SetMode(vtkWebGPUComputeBuffer::BufferMode::READ_ONLY_COMPUTE_STORAGE);
  inputBoundsBuffer->SetData(propsBounds);

  this->InputBoundsBufferIndex = this->FrustumCullingPass->AddBuffer(inputBoundsBuffer);
}

//------------------------------------------------------------------------------
void vtkWebGPUComputeFrustumCuller::CreateOutputIndicesBuffer(int nbProps)
{
  vtkNew<vtkWebGPUComputeBuffer> outputIndicesBuffer;
  outputIndicesBuffer->SetGroup(0);
  outputIndicesBuffer->SetBinding(1);
  outputIndicesBuffer->SetLabel("Frustum culler output indices buffer");
  outputIndicesBuffer->SetByteSize(nbProps * sizeof(unsigned int));
  outputIndicesBuffer->SetMode(vtkWebGPUComputeBuffer::BufferMode::READ_WRITE_MAP_COMPUTE_STORAGE);

  this->OutputIndicesBufferIndex = this->FrustumCullingPass->AddBuffer(outputIndicesBuffer);
}

void vtkWebGPUComputeFrustumCuller::CreateViewProjMatrixBuffer(const std::vector<float>& vpMat)
{
  vtkNew<vtkWebGPUComputeBuffer> viewProjMatBuffer;
  viewProjMatBuffer->SetGroup(0);
  viewProjMatBuffer->SetBinding(3);
  viewProjMatBuffer->SetMode(vtkWebGPUComputeBuffer::UNIFORM_BUFFER);
  viewProjMatBuffer->SetData(vpMat);
  viewProjMatBuffer->SetLabel("Camera view-projection matrix uniform buffer");

  this->CameraViewProjMatrixBufferIndex = this->FrustumCullingPass->AddBuffer(viewProjMatBuffer);
}

//------------------------------------------------------------------------------
void vtkWebGPUComputeFrustumCuller::ResizeBoundsBuffer(vtkProp** propList, int newPropsCount)
{
  if (this->InputBoundsBufferIndex != -1)
  {
    // If the buffer already exists, resizing it

    // WebGPU doesn't support double precision floating point numbers (as of may 2024). We're thus
    // resizing the buffer with floats size because the compute shader is going to use floats, not
    // doubles which would be the matching type for VTK's bounds.
    this->FrustumCullingPass->ResizeBuffer(
      this->InputBoundsBufferIndex, newPropsCount * 6 * sizeof(float));
  }
  else
  {
    // If the buffer doesn't already exist, creating it
    CreateInputBoundsBuffer(propList, newPropsCount);
  }
}

//------------------------------------------------------------------------------
void vtkWebGPUComputeFrustumCuller::ResizeOutputIndicesBuffer(vtkProp** propList, int newPropsCount)
{
  if (this->OutputIndicesBufferIndex != -1)
  {
    // If the buffer already exists, resizing it

    // WebGPU doesn't support double precision floating point numbers (as of may 2024). We're thus
    // resizing the buffer with floats size because the compute shader is going to use floats, not
    // doubles which would be the matching type for VTK's bounds.
    this->FrustumCullingPass->ResizeBuffer(
      this->OutputIndicesBufferIndex, newPropsCount * sizeof(int));
  }
  else
  {
    // If the buffer doesn't already exist, creating it
    CreateOutputIndicesBuffer(newPropsCount);
  }
}

//------------------------------------------------------------------------------
void vtkWebGPUComputeFrustumCuller::ResizeScratchList(vtkProp** propList, int listLength)
{
  this->CallbackScratchList.resize(listLength);
  for (int i = 0; i < listLength; i++)
    this->CallbackScratchList[i] = propList[i];
}

//------------------------------------------------------------------------------
void vtkWebGPUComputeFrustumCuller::UpdateActorToIndex(vtkProp** propList, int newPropsCount)
{
  this->ActorToIndex.clear();
  for (int i = 0; i < newPropsCount; i++)
  {
    vtkActor* actor = vtkActor::SafeDownCast(propList[i]);
    this->ActorToIndex.insert(std::make_pair(actor, i));
  }
}

//------------------------------------------------------------------------------
void vtkWebGPUComputeFrustumCuller::AddBoundsObservers(vtkProp** propList, int listLength)
{
  // Adding observers to the actors that don't have one already
  for (int i = 0; i < listLength; i++)
  {
    vtkProp* prop = propList[i];
    if (!prop->HasObserver(
          vtkCommand::ActorBoundsRecomputedEvent, this->ActorBoundsRecomputedCallback))
    {
      propList[i]->AddObserver(
        vtkCommand::ActorBoundsRecomputedEvent, this->ActorBoundsRecomputedCallback);
    }
  }
}

//------------------------------------------------------------------------------
void vtkWebGPUComputeFrustumCuller::AddCameraObserver(vtkRenderer* ren)
{
  vtkCamera* camera = ren->GetActiveCamera();
  if (!camera->HasObserver(vtkCommand::ModifiedEvent))
  {
    camera->AddObserver(vtkCommand::ModifiedEvent, this->CameraModifiedCallback);
  }
}

//------------------------------------------------------------------------------
void vtkWebGPUComputeFrustumCuller::TriggerBoundsRecomputation(vtkProp** propList, int listLength)
{
  for (int i = 0; i < listLength; i++)
  {
    vtkProp* prop = propList[i];
    prop->GetBounds();
  }
}

//------------------------------------------------------------------------------
void vtkWebGPUComputeFrustumCuller::UpdateBoundsBuffer(vtkProp** propList, int listLength)
{
  if (this->ActorsToReupload.size() == 0)
  {
    // No actors to reupload
    return;
  }
  else if (this->ActorsToReupload.size() > this->ReuploadBoundsBufferThreshold() ||
    this->PreviousPropsCount != listLength)
  {
    // We're only reuploading the whole buffer if there are too many actors bounds to reupload
    // (above the threshold).
    //
    // Also, if the previous size of the buffers is different than the current length of the prop
    // list, this means that the bounds buffer has been resized which means we'll have to reupload
    // everything because resizing a buffer (essentially recreating it) invalidates all the data it
    // contains since it could have been recreated at another memmory location.

    std::vector<float> allBounds(listLength * 6);
    for (int i = 0; i < listLength; i++)
    {
      vtkProp* prop = propList[i];
      double* propBounds = prop->GetBounds();

      for (int j = 0; j < 6; j++)
      {
        allBounds[i * 6 + j] = static_cast<float>(propBounds[j]);
      }
    }

    this->FrustumCullingPass->UpdateBufferData(this->InputBoundsBufferIndex, allBounds);

    vtkDebugWithObjectMacro(this, "All actor bounds reuploaded.");
  }
  else
  {
    // Below the threshold, we're only going to reupload the bounds one by one to avoid reuploading
    // the whole buffer.
    for (vtkActor* actor : this->ActorsToReupload)
    {
      auto mapFind = this->ActorToIndex.find(actor);
      if (mapFind == this->ActorToIndex.end())
      {
        vtkLog(ERROR,
          "Could not reupload the bounds of an actor that changed. The actor could not be found in "
          "the ActorToIndex map.");

        continue;
      }

      // Getting the bounds and converting them to float for WebGPU
      std::vector<float> boundsFloat(6);
      double bounds[6];
      actor->GetBounds(bounds);
      for (int i = 0; i < 6; i++)
      {
        boundsFloat[i] = static_cast<float>(bounds[i]);
      }

      int actorIndex = mapFind->second;
      this->FrustumCullingPass->UpdateBufferData(
        this->InputBoundsBufferIndex, actorIndex * 6 * sizeof(float), boundsFloat);
    }

    vtkDebugWithObjectMacro(this, << this->ActorsToReupload.size() << " actor bounds reuploaded.");
  }

  this->ActorsToReupload.clear();
}

//------------------------------------------------------------------------------
void vtkWebGPUComputeFrustumCuller::UpdateCamera(vtkRenderer* ren)
{
  if (!this->CameraDirty)
  {
    // No reupload needed
    return;
  }

  vtkCamera* camera = ren->GetActiveCamera();
  // Getting the view-projection matrix
  vtkMatrix4x4* viewMatrix = camera->GetModelViewTransformMatrix();

  // We're using [0, 1] for znear and zfar here to align with WebGPU convention but [-1, 1] as in
  // OpenGL would have worked too since we're not using the graphics pipeline (compute shader only)
  // that actually expects [0, 1]
  vtkMatrix4x4* projectionMatrix =
    camera->GetProjectionTransformMatrix(ren->GetTiledAspectRatio(), 0, 1);
  vtkNew<vtkMatrix4x4> viewProj;
  vtkMatrix4x4::Multiply4x4(projectionMatrix, viewMatrix, viewProj);
  // WebGPU uses column major matrices but VTK is row major
  viewProj->Transpose();

  // Getting the matrix data as floats
  std::vector<float> matrixData(16);
  for (int i = 0; i < 16; i++)
  {
    matrixData[i] = static_cast<float>(viewProj->GetData()[i]);
  }

  // Creating / updating the view-proj matrix uniform buffer
  if (this->CameraViewProjMatrixBufferIndex == -1)
  {
    this->CreateViewProjMatrixBuffer(matrixData);
  }
  else
  {
    this->FrustumCullingPass->UpdateBufferData(this->CameraViewProjMatrixBufferIndex, matrixData);
  }

  this->CameraDirty = false;
}

//------------------------------------------------------------------------------
int vtkWebGPUComputeFrustumCuller::ReuploadBoundsBufferThreshold()
{
  // Simply returns the constant for now but this function allows for a smarter implmentation in the
  // future
  return vtkWebGPUComputeFrustumCuller::WHOLE_BUFFER_REUPLOAD_THRESHOLD;
}

//------------------------------------------------------------------------------
void vtkWebGPUComputeFrustumCuller::BoundsRecomputedCallbackFunction(
  vtkObject*, unsigned long, void* clientdata, void* calldata)
{
  vtkActor* actor = static_cast<vtkActor*>(calldata);
  vtkWebGPUComputeFrustumCuller* culler = static_cast<vtkWebGPUComputeFrustumCuller*>(clientdata);

  culler->ActorsToReupload.push_back(actor);
}

//------------------------------------------------------------------------------
void vtkWebGPUComputeFrustumCuller::CameraModifiedCallbackFunction(
  vtkObject*, unsigned long, void* clientdata, void* calldata)
{
  vtkWebGPUComputeFrustumCuller* culler = static_cast<vtkWebGPUComputeFrustumCuller*>(clientdata);
  culler->CameraDirty = true;
}

//------------------------------------------------------------------------------
void vtkWebGPUComputeFrustumCuller::OutputObjectCountMapCallback(
  const void* mappedData, void* userdata)
{
  int* listLength = static_cast<int*>(userdata);

  const int* mappedIntData = static_cast<const int*>(mappedData);
  *listLength = *mappedIntData;
}

//------------------------------------------------------------------------------
void vtkWebGPUComputeFrustumCuller::OutputObjectIndicesMapCallback(
  const void* mappedData, void* userdata)
{
  vtkWebGPUComputeFrustumCuller::OutputIndicesCallbackData* data =
    reinterpret_cast<vtkWebGPUComputeFrustumCuller::OutputIndicesCallbackData*>(userdata);

  const int* nonCulledIndices = reinterpret_cast<const int*>(mappedData);
  for (int i = 0; i < *data->indicesCount; i++)
  {
    data->propList[i] = (*data->scratchList)[nonCulledIndices[i]];
  }
}

VTK_ABI_NAMESPACE_END
