//=========================================================================
//  Copyright (c) Kitware, Inc.
//  All rights reserved.
//  See LICENSE.txt 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 "vtkImageNarrowBand.h"

#include "vtkCellLocator.h"
#include "vtkCellTypes.h"
#include "vtkDoubleArray.h"
#include "vtkGeometryFilter.h"
#include "vtkImageData.h"
#include "vtkImplicitPolyDataDistance.h"
#include "vtkInformation.h"
#include "vtkInformationVector.h"
#include "vtkMultiThreader.h"
#include "vtkPointData.h"
#include "vtkUnstructuredGrid.h"

vtkStandardNewMacro(vtkImageNarrowBand);

namespace
{

VTK_THREAD_RETURN_TYPE threadedComputePolyData(void* /*args*/);
VTK_THREAD_RETURN_TYPE threadedComputeUnstructuredGrid(void* /*args*/);

struct ThreadArgs
{
public:
  ThreadArgs(vtkImageNarrowBand* self,
    vtkUnstructuredGrid* volumeMesh,
    vtkPolyData* surfMesh,
    vtkImageData* outputImage,
    vtkDoubleArray* distanceArray)
    : self(self)
    , volumeMesh(volumeMesh)
    , surfMesh(surfMesh)
    , outputImage(outputImage)
    , distanceArray(distanceArray)
  {
  }

  friend VTK_THREAD_RETURN_TYPE threadedComputePolyData(void* /*args*/);
  friend VTK_THREAD_RETURN_TYPE threadedComputeUnstructuredGrid(void* /*args*/);

private:
  vtkImageNarrowBand* self = nullptr;
  vtkUnstructuredGrid* volumeMesh = nullptr;
  vtkPolyData* surfMesh = nullptr;
  vtkImageData* outputImage = nullptr;
  vtkDoubleArray* distanceArray = nullptr;
};

VTK_THREAD_RETURN_TYPE threadedComputePolyData(void* args)
{
  // Unpack thread args
  auto* threadInfo = static_cast<vtkMultiThreader::ThreadInfo*>(args);
  auto* threadArgs = static_cast<ThreadArgs*>(threadInfo->UserData);

  // Have to make copy for thread (not particularly efficient, cell traversals fault)
  vtkNew<vtkPolyData> inputSurfMeshCopy;
  inputSurfMeshCopy->DeepCopy(threadArgs->surfMesh);
  vtkImageData* outputImage = threadArgs->outputImage;
  vtkDoubleArray* distancesArray = threadArgs->distanceArray;
  vtkImageNarrowBand* self = threadArgs->self;
  int stride = threadInfo->NumberOfThreads;
  int threadId = threadInfo->ThreadID;

  {
    // Setup the implicit function
    vtkNew<vtkImplicitPolyDataDistance> distFunc;
    distFunc->SetInput(inputSurfMeshCopy);

    int* dim = outputImage->GetDimensions();
    double* distancesPtr = distancesArray->GetPointer(0);

    const int numPx = dim[0] * dim[1] * dim[2];
    for (int i = threadId; i < numPx; i += stride)
    {
      int x = i % dim[0];
      int y = (i / dim[0]) % dim[1];
      int z = i / (dim[0] * dim[1]);
      double pos[3];
      outputImage->TransformIndexToPhysicalPoint(x, y, z, pos);

      distancesPtr[i] = distFunc->FunctionValue(pos);

      // Only report progress on one thread
      if (i % 10000 == 0 && threadId == 0)
      {
        self->UpdateProgress(static_cast<double>(i) / numPx);
      }
    }
  }

  return VTK_THREAD_RETURN_VALUE;
}

VTK_THREAD_RETURN_TYPE threadedComputeUnstructuredGrid(void* args)
{
  // Unpack thread args
  auto* threadInfo = static_cast<vtkMultiThreader::ThreadInfo*>(args);
  auto* threadArgs = static_cast<ThreadArgs*>(threadInfo->UserData);

  vtkNew<vtkUnstructuredGrid> inputVolumeMeshCopy;
  inputVolumeMeshCopy->DeepCopy(threadArgs->volumeMesh);
  vtkNew<vtkPolyData> inputSurfMeshCopy;
  inputSurfMeshCopy->DeepCopy(threadArgs->surfMesh);
  vtkImageData* outputImage = threadArgs->outputImage;
  vtkDoubleArray* distancesArray = threadArgs->distanceArray;
  vtkImageNarrowBand* self = threadArgs->self;
  int stride = threadInfo->NumberOfThreads;
  int threadId = threadInfo->ThreadID;

  {
    // Setup a cell locator to compute which cell we are in, and interpolation weights
    vtkNew<vtkCellLocator> locator;
    locator->SetDataSet(inputVolumeMeshCopy);
    locator->BuildLocator();

    // Setup a distance function for when not in a cell
    vtkNew<vtkImplicitPolyDataDistance> distFunc;
    distFunc->SetInput(inputSurfMeshCopy);

    int* dim = outputImage->GetDimensions();
    double* distancesPtr = distancesArray->GetPointer(0);

    auto* weights = new double[inputVolumeMeshCopy->GetMaxCellSize()];
    double closestDist;
    int subId;
    double closestPt[3];
    double pcoords[3];

    const int numPx = dim[0] * dim[1] * dim[2];
    for (int i = threadId; i < numPx; i += stride)
    {
      int x = i % dim[0];
      int y = (i / dim[0]) % dim[1];
      int z = i / (dim[0] * dim[1]);
      double pos[3];
      outputImage->TransformIndexToPhysicalPoint(x, y, z, pos);

      // Find the cell the point is in
      vtkIdType cellId = locator->FindCell(pos);
      // If not in a cell
      if (cellId == -1)
      {
        // Compute implicit poly data distance
        distancesPtr[i] = distFunc->FunctionValue(pos);
      }
      // If within a cell, interpolate
      else
      {
        distancesPtr[i] = 0.0;
        vtkCell* cell = inputVolumeMeshCopy->GetCell(cellId);
        cell->EvaluatePosition(pos, closestPt, subId, pcoords, closestDist, weights);
      }

      // Only report progress on one thread
      if (i % 10000 == 0 && threadId == 0)
      {
        self->UpdateProgress(static_cast<double>(i) / numPx);
      }
    }
  }

  return VTK_THREAD_RETURN_VALUE;
}

}

vtkImageNarrowBand::vtkImageNarrowBand()
{
  this->SetNumberOfInputPorts(2);
}

void vtkImageNarrowBand::PrintSelf(std::ostream& os, vtkIndent indent)
{
  this->Superclass::PrintSelf(os, indent);
  os << indent << "UseBinary: " << this->UseBinary << "\n";
  os << indent << "Bandwidth: " << this->BandWidth << "\n";
}

void vtkImageNarrowBand::SetReferenceImage(const vtkSmartPointer<vtkImageData>& refImage)
{
  this->SetInputData(1, refImage);
}

int vtkImageNarrowBand::FillInputPortInformation(int port, vtkInformation* info)
{
  if (port == 0)
  {
    info->Set(vtkAlgorithm::INPUT_REQUIRED_DATA_TYPE(), "vtkDataSet");
  }
  else if (port == 1)
  {
    info->Set(vtkAlgorithm::INPUT_REQUIRED_DATA_TYPE(), "vtkImageData");
  }
  return 1;
}

int vtkImageNarrowBand::RequestData(vtkInformation* request,
  vtkInformationVector** inputVec,
  vtkInformationVector* outputVec)
{
  (void)request;

  // Get the input geometry
  vtkInformation* inInfo1 = inputVec[0]->GetInformationObject(0);
  vtkPointSet* inputPointSet =
    vtkPointSet::SafeDownCast(inInfo1->Get(vtkDataObject::DATA_OBJECT()));
  // Get the input image
  vtkInformation* inInfo2 = inputVec[1]->GetInformationObject(0);
  vtkImageData* inputRefImage =
    vtkImageData::SafeDownCast(inInfo2->Get(vtkDataObject::DATA_OBJECT()));
  // Get the output image
  vtkInformation* outInfo = outputVec->GetInformationObject(0);
  vtkImageData* output = vtkImageData::SafeDownCast(outInfo->Get(vtkDataObject::DATA_OBJECT()));
  output->DeepCopy(inputRefImage);

  // Compute the bounding region within the image
  int* dims = inputRefImage->GetDimensions();

  vtkNew<vtkDoubleArray> distanceArray;
  distanceArray->SetName("Distances");
  distanceArray->SetNumberOfComponents(1);
  distanceArray->SetNumberOfValues(dims[0] * dims[1] * dims[2]);
  distanceArray->Fill(
    std::numeric_limits<double>::max()); // Cells outside of the bounds will have max assigned

  {
    // First try to acquire a vtkPolyData (either directly or through conversion of a vtkUnstructuredGrid)
    vtkSmartPointer<vtkPolyData> inputPolyData = vtkPolyData::SafeDownCast(inputPointSet);
    if (inputPolyData == nullptr)
    {
      // If we're instead dealing with an unstructured grid, we may still get a poly data
      vtkSmartPointer<vtkUnstructuredGrid> inputUnstructuredGrid =
        vtkUnstructuredGrid::SafeDownCast(inputPointSet);
      if (inputUnstructuredGrid != nullptr)
      {
        // Try to convert to polydata surface
        bool is2d = true;
        auto* cellTypes = inputUnstructuredGrid->GetDistinctCellTypesArray();
        for (vtkIdType i = 0; i < cellTypes->GetNumberOfTuples(); i++)
        {
          const int cellType = cellTypes->GetValue(i);
          if (cellType != VTK_QUAD && cellType != VTK_TRIANGLE && cellType != VTK_LINE &&
            cellType != VTK_VERTEX)
          {
            is2d = false;
            break;
          }
        }
        if (is2d)
        {
          vtkNew<vtkGeometryFilter> toPolyData;
          toPolyData->SetInputData(inputUnstructuredGrid);
          toPolyData->Update();
          inputPolyData = toPolyData->GetOutput();
        }
      }
    }

    // If we were able to acquire a vtkPolyData then we will compute distances to the
    // via vtkPolyDataImplicitDistance, which will give better results as it gives min
    // distance to points AND cells
    vtkNew<vtkMultiThreader> threader;
    if (inputPolyData != nullptr)
    {
      ThreadArgs args(this, nullptr, inputPolyData, output, distanceArray);
      threader->SetSingleMethod(threadedComputePolyData, &args);
      threader->SingleMethodExecute();
    }
    // Otherwise we are working with 3d cells
    else
    {
      vtkSmartPointer<vtkUnstructuredGrid> inputUnstructuredGrid =
        vtkUnstructuredGrid::SafeDownCast(inputPointSet);
      if (inputUnstructuredGrid != nullptr)
      {
        // Extract the surface
        vtkNew<vtkGeometryFilter> extractSurf;
        extractSurf->SetInputData(inputUnstructuredGrid);
        extractSurf->Update();

        ThreadArgs args(
          this, inputUnstructuredGrid, extractSurf->GetOutput(), output, distanceArray);
        threader->SetSingleMethod(threadedComputeUnstructuredGrid, &args);
        threader->SingleMethodExecute();
      }
      else
      {
        vtkWarningMacro(<< inputPointSet->GetClassName() << " geometry type not supported\n");
        return 1;
      }
    }
  }

  if (UseBinary)
  {
    vtkNew<vtkUnsignedCharArray> binaryArray;
    binaryArray->SetName("Band");
    binaryArray->SetNumberOfComponents(1);
    binaryArray->SetNumberOfValues(distanceArray->GetNumberOfValues());
    for (vtkIdType i = 0; i < binaryArray->GetNumberOfValues(); i++)
    {
      binaryArray->SetValue(i, std::abs(distanceArray->GetValue(i)) < BandWidth ? 1 : 0);
    }
    output->GetPointData()->AddArray(binaryArray);
    output->GetPointData()->SetActiveScalars("Band");
  }
  else
  {
    output->GetPointData()->AddArray(distanceArray);
    output->GetPointData()->SetActiveScalars("Distances");
  }

  return 1;
}
