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

  Program:   Visualization Toolkit
  Module:    vtkTexturePackingFilter.h

  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 "vtkTexturePackingFilter.h"

#include "vtkAppendPolyData.h"
#include "vtkInformation.h"
#include "vtkInformationVector.h"
#include "vtkMath.h"
#include "vtkMatrix4x4.h"
#include "vtkPartitionedDataSet.h"
#include "vtkPolyData.h"
#include "vtkTransform.h"
#include "vtkTransformPolyDataFilter.h"

#include <vector>

using namespace std;
using namespace Eigen;

vtkStandardNewMacro(vtkTexturePackingFilter);

vtkTexturePackingFilter::vtkTexturePackingFilter()
{
  this->SetNumberOfInputPorts(2);
  this->SetNumberOfOutputPorts(2);
}

int vtkTexturePackingFilter::RequestData(vtkInformation* /*request*/,
  vtkInformationVector** inputVector,
  vtkInformationVector* outputVector)
{
  // get the info objects
  vtkInformation* inPDSInfo0 = inputVector[ATLAS]->GetInformationObject(0);
  vtkInformation* inPDSInfo1 = inputVector[BOUNDARY]->GetInformationObject(0);
  vtkInformation* outPolyInfo0 = outputVector->GetInformationObject(UNIATLAS);
  vtkInformation* outPolyInfo1 = outputVector->GetInformationObject(HORIZON);

  // get the input and output
  auto* inputPDS0 = vtkPartitionedDataSet::GetData(inPDSInfo0);
  auto* inputPDS1 = vtkPartitionedDataSet::GetData(inPDSInfo1);
  auto* outputPD0 = vtkPolyData::GetData(outPolyInfo0);
  auto* outputPD1 = vtkPolyData::GetData(outPolyInfo1);

  // check inputs
  size_t nChart = inputPDS0->GetNumberOfPartitions();
  if (nChart != inputPDS1->GetNumberOfPartitions())
  {
    vtkErrorMacro("Numbers of partitions do not agree in inputPDSs.");
  }

  // ivars
  vector<vtkSmartPointer<vtkTransform> > transforms;  // transform for each polydatas
  vector<pair<vtkIdType, vtkIdType> > farthestPoints; // farthest points for re-orienting polydatas
  vector<vtkSmartPointer<vtkPolyData> > transformedPDs0;
  vector<vtkSmartPointer<vtkPolyData> > transformedPDs1;
  vector<TexelizedChart> texelizedCharts;
  VectorXi globalHorizon;

  // pre-allocation
  transforms.resize(nChart);
  farthestPoints.resize(nChart);
  transformedPDs0.resize(nChart);
  transformedPDs1.resize(nChart);
  texelizedCharts.resize(nChart);
  globalHorizon.resize(TextureMapWidth);
  globalHorizon.setZero();

  for (int i = 0; i < nChart; ++i)
  {
    transformedPDs0[i] = vtkPolyData::New();
    transformedPDs1[i] = vtkPolyData::New();
    transformedPDs0[i]->DeepCopy(vtkPolyData::SafeDownCast(inputPDS0->GetPartition(i)));
    transformedPDs1[i]->DeepCopy(vtkPolyData::SafeDownCast(inputPDS1->GetPartition(i)));
  }

  // re-position polydatas with the farthest points placed along y axis
  for (int i = 0; i < nChart; ++i)
  {
    auto curPD0 = transformedPDs0[i];
    auto curPD1 = transformedPDs1[i];

    // get point pair with the largest distance.
    pair<vtkIdType, vtkIdType> pointIds;
    Vector3d distanceVector;
    GetFarthestPoints(curPD1, pointIds, distanceVector);

    // 2d rotation matrix extraction
    Rotation2D<double> rot2(vtkMath::Pi() / 2);
    Matrix2d preRot;
    Matrix2d postRot;
    preRot.col(0) = Vector2d(distanceVector.x(), distanceVector.y());
    preRot.col(1) = rot2.toRotationMatrix() * Vector2d(distanceVector.x(), distanceVector.y());
    postRot.coeffRef(1, 0) = distanceVector.norm();
    postRot.coeffRef(0, 1) = -distanceVector.norm();
    Matrix2d rotMat = postRot * preRot.inverse();

    // construct transform
    vtkNew<vtkTransform> transform;
    vtkNew<vtkMatrix4x4> transformMatrix;
    transformMatrix->SetElement(0, 0, rotMat.coeff(0, 0));
    transformMatrix->SetElement(0, 1, rotMat.coeff(0, 1));
    transformMatrix->SetElement(1, 0, rotMat.coeff(1, 0));
    transformMatrix->SetElement(1, 1, rotMat.coeff(1, 1));
    transform->SetMatrix(transformMatrix);
    double translation[3];
    curPD1->GetPoint(pointIds.first, translation);
    transform->Translate(-translation[0], -translation[1], -translation[2]);
    transform->Update();

    // apply the same transform to both atlas and boundary
    vtkNew<vtkTransformPolyDataFilter> transformFilter;
    transformFilter->SetTransform(transform);
    transformFilter->SetInputDataObject(curPD0);
    transformFilter->Update();
    curPD0->ShallowCopy(transformFilter->GetOutputDataObject(0));

    // TODO: why transformFilter change cellTypes from polyline to line
    transformFilter->SetInputDataObject(curPD1);
    transformFilter->Update();
    curPD1->ShallowCopy(transformFilter->GetOutputDataObject(0));
  }

  // texelization
  for (int i = 0; i < nChart; ++i)
  {
    texelizedCharts[i].initialize(transformedPDs1[i], TexelSize, BoundaryTexel);
    if (!texelizedCharts[i].Compute())
    {
      vtkErrorMacro("Failed to compute texelized charts.");
      return 0;
    }
  }

  // Sort partitions with heights
  vector<pair<int, int> > partitionHeightPair;
  partitionHeightPair.resize(nChart);
  for (int i = 0; i < nChart; ++i)
  {
    partitionHeightPair[i] = make_pair(i, texelizedCharts[i].TopHorizon.maxCoeff());
  }

  sort(partitionHeightPair.begin(),
    partitionHeightPair.end(),
    [=](pair<int, int>& a, pair<int, int>& b) { return a.second > b.second; });

  // packing
  for (int i = 0; i < nChart; ++i)
  {
    auto curChart = texelizedCharts[partitionHeightPair[i].first];
    int maxOffset = TextureMapWidth - curChart.NumberOfTexel;
    int minGapArea = VTK_INT_MAX;
    int minGapOffset;

    // displacement for final chart position
    int minGapX;
    int minGapY;

    if (maxOffset < 0)
    {
      vtkErrorMacro("Chart width is larger than texture map width.");
      return 0;
    }

    // loop through all possible offsets
    for (int j = 0; j <= maxOffset; j += StepSize)
    {
      auto gap = curChart.BottomHorizon - globalHorizon.middleRows(j, curChart.NumberOfTexel);
      int curGapArea = gap.sum() - gap.minCoeff() * curChart.NumberOfTexel;
      if (curGapArea < minGapArea)
      {
        minGapArea = curGapArea;
        minGapOffset = j;
        minGapX = -curChart.LeftBound + j;
        minGapY = -gap.minCoeff();
      }
    }

    // update horizon
    globalHorizon.middleRows(minGapOffset, curChart.NumberOfTexel) =
      curChart.TopHorizon + (VectorXi::Ones(curChart.NumberOfTexel) * minGapY);

    // translate
    vtkNew<vtkTransform> transform;
    transform->Translate((double)minGapX * TexelSize, (double)minGapY * TexelSize, 0);
    vtkNew<vtkTransformPolyDataFilter> transformFilter;
    transformFilter->SetTransform(transform);

    auto curPD0 = transformedPDs0[partitionHeightPair[i].first];
    transformFilter->SetInputDataObject(curPD0);
    transformFilter->Update();
    curPD0->ShallowCopy(transformFilter->GetOutputDataObject(0));
  }

  // append combine multiple meshes into one polydata
  vtkNew<vtkAppendPolyData> appendFilter;
  for (int i = 0; i < nChart; ++i)
  {
    appendFilter->AddInputDataObject(transformedPDs0[i]);
  }
  appendFilter->Update();
  outputPD0->ShallowCopy(appendFilter->GetOutputDataObject(0));

  // point for horizon
  vtkNew<vtkPoints> horizonPoints;
  horizonPoints->Allocate(TextureMapWidth);
  for (int i = 0; i < TextureMapWidth; ++i)
  {
    horizonPoints->InsertNextPoint(i * TexelSize, globalHorizon[i] * TexelSize, 0);
  }
  outputPD1->SetPoints(horizonPoints);
  outputPD1->Allocate(TextureMapWidth - 1);
  for (int i = 0; i < TextureMapWidth - 1; ++i)
  {
    vtkNew<vtkIdList> cell;
    cell->SetNumberOfIds(2);
    cell->SetId(0, i);
    cell->SetId(1, i + 1);
    outputPD1->InsertNextCell(VTK_POLY_LINE, cell);
  }

  return 1;
}

int vtkTexturePackingFilter::FillInputPortInformation(int /* port */, vtkInformation* info)
{
  info->Set(vtkDataObject::DATA_TYPE_NAME(), "vtkPartitionedDataSet");
  return 1;
}

void vtkTexturePackingFilter::PrintSelf(ostream& os, vtkIndent indent)
{
  this->Superclass::PrintSelf(os, indent);
  os << indent << "Texel size: " << TexelSize << "\n";
  os << indent << "Texture map width: " << TextureMapWidth << " texels\n";
  //os << indent << "Texture map height: " << m_textureMapHeight << "\n";
}

void vtkTexturePackingFilter::GetFarthestPoints(const vtkSmartPointer<vtkPolyData>& curPD1,
  pair<vtkIdType, vtkIdType>& pointIds,
  Vector3d& distanceVector)
{
  distanceVector.setZero();
  vtkIdType nPoint = curPD1->GetNumberOfPoints();
  for (vtkIdType i = 0; i < nPoint - 1; ++i)
  {
    for (vtkIdType j = i + 1; j < nPoint; ++j)
    {
      auto point0 = Vector3d(curPD1->GetPoint(i));
      auto point1 = Vector3d(curPD1->GetPoint(j));
      auto curDistanceVector = point1 - point0;
      if (curDistanceVector.norm() > distanceVector.norm())
      {
        distanceVector = curDistanceVector;
        pointIds.first = i;
        pointIds.second = j;
      }
    }
  }
}

int vtkTexturePackingFilter::TexelizedChart::Compute()
{
  size_t nCell = CurPD1->GetNumberOfCells();
  vector<VectorXd> cellHorizons;
  vector<int> cellLeftBounds;
  vector<int> cellRightBounds;

  cellHorizons.resize(nCell);
  cellLeftBounds.resize(nCell);
  cellRightBounds.resize(nCell);

  // compute left and right bounds
  for (int i = 0; i < nCell; ++i)
  {
    auto* curCell = CurPD1->GetCell(i);

    /*
    // current vtkPolyDataTransformFilter changes polyline to line
    if (curCell->GetCellType() != VTK_POLY_LINE)
    {
      return 0;
    }
    */

    // TODO: extend to polyline with more than two points.

    double point0[3];
    double point1[3];
    curCell->GetPoints()->GetPoint(0, point0);
    curCell->GetPoints()->GetPoint(1, point1);
    double* pointLeft;
    double* pointRight;

    if (point0[0] < point1[0])
    {
      pointLeft = point0;
      pointRight = point1;
    }
    else
    {
      pointLeft = point1;
      pointRight = point0;
    }

    cellLeftBounds[i] = floor(pointLeft[0] / TexelSize);
    cellRightBounds[i] = ceil(pointRight[0] / TexelSize);

    LeftBound = (cellLeftBounds[i] < LeftBound ? cellLeftBounds[i] : LeftBound);
    RightBound = (cellRightBounds[i] > RightBound ? cellRightBounds[i] : RightBound);

    //compute cell horizon
    int nCellTexel = cellRightBounds[i] - cellLeftBounds[i] + 1;
    cellHorizons[i].resize(nCellTexel);
    for (int j = 0; j < nCellTexel; ++j)
    {
      // assume the left bound point is on the texel center
      cellHorizons[i][j] =
        ((pointRight[1] - pointLeft[1]) * j / (nCellTexel - 1.0) + pointLeft[1]) / TexelSize;
    }
  }

  // compute top and bottom horizons
  int preNumberOfTexel = RightBound - LeftBound + 1;
  VectorXi preBottomHorizon(preNumberOfTexel);
  VectorXi preTopHorizon(preNumberOfTexel);
  preBottomHorizon.setConstant(VTK_INT_MAX);
  preTopHorizon.setZero();

  for (int i = 0; i < nCell; ++i)
  {
    int offset = cellLeftBounds[i] - LeftBound;
    size_t nCellTexel = cellHorizons[i].size();
    for (int j = 0; j < nCellTexel; ++j)
    {
      if (floor(cellHorizons[i][j]) < preBottomHorizon[j + offset])
      {
        preBottomHorizon[j + offset] = floor(cellHorizons[i][j]);
      }
      if (ceil(cellHorizons[i][j]) > preTopHorizon[j + offset])
      {
        preTopHorizon[j + offset] = ceil(cellHorizons[i][j]);
      }
    }
  }

  // add boundary texels for spacing purpose
  RightBound += BoundaryTexel;
  LeftBound -= BoundaryTexel;
  NumberOfTexel = RightBound - LeftBound + 1;

  VectorXi corner(BoundaryTexel);
  for (int i = 0; i < BoundaryTexel; ++i)
  {
    corner[i] = i;
  }

  TopHorizon.resize(NumberOfTexel);
  TopHorizon.segment(BoundaryTexel, preNumberOfTexel) =
    preTopHorizon + (VectorXi::Ones(preNumberOfTexel) * BoundaryTexel);
  TopHorizon.head(BoundaryTexel) = preTopHorizon.head(BoundaryTexel);
  TopHorizon.tail(BoundaryTexel) = preTopHorizon.tail(BoundaryTexel);

  BottomHorizon.resize(NumberOfTexel);
  BottomHorizon.segment(BoundaryTexel, preNumberOfTexel) =
    preBottomHorizon - (VectorXi::Ones(preNumberOfTexel) * BoundaryTexel);
  BottomHorizon.head(BoundaryTexel) = preBottomHorizon.head(BoundaryTexel);
  BottomHorizon.tail(BoundaryTexel) = preBottomHorizon.tail(BoundaryTexel);

  return 1;
}
