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

#include "vtkCellArray.h"
#include "vtkExecutive.h"
#include "vtkImageData.h"
#include "vtkInformation.h"
#include "vtkInformationVector.h"
#include "vtkMath.h"
#include "vtkObjectFactory.h"
#include "vtkPointData.h"
#include "vtkPoints.h"
#include "vtkPolyData.h"

#include "vtkDijkstraGraphInternals.h"

#include <algorithm>

VTK_ABI_NAMESPACE_BEGIN
vtkStandardNewMacro(vtkDijkstraImageGeodesicPath);

//------------------------------------------------------------------------------
vtkDijkstraImageGeodesicPath::vtkDijkstraImageGeodesicPath()
{
  this->PixelSize = 1.0;
  this->ImageWeight = 1.0;
  this->EdgeLengthWeight = 0.0;
  this->CurvatureWeight = 0.0;
  this->RebuildStaticCosts = false;
}

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

//------------------------------------------------------------------------------
void vtkDijkstraImageGeodesicPath::SetImageWeight(double w)
{
  w = std::min(std::max(w, 0.0), 1.0);
  if (w != this->ImageWeight)
  {
    this->ImageWeight = w;
    this->RebuildStaticCosts = true;
    this->Modified();
  }
}

//------------------------------------------------------------------------------
void vtkDijkstraImageGeodesicPath::SetEdgeLengthWeight(double w)
{
  w = std::min(std::max(w, 0.0), 1.0);
  if (w != this->EdgeLengthWeight)
  {
    this->EdgeLengthWeight = w;
    this->RebuildStaticCosts = true;
    this->Modified();
  }
}

//------------------------------------------------------------------------------
void vtkDijkstraImageGeodesicPath::SetInputData(vtkDataObject* input)
{
  vtkImageData* image = vtkImageData::SafeDownCast(input);
  if (!image)
  {
    return;
  }

  int* dimensions = image->GetDimensions();
  int u[3];
  int n = 0;
  for (int i = 0; i < 3; ++i)
  {
    if (dimensions[i] != 1)
    {
      u[n++] = i;
    }
  }
  if (n != 2)
  {
    vtkErrorMacro(<< "Input cost image must be 2D: input dimensions " << dimensions[0] << ","
                  << dimensions[1] << "," << dimensions[2]);
    return;
  }

  double* spacing = image->GetSpacing();
  this->PixelSize = sqrt(spacing[u[0]] * spacing[u[0]] + spacing[u[1]] * spacing[u[1]]);
  this->Superclass::SetInputData(image);
}

//------------------------------------------------------------------------------
vtkImageData* vtkDijkstraImageGeodesicPath::GetImageDataInput()
{
  if (this->GetNumberOfInputConnections(0) < 1)
  {
    return nullptr;
  }
  return vtkImageData::SafeDownCast(this->GetExecutive()->GetInputData(0, 0));
}

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

  return 0;
}

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

  vtkImageData* image = vtkImageData::SafeDownCast(costInfo->Get(vtkDataObject::DATA_OBJECT()));

  if (!image)
  {
    return 0;
  }

  vtkPolyData* output = vtkPolyData::SafeDownCast(outInfo->Get(vtkDataObject::DATA_OBJECT()));

  if (!output)
  {
    return 0;
  }

  if (this->AdjacencyBuildTime.GetMTime() < image->GetMTime())
  {
    this->Initialize(image);
  }
  else
  {
    // if the filter's static cost weights change, then update them
    if (this->RebuildStaticCosts)
    {
      this->UpdateStaticCosts(image);
    }
    this->Reset();
  }

  this->ShortestPath(image, this->StartVertex, this->EndVertex);
  this->TraceShortestPath(image, output, this->StartVertex, this->EndVertex);

  return 1;
}

//------------------------------------------------------------------------------
double vtkDijkstraImageGeodesicPath::CalculateStaticEdgeCost(
  vtkDataSet* inData, vtkIdType u, vtkIdType v)
{
  vtkImageData* image = vtkImageData::SafeDownCast(inData);

  double p1[3];
  image->GetPoint(u, p1);
  double p2[3];
  image->GetPoint(v, p2);

  double pcoords[3];
  int ijk1[3];
  int ijk2[3];
  image->ComputeStructuredCoordinates(p1, ijk1, pcoords);
  image->ComputeStructuredCoordinates(p2, ijk2, pcoords);

  double cost = this->ImageWeight *
    (image->GetScalarComponentAsDouble(ijk1[0], ijk1[1], ijk1[2], 0) +
      image->GetScalarComponentAsDouble(ijk1[0], ijk1[1], ijk1[2], 0));

  if (this->EdgeLengthWeight != 0.0)
  {
    cost +=
      this->EdgeLengthWeight * (sqrt(vtkMath::Distance2BetweenPoints(p1, p2)) / this->PixelSize);
  }

  return cost;
}

//------------------------------------------------------------------------------
double vtkDijkstraImageGeodesicPath::CalculateDynamicEdgeCost(
  vtkDataSet* inData, vtkIdType u, vtkIdType v)
{
  double cost = 0.0;

  if (this->CurvatureWeight != 0.0)
  {
    int t = this->Internals->Predecessors[u];
    if (t != -1)
    {
      vtkImageData* image = vtkImageData::SafeDownCast(inData);

      double p0[3];
      image->GetPoint(t, p0);
      double p1[3];
      image->GetPoint(u, p1);
      double p2[3];
      image->GetPoint(v, p2);

      double p10[3] = { p1[0] - p0[0], p1[1] - p0[1], p1[2] - p0[2] };
      double p21[3] = { p2[0] - p1[0], p2[1] - p1[1], p2[2] - p1[2] };

      vtkMath::Normalize(p10);
      vtkMath::Normalize(p21);

      // the range of dot product of two unit vectors is [-1, 1] so normalize
      // the maximum curvature from 2 to 1
      cost = this->CurvatureWeight * (0.5 * fabs(vtkMath::Dot(p10, p21) - 1.0));
    }
  }

  return cost;
}

//------------------------------------------------------------------------------
// This is probably a horribly inefficient way to do it.
void vtkDijkstraImageGeodesicPath::BuildAdjacency(vtkDataSet* inData)
{
  vtkImageData* image = vtkImageData::SafeDownCast(inData);

  vtkIdType ncells = image->GetNumberOfCells();

  // optimized for cell type VTK_PIXEL
  //
  vtkIdList* ptIds = vtkIdList::New();
  vtkIdType uId[6] = { 0, 1, 2, 3, 0, 1 };
  vtkIdType vId[6] = { 1, 2, 3, 0, 2, 3 };
  double cost;

  for (vtkIdType i = 0; i < ncells; ++i)
  {
    image->GetCellPoints(i, ptIds);

    for (int j = 0; j < 6; ++j)
    {
      vtkIdType u = ptIds->GetId(vId[j]);
      vtkIdType v = ptIds->GetId(uId[j]);

      // before insert and calc, check if key map u has key v
      std::map<int, double>& mu = this->Internals->Adjacency[u];
      if (mu.find(v) == mu.end())
      {
        cost = this->CalculateStaticEdgeCost(image, u, v);
        mu.insert(std::pair<int, double>(v, cost));
      }

      std::map<int, double>& mv = this->Internals->Adjacency[v];
      if (mv.find(u) == mv.end())
      {
        cost = this->CalculateStaticEdgeCost(image, v, u);
        mv.insert(std::pair<int, double>(u, cost));
      }
    }
  }

  ptIds->Delete();

  this->RebuildStaticCosts = false;
  this->AdjacencyBuildTime.Modified();
}

//------------------------------------------------------------------------------
void vtkDijkstraImageGeodesicPath::UpdateStaticCosts(vtkImageData* image)
{
  for (int u = 0; u < static_cast<int>(this->Internals->Adjacency.size()); ++u)
  {
    std::map<int, double>& mu = this->Internals->Adjacency[u];
    std::map<int, double>::iterator it = mu.begin();
    for (; it != mu.end(); ++it)
    {
      int v = (*it).first;
      (*it).second = this->CalculateStaticEdgeCost(image, u, v);
    }
  }
  this->RebuildStaticCosts = false;
}

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

  os << indent << "ImageWeight: " << this->ImageWeight << endl;
  os << indent << "EdgeLengthWeight: " << this->EdgeLengthWeight << endl;
  os << indent << "CurvatureWeight: " << this->CurvatureWeight << endl;
}
VTK_ABI_NAMESPACE_END