// SPDX-FileCopyrightText: Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
// SPDX-License-Identifier: BSD-3-Clause
#include "vtkImageToAMR.h"

#include "vtkAMRBox.h"
#include "vtkAMRUtilities.h"
#include "vtkCellData.h"
#include "vtkCompositeDataPipeline.h"
#include "vtkImageData.h"
#include "vtkInformation.h"
#include "vtkInformationVector.h"
#include "vtkNew.h"
#include "vtkObjectFactory.h"
#include "vtkOverlappingAMR.h"
#include "vtkPointData.h"
#include "vtkSmartPointer.h"
#include "vtkTuple.h"
#include "vtkUniformGrid.h"

#include "vtkOverlappingAMRMetaData.h"
#include <algorithm>
#include <vector>
VTK_ABI_NAMESPACE_BEGIN
namespace
{
// Split one box to eight
int SplitXYZ(vtkAMRBox inBox, int refinementRatio, std::vector<vtkAMRBox>& out)
{
  inBox.Refine(refinementRatio);
  const int* lo = inBox.GetLoCorner();
  const int* hi = inBox.GetHiCorner();

  // the cartesian product A[0][0..n[0]] X A[1][0..n[1] X A[2][0..n[2]] is the refined grid
  int A[3][3], n[3];
  for (int d = 0; d < 3; d++)
  {
    A[d][0] = lo[d] - 1;
    A[d][2] = hi[d];
    if (inBox.EmptyDimension(d))
    {
      n[d] = 1;
      A[d][1] = hi[d];
    }
    else
    {
      n[d] = 2;
      A[d][1] = (lo[d] + hi[d]) / 2;
    }
  }

  // create the refined boxes and push them to the output stack
  int numOut(0);
  for (int i = 0; i < n[0]; i++)
  {
    for (int j = 0; j < n[1]; j++)
    {
      for (int k = 0; k < n[2]; k++)
      {
        vtkAMRBox box;
        box.SetDimensions(
          A[0][i] + 1, A[1][j] + 1, A[2][k] + 1, A[0][i + 1], A[1][j + 1], A[2][k + 1]);
        out.push_back(box);
        numOut++;
      }
    }
  }
  return numOut;
}

int ComputeTreeHeight(int maxNumNodes, int degree)
{
  if (maxNumNodes <= 0)
  {
    return 0;
  }
  // could have used a formula, but this is more clear
  int height = 1;
  int numNodes = 1;
  while (numNodes <= maxNumNodes)
  {
    numNodes = numNodes + degree * numNodes;
    height++;
  }
  height--;
  return height;
}

// split the blocks into a tree that starts out as a single stem
// than turn a full tree. This shape is designed so that numLevels and maxNumBlocks
// constraint can be satisfied
void Split(const vtkAMRBox& rootBox, int numLevels, int refinementRatio, int maxNumBlocks,
  std::vector<std::vector<vtkAMRBox>>& out)
{
  out.clear();
  out.resize(1);
  out.back().push_back(rootBox);
  maxNumBlocks--;

  int treeDegree = rootBox.ComputeDimension() * 2;
  int numTreeLevels =
    std::min(numLevels, ComputeTreeHeight(maxNumBlocks - (numLevels - 1), treeDegree)) -
    1; // minus one because root already has one
  int level = 1;
  for (; level < numLevels - numTreeLevels; level++)
  {
    out.emplace_back();
    const std::vector<vtkAMRBox>& parentBoxes = out[level - 1];
    std::vector<vtkAMRBox>& childBoxes = out[level];
    vtkAMRBox child = parentBoxes.back();
    child.Refine(refinementRatio);
    childBoxes.push_back(child);
  }

  for (; level < numLevels; level++)
  {
    out.emplace_back();
    const std::vector<vtkAMRBox>& parentBoxes = out[level - 1];
    std::vector<vtkAMRBox>& childBoxes = out[level];
    for (size_t i = 0; i < parentBoxes.size(); i++)
    {
      const vtkAMRBox& parent = parentBoxes[i];
      SplitXYZ(parent, refinementRatio, childBoxes);
    }
  }
}

// create a grid by sampling from input using the indices in box
vtkUniformGrid* ConstructGrid(
  vtkImageData* input, const vtkAMRBox& box, int coarsenRatio, double* origin, double* spacing)
{
  int numPoints[3];
  box.GetNumberOfNodes(numPoints);

  vtkUniformGrid* grid = vtkUniformGrid::New();
  grid->Initialize();
  grid->SetDimensions(numPoints);
  grid->SetSpacing(spacing);
  grid->SetOrigin(origin);

  vtkPointData *inPD = input->GetPointData(), *outPD = grid->GetPointData();
  vtkCellData *inCD = input->GetCellData(), *outCD = grid->GetCellData();

  outPD->CopyAllocate(inPD, grid->GetNumberOfPoints());
  outCD->CopyAllocate(inCD, grid->GetNumberOfCells());

  vtkAMRBox box0(box);
  box0.Refine(coarsenRatio); // refine it to the image data level

  int extents[6];
  input->GetExtent(extents);
  int imLo[3] = { extents[0], extents[2], extents[4] };
  const int* lo = box.GetLoCorner();

  for (int iz = 0; iz < numPoints[2]; iz++)
  {
    for (int iy = 0; iy < numPoints[1]; iy++)
    {
      for (int ix = 0; ix < numPoints[0]; ix++)
      {
        int ijkDst[3] = { ix, iy, iz };
        vtkIdType idDst = grid->ComputePointId(ijkDst);

        int ijkSrc[3] = { (lo[0] + ix) * coarsenRatio + imLo[0],
          (lo[1] + iy) * coarsenRatio + imLo[1], (lo[2] + iz) * coarsenRatio + imLo[2] };

        vtkIdType idSrc = input->ComputePointId(ijkSrc);

        outPD->CopyData(inPD, idSrc, idDst);
      }
    }
  }

  int numCells[3];
  for (int d = 0; d < 3; d++)
  {
    numCells[d] = std::max(numPoints[d] - 1, 1);
  }

  for (int iz = 0; iz < numCells[2]; iz++)
  {
    for (int iy = 0; iy < numCells[1]; iy++)
    {
      for (int ix = 0; ix < numCells[0]; ix++)
      {
        int ijkDst[3] = { ix, iy, iz };
        vtkIdType idDst = grid->ComputeCellId(ijkDst);

        int ijkSrc[3] = { (lo[0] + ix) * coarsenRatio + imLo[0],
          (lo[1] + iy) * coarsenRatio + imLo[1], (lo[2] + iz) * coarsenRatio + imLo[2] };

        vtkIdType idSrc = input->ComputeCellId(ijkSrc);

        outCD->CopyData(inCD, idSrc, idDst);
      }
    }
  }

  return grid;
}

}

vtkStandardNewMacro(vtkImageToAMR);
//------------------------------------------------------------------------------
vtkImageToAMR::vtkImageToAMR()
{
  this->NumberOfLevels = 2;
  this->RefinementRatio = 2;
  this->MaximumNumberOfBlocks = 100;
}

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

//------------------------------------------------------------------------------
int vtkImageToAMR::FillInputPortInformation(int, vtkInformation* info)
{
  info->Set(vtkAlgorithm::INPUT_REQUIRED_DATA_TYPE(), "vtkImageData");
  return 1;
}

//------------------------------------------------------------------------------
int vtkImageToAMR::RequestData(vtkInformation* vtkNotUsed(request),
  vtkInformationVector** inputVector, vtkInformationVector* outputVector)
{
  vtkInformation* inInfo = inputVector[0]->GetInformationObject(0);

  vtkImageData* input = vtkImageData::GetData(inputVector[0], 0);
  vtkOverlappingAMR* amr = vtkOverlappingAMR::GetData(outputVector);

  if (input->GetDataDimension() < 2)
  {
    vtkErrorMacro("Image dimension must be at least two.");
    return 0;
  }

  int whole_extent[6];
  inInfo->Get(vtkCompositeDataPipeline::WHOLE_EXTENT(), whole_extent);

  int dims[3] = { whole_extent[1] - whole_extent[0] + 1, whole_extent[3] - whole_extent[2] + 1,
    whole_extent[5] - whole_extent[4] + 1 };

  double inputBounds[6];
  input->GetBounds(inputBounds);

  double inputOrigin[3] = { inputBounds[0], inputBounds[2], inputBounds[4] };

  double inputSpacing[3];
  input->GetSpacing(inputSpacing);

  int gridDescription = vtkStructuredData::GetDataDescription(dims);

  // check whether the parameters are valid
  // and compute the base image resolution
  int dims0[3];
  double spacing0[3];
  for (int d = 0; d < 3; d++)
  {
    if (dims[d] <= 1)
    {
      if (dims[d] == 0)
      {
        vtkWarningMacro("Zero dimension? Really?");
      }
      dims0[d] = 1;
      spacing0[d] = 1.0;
    }
    else
    {
      int r = (int)(pow(static_cast<double>(this->RefinementRatio), this->NumberOfLevels - 1));
      if ((dims[d] - 1) % r != 0)
      {
        vtkErrorMacro("Image cannot be refined");
        return 0;
      }
      dims0[d] = 1 + (dims[d] - 1) / r;
      spacing0[d] = r * inputSpacing[d];
    }
  }

  vtkAMRBox rootBox(inputOrigin, dims0, spacing0, inputOrigin, gridDescription);

  std::vector<std::vector<vtkAMRBox>> amrBoxes;
  Split(
    rootBox, this->NumberOfLevels, this->RefinementRatio, this->MaximumNumberOfBlocks, amrBoxes);

  std::vector<unsigned int> blocksPerLevel;
  blocksPerLevel.reserve(amrBoxes.size());
  for (size_t i = 0; i < amrBoxes.size(); i++)
  {
    blocksPerLevel.push_back(static_cast<int>(amrBoxes[i].size()));
  }

  unsigned int numLevels = static_cast<unsigned int>(blocksPerLevel.size());
  bool abort = false;

  amr->Initialize(blocksPerLevel);
  amr->SetOrigin(inputOrigin);
  amr->SetGridDescription(gridDescription);

  double spacingi[3] = { spacing0[0], spacing0[1], spacing0[2] };
  unsigned int checkAbortInterval = std::min(numLevels / 10 + 1, (unsigned int)1000);
  for (unsigned int i = 0; i < numLevels && !abort; i++)
  {
    if (i % checkAbortInterval == 0)
    {
      abort = this->CheckAbort();
    }
    amr->SetSpacing(i, spacingi);
    for (int d = 0; d < 3; d++)
    {
      spacingi[d] /= this->RefinementRatio;
    }
  }

  for (unsigned int level = 0; level < numLevels && !abort; level++)
  {
    if (level % checkAbortInterval == 0)
    {
      abort = this->CheckAbort();
    }
    const std::vector<vtkAMRBox>& boxes = amrBoxes[level];
    for (size_t i = 0; i < boxes.size(); i++)
    {
      amr->SetAMRBox(level, static_cast<unsigned int>(i), boxes[i]);
    }
  }

  for (unsigned int level = 0; level < numLevels && !abort; level++)
  {
    if (level % checkAbortInterval == 0)
    {
      abort = this->CheckAbort();
    }
    double spacing[3];
    amr->GetSpacing(level, spacing);
    int coarsenRatio = (int)pow(static_cast<double>(this->RefinementRatio),
      static_cast<int>(numLevels - 1 - level)); // against the finest level
    for (size_t i = 0; i < amr->GetNumberOfBlocks(level) && !abort; i++)
    {
      const vtkAMRBox& box = amr->GetAMRBox(level, static_cast<unsigned int>(i));
      double origin[3];
      vtkAMRBox::GetBoxOrigin(box, inputOrigin, spacing, origin);
      vtkUniformGrid* grid = ConstructGrid(input, box, coarsenRatio, origin, spacing);
      amr->SetDataSet(level, static_cast<unsigned int>(i), grid);
      grid->Delete();
    }
  }

  // Skipping BlankCells in case amr is empty
  if (!this->CheckAbort())
  {
    vtkAMRUtilities::BlankCells(amr);
  }
  return 1;
}

//------------------------------------------------------------------------------
void vtkImageToAMR::PrintSelf(ostream& os, vtkIndent indent)
{
  this->Superclass::PrintSelf(os, indent);
  os << indent << "NumberOfLevels: " << this->NumberOfLevels << endl;
  os << indent << "RefinementRatio: " << this->RefinementRatio << endl;
}
VTK_ABI_NAMESPACE_END