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

/**
 * This test renders a point cloud with the WebGPU compute API and ensures that the resulting image
 * is correct
 */

#include "vtkActor.h"
#include "vtkCamera.h"
#include "vtkColorSeries.h"
#include "vtkLookupTable.h"
#include "vtkNamedColors.h"
#include "vtkNew.h"
#include "vtkPointData.h"
#include "vtkPolyDataMapper.h"
#include "vtkProperty.h"
#include "vtkRenderWindow.h"
#include "vtkRenderWindowInteractor.h"
#include "vtkRenderer.h"
#include "vtkScalarsToColors.h"
#include "vtkSmartPointer.h"
#include "vtkWebGPUComputePointCloudMapper.h"
#include "vtkWebGPURenderer.h"
#include "vtkXMLPolyDataReader.h"

#include <chrono>
#include <iostream>

//////////////////////////
bool USE_COMPUTE_SHADER = true;
//////////////////////////

std::chrono::time_point<std::chrono::high_resolution_clock> start, stop;

void startFunction(vtkObject* vtkNotUsed(caller), unsigned long vtkNotUsed(eid),
  void* vtkNotUsed(cliendata), void* vtkNotUsed(calldata))
{
  start = std::chrono::high_resolution_clock::now();
}

void endFunction(vtkObject* vtkNotUsed(caller), unsigned long vtkNotUsed(eid),
  void* vtkNotUsed(cliendata), void* vtkNotUsed(calldata))
{
  stop = std::chrono::high_resolution_clock::now();
}

namespace
{
void CreateLookupTableVTKBlue(vtkLookupTable* lut, int size)
{
  vtkNew<vtkColorSeries> myColors;
  myColors->SetColorSchemeByName("VTKBlueColors");

  vtkNew<vtkNamedColors> nc;

  myColors->AddColor(nc->GetColor3ub("alice_blue"));
  myColors->AddColor(nc->GetColor3ub("blue"));
  myColors->AddColor(nc->GetColor3ub("blue_light"));
  myColors->AddColor(nc->GetColor3ub("blue_medium"));
  myColors->AddColor(nc->GetColor3ub("cadet"));
  myColors->AddColor(nc->GetColor3ub("cobalt"));
  myColors->AddColor(nc->GetColor3ub("cornflower"));
  myColors->AddColor(nc->GetColor3ub("cerulean"));
  myColors->AddColor(nc->GetColor3ub("dodger_blue"));
  myColors->AddColor(nc->GetColor3ub("indigo"));
  myColors->AddColor(nc->GetColor3ub("manganese_blue"));
  myColors->AddColor(nc->GetColor3ub("midnight_blue"));
  myColors->AddColor(nc->GetColor3ub("navy"));
  myColors->AddColor(nc->GetColor3ub("peacock"));
  myColors->AddColor(nc->GetColor3ub("powder_blue"));
  myColors->AddColor(nc->GetColor3ub("royal_blue"));
  myColors->AddColor(nc->GetColor3ub("slate_blue"));
  myColors->AddColor(nc->GetColor3ub("slate_blue_dark"));
  myColors->AddColor(nc->GetColor3ub("slate_blue_light"));
  myColors->AddColor(nc->GetColor3ub("slate_blue_medium"));
  myColors->AddColor(nc->GetColor3ub("sky_blue"));
  myColors->AddColor(nc->GetColor3ub("sky_blue_deep"));
  myColors->AddColor(nc->GetColor3ub("sky_blue_light"));
  myColors->AddColor(nc->GetColor3ub("steel_blue"));
  myColors->AddColor(nc->GetColor3ub("steel_blue_light"));
  myColors->AddColor(nc->GetColor3ub("turquoise_blue"));
  myColors->AddColor(nc->GetColor3ub("ultramarine"));

  int numberOfColors = myColors->GetNumberOfColors();
  lut->SetNumberOfTableValues(size == 0 ? numberOfColors : size);
  lut->SetTableRange(0, lut->GetNumberOfTableValues());
  for (int i = 0; i < lut->GetNumberOfTableValues(); ++i)
  {
    vtkColor3ub color = myColors->GetColorRepeating(i);
    lut->SetTableValue(
      i, color.GetRed() / 255., color.GetGreen() / 255., color.GetBlue() / 255., 1.);
  }
}

vtkSmartPointer<vtkActor> CreatePointCube(
  double translationX, double translationY, double translationZ)
{
  vtkNew<vtkPolyData> polydata;

  vtkNew<vtkUnsignedCharArray> colors;
  colors->SetNumberOfComponents(4);

  vtkNew<vtkPoints> points;
  constexpr int sizeX = 40;
  constexpr int sizeY = 40;
  constexpr int sizeZ = 40;
  // 'divider' controls the space between the points. Higher values means points closer together
  constexpr float divider = 20.0f;
  constexpr float maxX = sizeX / divider;
  constexpr float maxY = sizeY / divider;
  constexpr float maxZ = sizeZ / divider;
  for (int i = 0; i < sizeX; i++)
  {
    for (int j = 0; j < sizeY; j++)
    {
      for (int k = 0; k < sizeZ; k++)
      {
        int pointIndex = k + j * sizeZ + i * sizeY * sizeZ;

        points->InsertNextPoint(
          i / divider + translationX, j / divider + translationY, k / divider + translationZ);
        colors->InsertComponent(
          pointIndex, 0, static_cast<unsigned char>(i / divider / maxX * 255.0f));
        colors->InsertComponent(
          pointIndex, 1, static_cast<unsigned char>(j / divider / maxY * 255.0f));
        colors->InsertComponent(
          pointIndex, 2, static_cast<unsigned char>(k / divider / maxZ * 127.0f + 127.0f));
        colors->InsertComponent(pointIndex, 3, 255);
      }
    }
  }

  std::vector<vtkIdType> pointIndices(sizeX * sizeY * sizeZ);
  for (int i = 0; i < sizeX * sizeY * sizeZ; i++)
    pointIndices[i] = i;
  vtkNew<vtkCellArray> pointsCellArray;
  pointsCellArray->InsertNextCell(sizeX * sizeY * sizeZ, pointIndices.data());

  vtkSmartPointer<vtkPolyDataMapper> mapper;
  if (USE_COMPUTE_SHADER)
    mapper = vtkSmartPointer<vtkWebGPUComputePointCloudMapper>::New();
  else
    mapper = vtkSmartPointer<vtkPolyDataMapper>::New();

  polydata->SetPoints(points);
  polydata->SetPolys(pointsCellArray);
  polydata->GetPointData()->SetScalars(colors);

  mapper->SetInputData(polydata);

  vtkSmartPointer<vtkActor> actor = vtkSmartPointer<vtkActor>::New();
  actor->SetMapper(mapper);

  if (!USE_COMPUTE_SHADER)
  {
    actor->GetProperty()->SetRepresentationToPoints();
    actor->GetProperty()->SetColor(1.0, 0.0f, 0.0f);
  }

  return actor;
}

vtkSmartPointer<vtkActor> ReadVTP(const std::string& filepath)
{
  vtkNew<vtkXMLPolyDataReader> reader;
  reader->SetFileName(filepath.c_str());
  reader->Update();

  vtkPolyData* polydata = reader->GetOutput();
  polydata->GetPointData()->SetScalars(polydata->GetPointData()->GetArray("Intensity"));

  vtkSmartPointer<vtkActor> actor = vtkSmartPointer<vtkActor>::New();

  vtkNew<vtkWebGPUComputePointCloudMapper> mapper;
  mapper->SetInputData(reader->GetOutput());

  actor->SetMapper(mapper);

  vtkNew<vtkLookupTable> lookupTable;
  CreateLookupTableVTKBlue(lookupTable, 20);
  lookupTable->Build();
  mapper->SetLookupTable(lookupTable);
  mapper->SetScalarRange(0, 6.6e4);

  return actor;
}
}

//------------------------------------------------------------------------------
int DEBUG_TestComputePointCloudRenderer(int, char*[])
{
  vtkNew<vtkRenderWindow> renWin;
  renWin->SetWindowName(__func__);
  renWin->SetMultiSamples(0);
  renWin->SetSize(512, 512);
  renWin->Initialize();

  vtkSmartPointer<vtkRenderer> renderer = vtkSmartPointer<vtkRenderer>::New();
  renderer->SetBackground(0.2, 0.3, 0.4);
  renWin->AddRenderer(renderer);

  vtkSmartPointer<vtkActor> actor =
    ReadVTP("/home/tom/Desktop/GeometricData/PointCloudDataArray.vtp");
  // vtkSmartPointer<vtkActor> actor = CreatePointCube(0, 0, 0);
  // vtkSmartPointer<vtkActor> actor2 = CreatePointCube(0, 0, -10);

  renderer->ResetCamera(actor->GetBounds());

  renderer->AddActor(actor);
  // renderer->AddActor(actor2);

  vtkNew<vtkCallbackCommand> vtkStartCallbackCommand;
  vtkStartCallbackCommand->SetCallback(startFunction);
  renWin->AddObserver(vtkCommand::EndEvent, vtkStartCallbackCommand);

  vtkNew<vtkCallbackCommand> vtkEndCallbackCommand;
  vtkEndCallbackCommand->SetCallback(endFunction);
  renWin->AddObserver(vtkCommand::EndEvent, vtkEndCallbackCommand);

  vtkNew<vtkRenderWindowInteractor> iren;
  iren->SetRenderWindow(renWin);
  iren->Start();

  return EXIT_SUCCESS;
}
