/*=========================================================================
 *
 *  Copyright Insight Software Consortium
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *         http://www.apache.org/licenses/LICENSE-2.0.txt
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
 *=========================================================================*/
/*=========================================================================
 *
 *  Portions of this file are subject to the VTK Toolkit Version 3 copyright.
 *
 *  Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
 *
 *  For complete copyright, license and disclaimer of warranty information
 *  please refer to the NOTICE file at the top of the ITK source tree.
 *
 *=========================================================================*/
#ifndef itkDeformableSimplexMesh3DFilter_hxx
#define itkDeformableSimplexMesh3DFilter_hxx

#include "itkDeformableSimplexMesh3DFilter.h"
#include "itkNumericTraits.h"
#include "itkMath.h"

#include <set>

#include "vxl_version.h"
#include "vnl/vnl_cross.h"

namespace itk
{
/* Constructor. */
template< typename TInputMesh, typename TOutputMesh >
DeformableSimplexMesh3DFilter< TInputMesh, TOutputMesh >
::DeformableSimplexMesh3DFilter()
{
  m_Step = 0;
  m_Iterations = 20;
  m_Alpha = 0.2;
  m_Beta  = 0.01;
  m_Gamma = 0.05;
  m_Damping = 0.65;
  m_Rigidity = 1;

  m_ImageDepth = 0;
  m_ImageHeight = 0;
  m_ImageWidth = 0;

  this->ProcessObject::SetNumberOfRequiredInputs(1);

  OutputMeshPointer output = OutputMeshType::New();
  this->ProcessObject::SetNumberOfRequiredOutputs(1);
  this->ProcessObject::SetNthOutput( 0, output.GetPointer() );

  this->m_Data = nullptr;
}

/* PrintSelf. */
template< typename TInputMesh, typename TOutputMesh >
void
DeformableSimplexMesh3DFilter< TInputMesh, TOutputMesh >
::PrintSelf(std::ostream & os, Indent indent) const
{
  Superclass::PrintSelf(os, indent);
  os << indent << "Alpha = " << this->GetAlpha() << std::endl;
  os << indent << "Beta = " << this->GetBeta() << std::endl;
  os << indent << "Gamma = " << this->GetGamma() << std::endl;
  os << indent << "Rigidity = " << this->GetRigidity() << std::endl;
  os << indent << "Iterations = " << this->GetIterations() << std::endl;
  os << indent << "Step = " << this->GetStep() << std::endl;
  os << indent << "ImageDepth = " << this->GetImageDepth() << std::endl;

  const GradientImageType * gradientImage = this->GetGradient();

  if ( gradientImage )
    {
    os << indent << "Gradient = " << gradientImage << std::endl;
    }
  else
    {
    os << indent << "Gradient = " << "(None)" << std::endl;
    }
  os << indent << "ImageHeight = " << this->GetImageHeight() << std::endl;
  os << indent << "ImageWidth = " << this->GetImageWidth() << std::endl;
  os << indent << "Damping = " << this->GetDamping() << std::endl;
  if ( this->m_Data.IsNotNull() )
    {
    os << indent << "Data = " << this->GetData() << std::endl;
    }
  else
    {
    os << indent << "Data = " << "(None)" << std::endl;
    }
} /* End PrintSelf. */

/* Generate Data */
template< typename TInputMesh, typename TOutputMesh >
void
DeformableSimplexMesh3DFilter< TInputMesh, TOutputMesh >
::GenerateData()
{
  this->Initialize();

  m_Step = 0;

  while ( m_Step < m_Iterations )
    {
    const float progress = static_cast< float >( m_Step )
                           / static_cast< float >( m_Iterations );

    this->UpdateProgress(progress);

    this->ComputeGeometry();

    if ( m_Step % 10 == 0 && m_Step > 0 )
      {
      this->UpdateReferenceMetrics();
      }

    this->ComputeDisplacement();
    m_Step++;
    }

  const InputMeshType *             inputMesh = this->GetInput(0);
  const InputPointsContainer *      points = inputMesh->GetPoints();
  InputPointsContainerConstIterator pointItr = points->Begin();

  while ( pointItr != points->End() )
    {
    SimplexMeshGeometry *data;
    IdentifierType       idx = pointItr.Index();
    data = this->m_Data->GetElement(idx);
    delete data->neighborSet;
    data->neighborSet = nullptr;
    ++pointItr;
    }

  this->ComputeOutput();
}

/* Set default value of parameters and initialize local data container
 *  such as forces, displacements and displacement derivatives. */
template< typename TInputMesh, typename TOutputMesh >
void
DeformableSimplexMesh3DFilter< TInputMesh, TOutputMesh >
::Initialize()
{
  const InputMeshType *             inputMesh = this->GetInput(0);
  const InputPointsContainer *      points = inputMesh->GetPoints();
  InputPointsContainerConstIterator pointItr = points->Begin();

  const GradientImageType * gradientImage = this->GetGradient();
  if (gradientImage->GetSpacing()[0] != 1.0
      || gradientImage->GetSpacing()[1] != 1.0
      || gradientImage->GetSpacing()[2] != 1.0
      || gradientImage->GetOrigin()[0] != 0.0
      || gradientImage->GetOrigin()[1] != 0.0
      || gradientImage->GetOrigin()[2] != 0.0)
    {
    itkExceptionMacro("This filter assumes [0,0,0] origin and spacing of [1,1,1]."
        << " Gradient image has origin: " << gradientImage->GetOrigin()
        << " and spacing: " << gradientImage->GetSpacing());
    }

  if ( gradientImage )
    {
    GradientImageSizeType imageSize = gradientImage->GetBufferedRegion().GetSize();

    m_ImageWidth  = imageSize[0];
    m_ImageHeight = imageSize[1];
    m_ImageDepth  = imageSize[2];
    }
  else
    {
    m_ImageWidth  = 0;
    m_ImageHeight = 0;
    m_ImageDepth  = 0;
    }

  if ( this->m_Data.IsNull() )
    {
    this->m_Data = inputMesh->GetGeometryData();
    }

  while ( pointItr != points->End() )
    {
    SimplexMeshGeometry *data;
    IdentifierType       idx = pointItr.Index();

    data = this->m_Data->GetElement(idx);
    data->pos = pointItr.Value();

    //        InputMeshType::ArrayType neighbors =
    // this->GetInput(0)->GetNeighbors( pointItr.Index() );

    data->neighbors[0] = points->GetElement(data->neighborIndices[0]);
    data->neighbors[1] = points->GetElement(data->neighborIndices[1]);
    data->neighbors[2] = points->GetElement(data->neighborIndices[2]);

    // store neighborset with a specific radius
    InputNeighbors *       neighborsList = inputMesh->GetNeighbors(pointItr.Index(), m_Rigidity);
    auto neighborIt = neighborsList->begin();

    auto * neighborSet = new NeighborSetType();
    while ( neighborIt != neighborsList->end() )
      {
      neighborSet->insert(*neighborIt++);
      }
    // garbage collection (from itkSimplexMesh)
    delete neighborsList;
    data->neighborSet =  neighborSet;

    ++pointItr;
    }

}

/* Set the gradient image as an input */
template< typename TInputMesh, typename TOutputMesh >
void
DeformableSimplexMesh3DFilter< TInputMesh, TOutputMesh >
::SetGradient( const GradientImageType * gradientImage )
{
  this->SetNthInput( 1, const_cast< GradientImageType * >( gradientImage ) );
}

/* Get the gradient image as an input */
template< typename TInputMesh, typename TOutputMesh >
const typename
DeformableSimplexMesh3DFilter< TInputMesh, TOutputMesh >::GradientImageType *
DeformableSimplexMesh3DFilter< TInputMesh, TOutputMesh >
::GetGradient() const
{
  const auto * gradientImage = dynamic_cast< const GradientImageType * >( this->ProcessObject::GetInput(1) );

  return gradientImage;
}

/* Compute normals. */
template< typename TInputMesh, typename TOutputMesh >
void
DeformableSimplexMesh3DFilter< TInputMesh, TOutputMesh >
::ComputeGeometry()
{
  PointType           Foot;
  CovariantVectorType normal;
  CovariantVectorType z;
  z.Fill( 0.0 );
  VectorType          tmp;
  //   IdentifierType idx = 0;

  const InputMeshType *       inputMesh = this->GetInput(0);
  const InputPointsContainer *points  = inputMesh->GetPoints();

  typename GeometryMapType::Iterator dataIt = this->m_Data->Begin();

  SimplexMeshGeometry *data;

  while ( dataIt != this->m_Data->End() )
    {
    //      idx = dataIt.Index();
    data = dataIt.Value();

    data->neighbors[0] = points->GetElement(data->neighborIndices[0]);
    data->neighbors[1] = points->GetElement(data->neighborIndices[1]);
    data->neighbors[2] = points->GetElement(data->neighborIndices[2]);

    // compute normal
    normal.Fill(0.0);

    z.SetVnlVector( vnl_cross_3d( ( data->neighbors[1] - data->neighbors[0] ).GetVnlVector(),
                                    ( data->neighbors[2] - data->neighbors[0] ).GetVnlVector() ) );
    z.Normalize();
    normal += z;

    // copy normal
    data->normal = normal;

    // compute the simplex angle
    data->ComputeGeometry();

    tmp = data->neighbors[0] - data->pos;

    double D = 1.0 / ( 2 * data->sphereRadius ); /* */

    double tmpNormalProd = dot_product( tmp.GetVnlVector(), data->normal.GetVnlVector() );

    double sinphi =  2 *data->circleRadius *D *itk::Math::sgn(tmpNormalProd);
    double phi = std::asin(sinphi);

    data->phi = phi;
    data->meanCurvature = std::abs(sinphi / data->circleRadius);
    tmp = data->pos - data->neighbors[0];

    //compute the foot of p projection of p onto the triangle spanned by its
    // neighbors
    double distance = -tmpNormalProd;
    tmp.SetVnlVector( ( data->pos ).GetVnlVector() - distance * normal.GetVnlVector() );
    Foot.Fill(0.0);
    Foot += tmp;

    data->distance = ( ( data->circleCenter ) - Foot ).GetNorm();

    data->eps = ComputeBarycentricCoordinates(Foot, data);
    dataIt.Value() = data;
    ++dataIt;
    }
}

template< typename TInputMesh, typename TOutputMesh >
void
DeformableSimplexMesh3DFilter< TInputMesh, TOutputMesh >
::ComputeDisplacement()
{
  const InputMeshType *inputMesh = this->GetInput(0);

  const GradientImageType *gradientImage = this->GetGradient();

  // Filters should not modify their input...
  // There is a design flaw here.
  auto * nonConstPoints = const_cast< InputPointsContainer * >( inputMesh->GetPoints() );

  typename GeometryMapType::Iterator dataIt = this->m_Data->Begin();
  SimplexMeshGeometry *data;
  VectorType           displacement;
  displacement.Fill( 0.0 );

  while ( dataIt != this->m_Data->End() )
    {
    data = dataIt.Value();

    this->ComputeInternalForce(data);

    this->ComputeExternalForce(data,gradientImage);

    displacement.SetVnlVector( m_Alpha * ( data->internalForce ).GetVnlVector()
                                 + ( data->externalForce ).GetVnlVector() );

    data->pos += displacement;
    nonConstPoints->InsertElement(dataIt.Index(), data->pos);

    ++dataIt;
    }
}

/* */
template< typename TInputMesh, typename TOutputMesh >
void
DeformableSimplexMesh3DFilter< TInputMesh, TOutputMesh >
::ComputeInternalForce(SimplexMeshGeometry *data)
{
  const double r = data->circleRadius;
  const double d = data->distance;
  const double phi = data->phi;

  const NeighborSetType *const neighborSet = data->neighborSet;

  auto neighborIt = neighborSet->begin();

  double phiRef = 0.0;
  while ( neighborIt != neighborSet->end() )
    {
    phiRef += this->m_Data->GetElement(*neighborIt++)->phi;
    }
  phiRef /= static_cast<double>(neighborSet->size());

  double L=0.0;
  double L_Ref=0.0;
  if ( L_Func(r,d,phi,L) && L_Func(r,d,phi,L_Ref) )
    {
    const PointType eps = data->eps;
    const PointType epsRef = data->referenceMetrics;

    const double eps1Diff = epsRef[0] - eps[0];
    const double eps2Diff = epsRef[1] - eps[1];
    const double eps3Diff = epsRef[2] - eps[2];
    //    diffAbsSum = std::abs(eps1Diff)+std::abs(eps2Diff)+std::abs(eps3Diff);

    VectorType tangentForce;
    tangentForce.SetVnlVector(eps1Diff * ( data->neighbors[0] ).GetVnlVector()
      + eps2Diff * ( data->neighbors[1] ).GetVnlVector()
      + eps3Diff * ( data->neighbors[2] ).GetVnlVector()
    );


    VectorType normalForce;
    normalForce.SetVnlVector( -1.0 * ( L_Ref - L ) * ( data->normal ).GetVnlVector() );

    data->internalForce.Fill(0.0);

    data->internalForce += tangentForce + normalForce;
    }
}

/** Compute model Displacement according to image gradient forces */
template< typename TInputMesh, typename TOutputMesh >
void
DeformableSimplexMesh3DFilter< TInputMesh, TOutputMesh >
::ComputeExternalForce(SimplexMeshGeometry *data,const GradientImageType *gradientImage)
{
  PointType         vec_for, tmp_vec_1, tmp_vec_2, tmp_vec_3;
  GradientIndexType coord, coord2, tmp_co_1, tmp_co_2, tmp_co_3;

  coord[0] = static_cast< GradientIndexValueType >( data->pos[0] );
  coord[1] = static_cast< GradientIndexValueType >( data->pos[1] );
  coord[2] = static_cast< GradientIndexValueType >( data->pos[2] );

  coord2[0] = static_cast< GradientIndexValueType >( std::ceil(data->pos[0]) );
  coord2[1] = static_cast< GradientIndexValueType >( std::ceil(data->pos[1]) );
  coord2[2] = static_cast< GradientIndexValueType >( std::ceil(data->pos[2]) );

  tmp_co_1[0] = coord2[0];
  tmp_co_1[1] = coord[1];
  tmp_co_1[2] = coord[2];

  tmp_co_2[0] = coord[0];
  tmp_co_2[1] = coord2[1];
  tmp_co_2[2] = coord[2];

  tmp_co_3[0] = coord[0];
  tmp_co_3[1] = coord[1];
  tmp_co_3[2] = coord2[2];

  if ( ( coord[0] >= 0 ) && ( coord[1] >= 0 ) && ( coord[2] >= 0 )
       && ( coord2[0] < m_ImageWidth ) && ( coord2[1] < m_ImageHeight ) && ( coord2[2] < m_ImageDepth ) )
    {
    const GradientType & gradient0 = gradientImage->GetPixel(coord);
    const GradientType & gradient1 = gradientImage->GetPixel(tmp_co_1);
    const GradientType & gradient2 = gradientImage->GetPixel(tmp_co_2);
    const GradientType & gradient3 = gradientImage->GetPixel(tmp_co_3);

    vec_for[0] = gradient0[0];
    vec_for[1] = gradient0[1];
    vec_for[2] = gradient0[2];

    tmp_vec_1[0] = gradient1[0] - gradient0[0];
    tmp_vec_1[1] = gradient1[1] - gradient0[1];
    tmp_vec_1[2] = gradient1[2] - gradient0[2];
    tmp_vec_2[0] = gradient2[0] - gradient0[0];
    tmp_vec_2[1] = gradient2[1] - gradient0[1];
    tmp_vec_2[2] = gradient2[2] - gradient0[2];
    tmp_vec_3[0] = gradient3[0] - gradient0[0];
    tmp_vec_3[1] = gradient3[1] - gradient0[1];
    tmp_vec_3[2] = gradient3[2] - gradient0[2];

    vec_for[0] = vec_for[0] + ( ( data->pos )[0] - coord[0] ) * tmp_vec_1[0]
                 + ( ( data->pos )[1] - coord[1] ) * tmp_vec_2[0] + ( ( data->pos )[2] - coord[2] ) * tmp_vec_3[0];
    vec_for[1] = vec_for[1] + ( ( data->pos )[1] - coord[1] ) * tmp_vec_2[1]
                 + ( ( data->pos )[0] - coord[0] ) * tmp_vec_1[1] + ( ( data->pos )[2] - coord[2] ) * tmp_vec_3[1];
    vec_for[2] = vec_for[2] + ( ( data->pos )[2] - coord[2] ) * tmp_vec_3[2]
                 + ( ( data->pos )[1] - coord[1] ) * tmp_vec_2[2] + ( ( data->pos )[0] - coord[0] ) * tmp_vec_1[2];
    }
  else
    {
    vec_for.Fill(0);
    }

  double mag = dot_product( data->normal.GetVnlVector(), vec_for.GetVnlVector() );
  //      double mag = vec_for[0]*(data->normal)[0] +
  // vec_for[1]*(data->normal)[1]+ vec_for[2]*(data->normal)[2];

  vec_for[0] = mag * ( data->normal )[0]; /*num_for*/
  vec_for[1] = mag * ( data->normal )[1]; /*num_for*/
  vec_for[2] = mag * ( data->normal )[2]; /*num_for*/

  mag = vec_for.GetVectorFromOrigin().GetNorm();

  if ( mag > 0.5 )
    {
    for ( int i = 0; i < 3; i++ )
      {
      vec_for[i] = ( 0.5 * vec_for[i] ) / mag;
      }
    }

  data->externalForce[0] = m_Beta * vec_for[0];
  data->externalForce[1] = m_Beta * vec_for[1];
  data->externalForce[2] = m_Beta * vec_for[2];
}

/* Copy the content of m_Location into the Output. */
template< typename TInputMesh, typename TOutputMesh >
void
DeformableSimplexMesh3DFilter< TInputMesh, TOutputMesh >
::ComputeOutput()
{
  OutputMeshType *output = this->GetOutput();

  this->CopyInputMeshToOutputMeshPoints();
  this->CopyInputMeshToOutputMeshPointData();
  this->CopyInputMeshToOutputMeshCells();
  this->CopyInputMeshToOutputMeshCellData();

  output->SetGeometryData(this->m_Data);
  output->SetLastCellId( this->GetInput(0)->GetLastCellId() );
}

/*  */
template< typename TInputMesh, typename TOutputMesh >
void
DeformableSimplexMesh3DFilter< TInputMesh, TOutputMesh >
::UpdateReferenceMetrics()
{
  const InputMeshType *inputMesh = this->GetInput(0);

  // Filters should not change their input.
  // There is a design flaw here.
  auto * nonConstInputMesh = const_cast< InputMeshType * >( inputMesh );

  double H;
  double H_N1;
  double H_N2;
  double H_N3;
  double H_Mean;

  GeometryMapIterator dataIt = this->m_Data->Begin();

  SimplexMeshGeometry *data;

  while ( dataIt != this->m_Data->End() )
    {
    data = dataIt->Value();
    H_N1 = ( (SimplexMeshGeometry *)( this->m_Data->GetElement(data->neighborIndices[0]) ) )->meanCurvature;
    H_N2 = ( (SimplexMeshGeometry *)( this->m_Data->GetElement(data->neighborIndices[1]) ) )->meanCurvature;
    H_N3 = ( (SimplexMeshGeometry *)( this->m_Data->GetElement(data->neighborIndices[2]) ) )->meanCurvature;
    H = data->meanCurvature;

    H_Mean = ( H_N1 + H_N2 + H_N3 ) / 3.0;

    PointType deltaH;

    deltaH[0] = ( H_N1 - H_Mean ) / H_Mean;
    deltaH[1] = ( H_N2 - H_Mean ) / H_Mean;
    deltaH[2] = ( H_N3 - H_Mean ) / H_Mean;

    //deltaH[0] = (H_N1 - H_Mean)/H;
    //deltaH[1] = (H_N2 - H_Mean)/H;
    //deltaH[2] = (H_N3 - H_Mean)/H;

    PointType eps, eps_opt;
    // compute optimal reference metrics
    eps_opt[0] = ( 1.0 / 3.0 ) + m_Gamma * deltaH[0];
    eps_opt[1] = ( 1.0 / 3.0 ) + m_Gamma * deltaH[1];
    eps_opt[2] = ( 1.0 / 3.0 ) + m_Gamma * deltaH[2];

    eps = data->referenceMetrics;

    eps[0] = eps[0] + 0.5 * ( eps_opt[0] - eps[0] );
    eps[1] = eps[1] + 0.5 * ( eps_opt[1] - eps[1] );
    eps[2] = eps[2] + 0.5 * ( eps_opt[2] - eps[2] );

    // set current reference metrics
    data->referenceMetrics = eps;
    nonConstInputMesh->SetPointData(dataIt->Index(), H);
    dataIt.Value() = data;
    //      m_Data->InsertElement(dataIt->Index(),data);
    ++dataIt;
    }
}

template< typename TInputMesh, typename TOutputMesh >
bool DeformableSimplexMesh3DFilter< TInputMesh, TOutputMesh >
::L_Func(const double r, const double d, const double phi, double & output)
{
  const double r2 = r * r;
  const double d2 = d * d;
  const double r2Minusd2 = r2 - d2;
  const double tanPhi = std::tan(phi);

  double eps = 1.0;

  if ( phi * itk::Math::sgn(phi) > itk::Math::pi_over_2 )
    {
    eps = -1.0;
    }
  double tmpSqr = r2 + r2Minusd2 * tanPhi * tanPhi;
  if ( tmpSqr > 0 )
    {
    const double denom = eps * ( std::sqrt(tmpSqr) + r );
    if ( Math::NotAlmostEquals( denom, 0.0 ) )
      {
      output = ( r2Minusd2 * tanPhi ) / denom;
      return true;
      }
    }
  return false;
}

template< typename TInputMesh, typename TOutputMesh >
typename DeformableSimplexMesh3DFilter< TInputMesh, TOutputMesh >::PointType
DeformableSimplexMesh3DFilter< TInputMesh, TOutputMesh >
::ComputeBarycentricCoordinates(PointType p, SimplexMeshGeometry *data)
{
  const PointType a = data->neighbors[0];
  const PointType b = data->neighbors[1];
  const PointType c = data->neighbors[2];

  VectorType n, na, nb, nc;
  n.SetVnlVector( vnl_cross_3d( ( b - a ).GetVnlVector(), ( c - a ).GetVnlVector() ) );
  na.SetVnlVector( vnl_cross_3d( ( c - b ).GetVnlVector(), ( p - b ).GetVnlVector() ) );
  nb.SetVnlVector( vnl_cross_3d( ( a - c ).GetVnlVector(), ( p - c ).GetVnlVector() ) );
  nc.SetVnlVector( vnl_cross_3d( ( b - a ).GetVnlVector(), ( p - a ).GetVnlVector() ) );

  PointType eps;
  eps[0] = dot_product( n.GetVnlVector(), na.GetVnlVector() ) / n.GetSquaredNorm();
  eps[1] = dot_product( n.GetVnlVector(), nb.GetVnlVector() ) / n.GetSquaredNorm();
  eps[2] = dot_product( n.GetVnlVector(), nc.GetVnlVector() ) / n.GetSquaredNorm();

  return eps;
}
} /* end namespace itk. */

#endif // itkDeformableSimplexMesh3DFilter_hxx
