/*=========================================================================

  Program:   Visualization Toolkit
  Module:    vtkProportionalEditFilter.cxx

  Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
  All rights reserved.
  See Copyright.txt or http://www.kitware.com/Copyright.htm 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 "vtk/aeva/ext/vtkProportionalEditFilter.h"

#include <cstdlib>
#include <ctime>

#include "vtkCellData.h"
#include "vtkFloatArray.h"
#include "vtkInformation.h"
#include "vtkInformationVector.h"
#include "vtkMath.h"
#include "vtkPointData.h"
#include "vtkPolyData.h"

namespace
{
std::map<std::string, vtkProportionalEditFilter::FalloffType> initNameToType()
{
  using Pair = std::pair<std::string, vtkProportionalEditFilter::FalloffType>;
  std::map<std::string, vtkProportionalEditFilter::FalloffType> nameToType;

  nameToType.insert(Pair(std::string("Linear"), vtkProportionalEditFilter::Linear));
  nameToType.insert(Pair(std::string("Smooth"), vtkProportionalEditFilter::Smooth));
  nameToType.insert(Pair(std::string("Constant"), vtkProportionalEditFilter::Constant));
  nameToType.insert(Pair(std::string("Random"), vtkProportionalEditFilter::Random));
  nameToType.insert(Pair(std::string("Sphere"), vtkProportionalEditFilter::Sphere));
  nameToType.insert(Pair(std::string("Sharp"), vtkProportionalEditFilter::Sharp));
  nameToType.insert(Pair(std::string("InverseSquare"), vtkProportionalEditFilter::InverseSquare));
  return nameToType;
}

std::map<vtkProportionalEditFilter::FalloffType, std::string> initTypeToName()
{
  using Pair = std::pair<vtkProportionalEditFilter::FalloffType, std::string>;
  std::map<vtkProportionalEditFilter::FalloffType, std::string> typeToName;

  typeToName.insert(Pair(vtkProportionalEditFilter::Linear, std::string("Linear")));
  typeToName.insert(Pair(vtkProportionalEditFilter::Smooth, std::string("Smooth")));
  typeToName.insert(Pair(vtkProportionalEditFilter::Constant, std::string("Constant")));
  typeToName.insert(Pair(vtkProportionalEditFilter::Random, std::string("Random")));
  typeToName.insert(Pair(vtkProportionalEditFilter::Sphere, std::string("Sphere")));
  typeToName.insert(Pair(vtkProportionalEditFilter::Sharp, std::string("Sharp")));
  typeToName.insert(Pair(vtkProportionalEditFilter::InverseSquare, std::string("InverseSquare")));
  return typeToName;
}
} // namespace

const std::map<std::string, vtkProportionalEditFilter::FalloffType>
  vtkProportionalEditFilter::Falloff::nameToType(initNameToType());
const std::map<vtkProportionalEditFilter::FalloffType, std::string>
  vtkProportionalEditFilter::Falloff::typeToName(initTypeToName());

void vtkProportionalEditFilter::Falloff::SetType(const std::string& name)
{
  auto it = Falloff::nameToType.find(name);
  if (it != Falloff::nameToType.end())
  {
    this->type = it->second;
  }
  else
  {
    // used Smooth as default
    // todo: print warning info here, but not sure how to do it.
    this->type = Linear;
  }

  this->SetType(this->type);
}

void vtkProportionalEditFilter::Falloff::SetType(const FalloffType& type)
{
  switch (type)
  {
    case Smooth:
      this->eval = Falloff::SmoothEval;
      return;
    case Linear:
      this->eval = Falloff::LinearEval;
      return;
    case Constant:
      this->eval = Falloff::ConstantEval;
      return;
    case Random:
      this->eval = Falloff::RandomEval;
      return;
    case Sphere:
      this->eval = Falloff::SphereEval;
      return;
    case InverseSquare:
      this->eval = Falloff::InverseSquareEval;
      return;
    case Sharp:
      this->eval = Falloff::SharpEval;
      return;
    default:
      this->eval = Falloff::SmoothEval;
      return;
  }
}

vtkStandardNewMacro(vtkProportionalEditFilter);

// The following code defines methods for the vtkProportionalEditFilter class
//

vtkProportionalEditFilter::vtkProportionalEditFilter()
{
  this->AnchorPointCoords[0] = 0.0;
  this->AnchorPointCoords[1] = 0.0;
  this->AnchorPointCoords[2] = 0.0;
  this->InfluenceRadius = 0.0;
  this->Projected = false;
  this->falloff = Falloff();
  this->GenerateEditDistance = false;
  this->GenerateEditVector = false;
  this->OutputPointsPrecision = vtkAlgorithm::DEFAULT_PRECISION;
  this->SetNumberOfInputPorts(1);
}

int vtkProportionalEditFilter::RequestData(vtkInformation* /*request*/,
  vtkInformationVector** inputVector,
  vtkInformationVector* outputVector)
{
  // get the info objects
  vtkInformation* inInfo = inputVector[0]->GetInformationObject(0);
  vtkInformation* outInfo = outputVector->GetInformationObject(0);

  // get the input and output
  vtkPolyData* input = vtkPolyData::SafeDownCast(inInfo->Get(vtkDataObject::DATA_OBJECT()));
  vtkPolyData* output = vtkPolyData::SafeDownCast(outInfo->Get(vtkDataObject::DATA_OBJECT()));

  // Check input
  //
  vtkIdType numPts = input->GetNumberOfPoints();
  if (numPts < 1)
  {
    vtkDebugMacro(<< "No points to be edited!");
    return 0;
  }

  vtkDebugMacro(<< "Proportional Editing with\n"
                << "\tCoordinates of Anchor Point = [" << this->AnchorPointCoords[0] << ", "
                << this->AnchorPointCoords[1] << ", " << this->AnchorPointCoords[2] << "]\n"
                << "\tDisplacement Of Anchor Point = [" << this->AnchorPointDisplacement[0] << ", "
                << this->AnchorPointDisplacement[1] << ", " << this->AnchorPointDisplacement[2]
                << "]\n"
                << "\tFalloff = " << this->falloff.GetTypeName() << "\n"
                << "\tInfluence Radius = " << this->InfluenceRadius << "\n"
                << "\tProjected = " << (this->Projected ? "On\n" : "Off\n"));
  if (this->Projected)
  {
    vtkDebugMacro(<< "\tProjectedDirection = [" << this->ProjectionDirection[0] << ", "
                  << this->ProjectionDirection[1] << ", " << this->ProjectionDirection[2] << "]\n");
  }

  vtkPoints* newPts = vtkPoints::New();
  vtkPoints* inPts;
  inPts = input->GetPoints();

  // Set the desired precision for the points in the output.
  if (this->OutputPointsPrecision == vtkAlgorithm::DEFAULT_PRECISION)
  {
    newPts->SetDataType(inPts->GetDataType());
  }
  else if (this->OutputPointsPrecision == vtkAlgorithm::SINGLE_PRECISION)
  {
    newPts->SetDataType(VTK_FLOAT);
  }
  else if (this->OutputPointsPrecision == vtkAlgorithm::DOUBLE_PRECISION)
  {
    newPts->SetDataType(VTK_DOUBLE);
  }

  newPts->SetNumberOfPoints(numPts);
  for (vtkIdType i = 0; i < numPts; i++) // initialize to old coordinates
  {
    newPts->SetPoint(i, inPts->GetPoint(i));
  }

  // Update output. Only point coordinates have changed.
  //
  // output->CopyStructure(input);
  output->GetPointData()->PassData(input->GetPointData());
  output->GetCellData()->PassData(input->GetCellData());

  double center[3] = {
    this->AnchorPointCoords[0], this->AnchorPointCoords[1], this->AnchorPointCoords[2]
  };
  // in case InfluenceRadius<=0
  const double realRadius = std::max(1e-12, this->InfluenceRadius);
  const double realRadius2 = realRadius * realRadius;

  // normalized ProjectionDirection
  double dir[3];
  if (this->Projected)
  {
    double mag = vtkMath::Norm(this->ProjectionDirection);

    dir[0] = this->ProjectionDirection[0] / mag;
    dir[1] = this->ProjectionDirection[1] / mag;
    dir[2] = this->ProjectionDirection[2] / mag;
  }

  std::function<double(const double[3], const double[3], const double[3])> distance2;
  // if projected, use point-to-line distance
  if (this->Projected)
  {
    distance2 = [](const double dir[3], const double p[3], const double q[3]) {
      double qp[3] = { q[0] - p[0], q[1] - p[1], q[2] - p[2] }; // vector q->p
      double cprod[3];
      vtkMath::Cross(qp, dir, cprod);
      return vtkMath::Dot(cprod, cprod);
    };
  }
  // else use point-to-point distance
  else
  {
    distance2 = [](const double /*dir*/[3], const double p[3], const double q[3]) {
      return vtkMath::Distance2BetweenPoints(p, q);
    };
  }

  double xyz[3];
  double falloffCoef;
  double dist;
  double dist2;
  for (vtkIdType i = 0; i < numPts; ++i)
  {
    newPts->GetPoint(i, xyz);

    dist2 = distance2(dir, center, xyz);

    if (dist2 > realRadius2)
    {
      continue;
    }
    dist = sqrt(dist2) / realRadius;

    falloffCoef = this->falloff(dist);

    xyz[0] += falloffCoef * this->AnchorPointDisplacement[0];
    xyz[1] += falloffCoef * this->AnchorPointDisplacement[1];
    xyz[2] += falloffCoef * this->AnchorPointDisplacement[2];

    newPts->SetPoint(i, xyz);
  }

  if (this->GenerateEditDistance)
  {
    double x0[3];
    double x1[3];
    vtkFloatArray* newScalars = vtkFloatArray::New();
    newScalars->SetNumberOfTuples(numPts);
    for (vtkIdType i = 0; i < numPts; i++)
    {
      inPts->GetPoint(i, x0);
      newPts->GetPoint(i, x1);
      newScalars->SetComponent(i, 0, sqrt(vtkMath::Distance2BetweenPoints(x0, x1)));
    }
    int idx = output->GetPointData()->AddArray(newScalars);
    output->GetPointData()->SetActiveAttribute(idx, vtkDataSetAttributes::SCALARS);
    newScalars->Delete();
  }

  if (this->GenerateEditVector)
  {
    double x0[3];
    double x1[3];
    double dx[3];
    vtkFloatArray* newVectors = vtkFloatArray::New();
    newVectors->SetNumberOfComponents(3);
    newVectors->SetNumberOfTuples(numPts);
    for (vtkIdType i = 0; i < numPts; i++)
    {
      inPts->GetPoint(i, x0);
      newPts->GetPoint(i, x1);
      for (int j = 0; j < 3; j++)
      {
        dx[j] = x1[j] - x0[j];
      }
      newVectors->SetTuple(i, dx);
    }
    output->GetPointData()->SetVectors(newVectors);
    newVectors->Delete();
  }

  output->SetPoints(newPts);
  newPts->Delete();

  output->SetVerts(input->GetVerts());
  output->SetLines(input->GetLines());
  output->SetPolys(input->GetPolys());
  output->SetStrips(input->GetStrips());

  return 1;
}

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

  os << indent << "Proportional Editing with\n";
  os << indent << "\tCoordinates of Anchor Point = [" << this->AnchorPointCoords[0] << ", "
     << this->AnchorPointCoords[1] << ", " << this->AnchorPointCoords[2] << "]\n";
  os << indent << "\tDisplacement Of Anchor Point = [" << this->AnchorPointDisplacement[0] << ", "
     << this->AnchorPointDisplacement[1] << ", " << this->AnchorPointDisplacement[2] << "]\n";
  os << indent << "\tFalloff = " << this->falloff.GetTypeName() << "\n";
  os << indent << "\tInfluence Radius = " << this->InfluenceRadius << "\n";
  os << indent << "\tProjected = " << (this->Projected ? "On\n" : "Off\n") << indent;

  if (this->Projected)
  {
    os << "\tProjectedDirection = [" << this->ProjectionDirection[0] << ", "
       << this->ProjectionDirection[1] << ", " << this->ProjectionDirection[2] << "]\n";
  }

  os << indent << "Generate Edit Distance: " << (this->GenerateEditDistance ? "On\n" : "Off\n");
  os << indent << "Generate Edit Vector: " << (this->GenerateEditVector ? "On\n" : "Off\n");
  os << indent << "Output Points Precision: " << this->OutputPointsPrecision << "\n";
}
