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

  This file is part of the elastix software.

  Copyright (c) University Medical Center Utrecht. All rights reserved.
  See src/CopyrightElastix.txt or http://elastix.isi.uu.nl/legal.php 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 notices for more information.

======================================================================*/
#include "SplineKernelTransform/itkThinPlateSplineKernelTransform2.h"
#include "itkTransformixInputPointFileReader.h"

#include <ctime>
#include <fstream>
#include <iomanip>

#include "vnl/algo/vnl_qr.h"
//#include "vnl/algo/vnl_sparse_lu.h"
//#include "vnl/algo/vnl_cholesky.h"
#include "vnl/vnl_matlab_filewrite.h"
#include "vnl/vnl_matrix_fixed.h"
#include "vnl/vnl_sparse_matrix.h"

//-------------------------------------------------------------------------------------
// Helper class to be able to access protected functions and variables.

namespace itk {

template <class TScalarType, unsigned int NDimensions>
class KernelTransformPublic
  : public ThinPlateSplineKernelTransform2<TScalarType, NDimensions>
{
public:
  typedef KernelTransformPublic               Self;
  typedef ThinPlateSplineKernelTransform2<
    TScalarType, NDimensions >                Superclass;
  typedef SmartPointer<Self>                  Pointer;
  typedef SmartPointer<const Self>            ConstPointer;
  itkTypeMacro( KernelTransformPublic, ThinPlateSplineKernelTransform2 );
  itkNewMacro( Self );

  typedef typename Superclass::PointSetType   PointSetType;
  typedef typename Superclass::LMatrixType    LMatrixType;
  typedef typename Superclass::GMatrixType    GMatrixType;
  typedef typename Superclass::InputVectorType  InputVectorType;

  void SetSourceLandmarksPublic( PointSetType * landmarks )
  {
    this->m_SourceLandmarks = landmarks;
    this->m_WMatrixComputed = false;
    this->m_LMatrixComputed = false;
    this->m_LInverseComputed = false;
  }

  void ComputeLPublic( void )
  {
    this->ComputeL();
  }

  LMatrixType GetLMatrix( void ) const
  {
    return this->m_LMatrix;
  }

  void ComputeGPublic( const InputVectorType & landmarkVector,
    GMatrixType & GMatrix ) const
  {
    this->ComputeG( landmarkVector, GMatrix );
  }

}; // end helper class
} // end namespace itk

//-------------------------------------------------------------------------------------

// Test matrix inversion performance
// Test Jacobian computation performance
int main( int argc, char *argv[] )
{
  /** Some basic type definitions. */
  const unsigned int Dimension = 3;
  // ScalarType double needed for Cholesky. Double is used in elastix.
  typedef double   ScalarType;
  const unsigned long maxTestedLandmarksForSVD = 401;
  const ScalarType tolerance = 1e-8; // for double

  /** Check. */
  if ( argc != 3 )
  {
    std::cerr << "ERROR: You should specify a text file with the thin plate spline "
      << "source (fixed image) landmarks." << std::endl;
    return 1;
  }

  /** Other typedefs. */
  typedef itk::KernelTransformPublic<
    ScalarType, Dimension >                             TransformType;
  typedef TransformType::JacobianType                   JacobianType;
  typedef TransformType::NonZeroJacobianIndicesType     NonZeroJacobianIndicesType;
  typedef TransformType::PointSetType                   PointSetType;
  typedef itk::TransformixInputPointFileReader<
    PointSetType >                                      IPPReaderType;

  typedef PointSetType::PointsContainer                 PointsContainerType;
  typedef PointsContainerType::Pointer                  PointsContainerPointer;
  typedef PointSetType::PointType                       PointType;
  typedef TransformType::LMatrixType                    LMatrixType;
  typedef vnl_sparse_matrix<ScalarType>                 LSparseMatrixType;

  PointSetType::Pointer dummyLandmarks = PointSetType::New();

  /** Create the kernel transform. */
  TransformType::Pointer kernelTransform = TransformType::New();
  kernelTransform->SetStiffness( 0.0 ); // interpolating

  /** Read landmarks. */
  IPPReaderType::Pointer ippReader = IPPReaderType::New();
  ippReader->SetFileName( argv[ 1 ] );
  try
  {
    ippReader->Update();
  }
  catch ( itk::ExceptionObject & excp )
  {
    std::cerr << "  Error while opening input point file." << std::endl;
    std::cerr << excp << std::endl;
    return 1;
  }

  // Expect points, not indices.
  if ( ippReader->GetPointsAreIndices() )
  {
    std::cerr << "ERROR: landmarks should be specified as points (not indices)"
      << std::endl;
    return 1;
  }

  /** Get the set of input points. */
  PointSetType::Pointer sourceLandmarks = ippReader->GetOutput();
  //const unsigned long realNumberOfLandmarks = ippReader->GetNumberOfPoints();

  std::vector<unsigned long> usedNumberOfLandmarks;
  usedNumberOfLandmarks.push_back( 100 );
  usedNumberOfLandmarks.push_back( 200 );
//   usedNumberOfLandmarks.push_back( 500 );
//   usedNumberOfLandmarks.push_back( 1000 );
//   usedNumberOfLandmarks.push_back( realNumberOfLandmarks );

  std::cerr << "Matrix scalar type: "
    << typeid( ScalarType ).name() << std::endl;

  // Loop over usedNumberOfLandmarks
  for ( std::size_t i = 0; i < usedNumberOfLandmarks.size(); i++ )
  {
    unsigned long numberOfLandmarks = usedNumberOfLandmarks[ i ];
    std::cerr << "----------------------------------------\n";
    std::cerr << "Number of specified landmarks: "
      << numberOfLandmarks << std::endl;

    /** Get subset. */
    PointsContainerPointer usedLandmarkPoints = PointsContainerType::New();
    PointSetType::Pointer usedLandmarks = PointSetType::New();
    for ( unsigned long j = 0; j < numberOfLandmarks; j++ )
    {
      PointType tmp = (*sourceLandmarks->GetPoints())[ j ];
      usedLandmarkPoints->push_back( tmp );
    }
    usedLandmarks->SetPoints( usedLandmarkPoints );

    /** Set the ipp as source landmarks.
     * 1) Compute L matrix
     * 2) Compute inverse of L
     */

    LMatrixType lMatrixInverse1, lMatrixInverse2; //, lMatrixInverse4;

    /** Task 1: compute L. */
    clock_t startClock = clock();
    kernelTransform->SetSourceLandmarksPublic( usedLandmarks );
    kernelTransform->ComputeLPublic();
    LMatrixType lMatrix = kernelTransform->GetLMatrix();
    std::cerr << "Computing L matrix took "
      << clock() - startClock << " ms." << std::endl;

    /** Task 2: Compute L inverse. */
    if ( numberOfLandmarks < maxTestedLandmarksForSVD )
    {
      // Method 1: Singular Value Decomposition
      startClock = clock();
      lMatrixInverse1 = vnl_svd<ScalarType>( lMatrix ).inverse();
      std::cerr << "L matrix inversion (method 1, svd) took: "
        << clock() - startClock << " ms." << std::endl;
    }
    else
    {
      std::cerr << "L matrix inversion (method 1, svd) took: too long" << std::endl;
    }

    // Method 2: QR Decomposition
    startClock = clock();
    lMatrixInverse2 = vnl_qr<ScalarType>( lMatrix ).inverse();
    std::cerr << "L matrix inversion (method 2,  qr) took: "
      << clock() - startClock << " ms." << std::endl;

    // Method 3: Cholesky decomposition
    // Cholesky decomposition does not work due to lMatrix not being positive definite.
    //   startClock = clock();
    //   LMatrixType lMatrixInverse3 = vnl_cholesky( lMatrix,
    //     vnl_cholesky::Operation::estimate_condition ).inverse();
    //   std::cerr << "L matrix inversion (method 3, cholesky ) took: "
    //     << clock() - startClock << " ms." << std::endl;

    /** The following code is out-commented.
     * It is used to test LU decomposition, which in vnl is only implemented
     * for sparse matrices. It also depends on a local modification of the
     * vnl_sparse_lu claas, where a method invert() was implemented similar
     * to the invert() of vnl_qr.inverse().
     */
//     // Convert to sparse matrix
//     startClock = clock();
//     LSparseMatrixType lSparseMatrix( lMatrix.rows(), lMatrix.cols() );
//     for ( unsigned int r = 0; r < lMatrix.rows(); r++ )
//     {
//       for ( unsigned int c = 0; c < lMatrix.cols(); c++ )
//       {
//         ScalarType val = lMatrix.get( r, c );
//         if ( val != 0 )
//         {
//           lSparseMatrix( r, c ) = val;
//         }
//       }
//     }
//     std::cerr << "Conversion to sparse matrix took: "
//       << clock() - startClock << " ms." << std::endl;
//
//     // Method 4: LU Decomposition
//     // Depends on local ITK vnl_sparse_lu modification
//     startClock = clock();
//     lMatrixInverse4 = vnl_sparse_lu( lSparseMatrix ).inverse();
//     std::cerr << "L matrix inversion (method 4,  lu) took: "
//       << clock() - startClock << " ms." << std::endl;

    /** Compute error compared to SVD. */
    if ( numberOfLandmarks < maxTestedLandmarksForSVD )
    {
      double diff_qr = (lMatrixInverse1 - lMatrixInverse2).frobenius_norm();
      //double diff_lu = (lMatrixInverse1a - lMatrixInverse4).frobenius_norm();

      std::cerr << "Frobenius difference of method 2 with SVD: "
        << diff_qr << std::endl;
      //std::cerr << "Frobenius difference of method 4 with SVD: " << diff_lu << std::endl;

      if ( diff_qr > tolerance )
      {
        std::cerr
          << "ERROR: Frobenius difference of matrix inversion methods too big: "
          << diff_qr << std::endl;
        return 1;
      }
    }
    else
    {
      std::cerr << "Frobenius difference of method 2,4 with SVD: unknown" << std::endl;
    }

    //   startClock = clock();
    //   LMatrixType lMatrixInverse3 = vnl_lu<ScalarType>( kernelTransform->GetLMatrix() ).inverse();
    //   std::cerr << "L matrix inversion (method 2, lu ) took: "
    //     << clock() - startClock << " ms." << std::endl;

    // To do: Add SuiteSparse tests.

    // Write L Matrix to Matlab file. For inspection of matrix appearance.
    std::ostringstream makeFileName( "" );
    makeFileName << argv[ 2 ]
      << "/LMatrix_N"
      << numberOfLandmarks << ".mat";
    vnl_matlab_filewrite matlabWriter( makeFileName.str().c_str() );
    matlabWriter.write( lMatrix, "lMatrix" );
    matlabWriter.write( lMatrixInverse2, "lMatrixInverseQR" );

    //
    // Test Jacobian computation performance

    typedef vnl_matrix_fixed<ScalarType, Dimension, Dimension> GMatrixType;
    GMatrixType Gmatrix; // dim x dim
    typedef PointSetType::PointsContainerIterator      PointsIterator;

    // OLD way:
    PointType p; p[0] = 10.0; p[1] = 13.0; p[2] = 11.0;
    startClock = clock();
    JacobianType jac1;
    jac1.SetSize( Dimension, numberOfLandmarks * Dimension );
    jac1.Fill( 0.0 );
    PointsIterator sp = usedLandmarks->GetPoints()->Begin();
    for ( unsigned int lnd = 0; lnd < numberOfLandmarks; lnd++ )
    {
      kernelTransform->ComputeGPublic( p - sp->Value(), Gmatrix );
      for ( unsigned int dim = 0; dim < Dimension; dim++ )
      {
        for ( unsigned int odim = 0; odim < Dimension; odim++ )
        {
          for ( unsigned int lidx = 0; lidx < numberOfLandmarks * Dimension; lidx++ )
          {
            jac1[ odim ][ lidx ] += Gmatrix( dim, odim )
              * lMatrixInverse2[ lnd * Dimension + dim ][ lidx ];
          }
        }
      }
      ++sp;
    }

    for ( unsigned int odim = 0; odim < Dimension; odim++ )
    {
      for ( unsigned long lidx = 0; lidx < numberOfLandmarks * Dimension; lidx++ )
      {
        for ( unsigned int dim = 0; dim < Dimension; dim++ )
        {
          jac1[ odim ][ lidx ] += p[ dim ]
          * lMatrixInverse2[ ( numberOfLandmarks + dim ) * Dimension + odim ][ lidx ];
        }
        const unsigned long index = ( numberOfLandmarks + Dimension ) * Dimension + odim;
        jac1[ odim ][ lidx ] += lMatrixInverse2[ index ][ lidx ];
      }
    }
    std::cerr << "\nJacobian computation (OLD) took: "
      << clock() - startClock << " ms." << std::endl;

    // NEW way:

    /** Reset source landmarks, otherwise L is not recomputed. */
    kernelTransform->SetSourceLandmarks( dummyLandmarks );
    kernelTransform->SetSourceLandmarks( usedLandmarks );
    startClock = clock();
    JacobianType jac2;
    NonZeroJacobianIndicesType nzji;
    kernelTransform->GetJacobian( p, jac2, nzji );
    std::cerr << "Jacobian computation (NEW) took: "
      << clock() - startClock << " ms." << std::endl;

    // diff
    double diff_jac = (jac1 - jac2).frobenius_norm();
    std::cerr << "Frobenius difference of jacs: " << diff_jac << std::endl;
    if ( diff_jac > tolerance )
    {
      std::cerr << "ERROR: Frobenius difference of Jacobian computation too big: " << diff_jac << std::endl;
      return 1;
    }

  } // end loop

  /** Return a value. */
  return 0;

} // end main
