/** ANTS Landmarks used to initialize an b-spline displacement field ... */

#include "antsUtilities.h"

#include "itkAffineTransform.h"
#include "itkBSplineScatteredDataPointSetToImageFilter.h"
#include "itkContinuousIndex.h"
#include "itkImageFileReader.h"
#include "itkImageFileWriter.h"
#include "itkImageRegionIteratorWithIndex.h"
#include "itkImportImageFilter.h"
#include "itkLandmarkBasedTransformInitializer.h"
#include "itkRigid2DTransform.h"
#include "itkVersorRigid3DTransform.h"
#include "itkTransformFileWriter.h"

#include <string>
#include <vector>
#include <algorithm>

namespace ants
{

template<class ImageType, class PointSetType>
void ReadLabeledPointSetFromImage( typename ImageType::Pointer image, typename PointSetType::Pointer pointSet, std::vector<typename ImageType::PixelType> & labels )
{
  labels.clear();

  typename PointSetType::Pointer allPoints = PointSetType::New();
  allPoints->Initialize();

  itk::ImageRegionIteratorWithIndex<ImageType> It( image, image->GetLargestPossibleRegion() );

  unsigned int count = 0;
  for( It.GoToBegin(); !It.IsAtEnd(); ++It )
    {
    typename ImageType::PixelType value = It.Get();
    if( value != 0 )
      {
      if( std::find( labels.begin(), labels.end(), value ) == labels.end() )
        {
        labels.push_back( value );
        }
      typename PointSetType::PointType point;
      image->TransformIndexToPhysicalPoint( It.GetIndex(), point );
      allPoints->SetPointData( count, value );
      allPoints->SetPoint( count++, point );
      }
    }
  std::sort( labels.begin(), labels.end() );

  pointSet->Initialize();

  for( unsigned int n = 0; n < labels.size(); n++ )
    {
    typename ImageType::PixelType currentLabel = labels[n];
    typename PointSetType::PointType center;
    center.Fill( 0 );
    float N = 0;
    typename PointSetType::PointsContainerConstIterator ItP =
      allPoints->GetPoints()->Begin();
    typename PointSetType::PointDataContainerIterator ItD =
      allPoints->GetPointData()->Begin();
    while( ItP != allPoints->GetPoints()->End() )
      {
      if( ItD.Value() == currentLabel )
        {
        typename PointSetType::PointType point = ItP.Value();
        for( unsigned int d = 0; d < ImageType::ImageDimension; d++ )
          {
          center[d] += point[d];
          }
        N += 1.0;
        }
      ++ItP;
      ++ItD;
      }
    for( unsigned int d = 0; d < ImageType::ImageDimension; d++ )
      {
      center[d] /= N;
      }
    pointSet->SetPoint( n, center );
    pointSet->SetPointData( n, currentLabel );
    }
}

template <unsigned int ImageDimension, class TransformType>
int InitializeLinearTransform( int itkNotUsed( argc ), char *argv[] )
{
  typedef unsigned int                          LabelType;
  typedef itk::Image<LabelType, ImageDimension> LabelImageType;

  typedef itk::LandmarkBasedTransformInitializer<TransformType,
    LabelImageType, LabelImageType> TransformInitializerType;
  typedef typename TransformInitializerType::LandmarkPointContainer  LandmarkContainerType;

  typedef itk::PointSet<LabelType, ImageDimension> PointSetType;

  //
  // Read in the fixed and moving images and convert to point sets
  //

  typedef itk::ImageFileReader<LabelImageType> ImageReaderType;
  typename ImageReaderType::Pointer fixedReader = ImageReaderType::New();
  fixedReader->SetFileName( argv[2] );
  fixedReader->Update();
  typename LabelImageType::Pointer fixedImage = fixedReader->GetOutput();

  typename PointSetType::Pointer fixedPoints = PointSetType::New();
  std::vector<LabelType> fixedLabels;
  ReadLabeledPointSetFromImage<LabelImageType, PointSetType>( fixedImage, fixedPoints, fixedLabels );

  typename ImageReaderType::Pointer movingReader = ImageReaderType::New();
  movingReader->SetFileName( argv[3] );
  movingReader->Update();
  typename LabelImageType::Pointer movingImage = movingReader->GetOutput();

  typename PointSetType::Pointer movingPoints = PointSetType::New();
  std::vector<LabelType> movingLabels;
  ReadLabeledPointSetFromImage<LabelImageType, PointSetType>( movingImage, movingPoints, movingLabels );


  LandmarkContainerType fixedLandmarks;
  typename PointSetType::PointsContainerConstIterator ItF = fixedPoints->GetPoints()->Begin();
  while( ItF != fixedPoints->GetPoints()->End() )
    {
    fixedLandmarks.push_back( ItF.Value() );
    ++ItF;
    }


  LandmarkContainerType movingLandmarks;
  typename PointSetType::PointsContainerConstIterator ItM = movingPoints->GetPoints()->Begin();
  while( ItM != movingPoints->GetPoints()->End() )
    {
    movingLandmarks.push_back( ItM.Value() );
    ++ItM;
    }

  if( fixedLandmarks.size() != movingLandmarks.size() )
    {
    std::cerr << "The number of fixed points and moving points must be the same." << std::endl;
    return EXIT_FAILURE;
    }

  typename std::vector<LabelType>::const_iterator itf;
  for( itf = fixedLabels.begin(); itf != fixedLabels.end(); ++itf )
    {
    if( std::find( movingLabels.begin(), movingLabels.end(), *itf ) == movingLabels.end() )
      {
      std::cerr << "Labels do not match." << std::endl;
      return EXIT_FAILURE;
      }
    }

  //
  // Calculate the transform
  //

  typename TransformType::Pointer transform = TransformType::New();
  transform->SetIdentity();

  typename TransformInitializerType::Pointer initializer = TransformInitializerType::New();
  initializer->SetFixedLandmarks( fixedLandmarks );
  initializer->SetMovingLandmarks( movingLandmarks );
  initializer->SetTransform( transform );
  initializer->InitializeTransform();

  //
  // Write the transform
  //

  typename itk::TransformFileWriter::Pointer transformWriter = itk::TransformFileWriter::New();
  transformWriter->SetFileName( argv[5] );
  transformWriter->SetInput( transform );

  try
    {
    transformWriter->Update();
    }
  catch( itk::ExceptionObject & err )
    {
    std::cerr << "Exception in writing tranform file: " << argv[5] << std::endl;
    return EXIT_FAILURE;
    }

  return EXIT_SUCCESS;
}

template <unsigned int ImageDimension>
int InitializeBSplineTransform( int argc, char *argv[] )
{
  typedef float                                 RealType;
  typedef unsigned int                          LabelType;
  typedef itk::Image<LabelType, ImageDimension> LabelImageType;

  typedef itk::PointSet<LabelType, ImageDimension> PointSetType;

  //
  // Read in the fixed and moving images and convert to point sets
  //

  typedef itk::ImageFileReader<LabelImageType> ImageReaderType;
  typename ImageReaderType::Pointer fixedReader = ImageReaderType::New();
  fixedReader->SetFileName( argv[2] );
  fixedReader->Update();
  typename LabelImageType::Pointer fixedImage = fixedReader->GetOutput();

  itk::ImageRegionIteratorWithIndex<LabelImageType> ItF( fixedImage, fixedImage->GetLargestPossibleRegion() );

  typename PointSetType::Pointer fixedLandmarks = PointSetType::New();
  typename std::vector<LabelType> fixedLabels;

  ReadLabeledPointSetFromImage<LabelImageType, PointSetType>( fixedImage, fixedLandmarks, fixedLabels );

  typename ImageReaderType::Pointer movingReader = ImageReaderType::New();
  movingReader->SetFileName( argv[3] );
  movingReader->Update();
  typename LabelImageType::Pointer movingImage = movingReader->GetOutput();

  typename PointSetType::Pointer movingLandmarks = PointSetType::New();
  typename std::vector<LabelType> movingLabels;
  ReadLabeledPointSetFromImage<LabelImageType, PointSetType>( movingImage, movingLandmarks, movingLabels );

  if( fixedLandmarks->GetNumberOfPoints() != movingLandmarks->GetNumberOfPoints() )
    {
    std::cerr << "The number of fixed points and moving points must be the same." << std::endl;
    return EXIT_FAILURE;
    }

  typename std::vector<LabelType>::const_iterator itf;
  for( itf = fixedLabels.begin(); itf != fixedLabels.end(); ++itf )
    {
    if( std::find( movingLabels.begin(), movingLabels.end(), *itf ) == movingLabels.end() )
      {
      std::cerr << "Labels do not match." << std::endl;
      return EXIT_FAILURE;
      }
    }

  //
  // Calculate the transform
  //

  typename LabelImageType::DirectionType fixedDirection = fixedImage->GetDirection();
  typename LabelImageType::DirectionType fixedDirectionInverse( fixedDirection.GetInverse() );

  typename LabelImageType::DirectionType identityDirection;
  identityDirection.SetIdentity();

  const typename LabelImageType::RegionType & bufferedRegion = fixedImage->GetBufferedRegion();
  const itk::SizeValueType numberOfPixels = bufferedRegion.GetNumberOfPixels();

  const bool filterHandlesMemory = false;

  typedef itk::ImportImageFilter<LabelType, ImageDimension> ImporterType;
  typename ImporterType::Pointer importer = ImporterType::New();
  importer->SetImportPointer( const_cast<LabelType *>( fixedImage->GetBufferPointer() ), numberOfPixels, filterHandlesMemory );
  importer->SetRegion( fixedImage->GetBufferedRegion() );
  importer->SetOrigin( fixedImage->GetOrigin() );
  importer->SetSpacing( fixedImage->GetSpacing() );
  importer->SetDirection( identityDirection );
  importer->Update();

  const typename ImporterType::OutputImageType * parametricInputImage = importer->GetOutput();

  typedef itk::Vector<RealType, ImageDimension>  VectorType;
  typedef itk::Image<VectorType, ImageDimension> DisplacementFieldType;

  // Read in the optional label weights

  std::vector<float>     labelWeights;
  std::vector<LabelType> userLabels;

  bool useWeights = false;

  unsigned int labelCount = 0;
  if( argc > 10 )
    {
    useWeights = true;

    std::fstream labelStr( argv[10] );

    if( labelStr.is_open() )
      {
      while( !labelStr.eof() )
        {
        char line[256];
        labelStr.getline( line, 256 );

        std::string lineString = std::string( line );
        std::size_t pos = lineString.find( ',' );

        RealType value;
        if( pos == std::string::npos )
          {
          std::istringstream iss( lineString );
          iss >> value;
          labelWeights.push_back( value );
          userLabels.push_back( movingLabels[labelCount++] );
          }
        else
          {
          unsigned int localLabel;

          std::string        element = lineString.substr( 0, pos );
          std::istringstream iss( element );
          iss >> localLabel;
          userLabels.push_back( localLabel );

          element = lineString.substr( pos + 1, lineString.length() );
          std::istringstream iss2( element );
          iss2 >> value;
          labelWeights.push_back( value );
          }
        }

      labelStr.close();
      }
    else
      {
      std::cerr << "File " << argv[10] << " cannot be opened." << std::endl;
      return EXIT_FAILURE;
      }
    }

  // Now match up the center points

  typedef itk::PointSet<VectorType, ImageDimension> DisplacementFieldPointSetType;
  typedef itk::BSplineScatteredDataPointSetToImageFilter
    <DisplacementFieldPointSetType, DisplacementFieldType> BSplineFilterType;
  typedef typename BSplineFilterType::WeightsContainerType WeightsContainerType;

  typename WeightsContainerType::Pointer weights = WeightsContainerType::New();
  weights->Initialize();
  const typename WeightsContainerType::Element boundaryWeight = 1.0e10;
  const typename WeightsContainerType::Element weight = 1.0;

  typename DisplacementFieldPointSetType::Pointer fieldPoints = DisplacementFieldPointSetType::New();
  fieldPoints->Initialize();
  unsigned long count = 0;

  typename PointSetType::PointsContainerConstIterator mIt =
    movingLandmarks->GetPoints()->Begin();
  typename PointSetType::PointDataContainerIterator mItD =
    movingLandmarks->GetPointData()->Begin();

  while( mItD != movingLandmarks->GetPointData()->End() )
    {
    typename PointSetType::PointsContainerConstIterator fIt =
      fixedLandmarks->GetPoints()->Begin();
    typename PointSetType::PointDataContainerIterator fItD =
      fixedLandmarks->GetPointData()->Begin();

    while( fItD != fixedLandmarks->GetPointData()->End() )
      {
      if( fItD.Value() == mItD.Value() )
        {
        typename PointSetType::PointType fpoint = fIt.Value();
        typename PointSetType::PointType mpoint = mIt.Value();

        VectorType vector;

        typename LabelImageType::PointType fixedPhysicalPoint;
        for( unsigned int i = 0; i < ImageDimension; i++ )
          {
          fixedPhysicalPoint[i] = fpoint[i];
          vector[i] = mpoint[i] - fpoint[i];
          }

        itk::ContinuousIndex<double, ImageDimension> fixedCidx;
        fixedImage->TransformPhysicalPointToContinuousIndex( fixedPhysicalPoint, fixedCidx );

        typename DisplacementFieldType::PointType fieldPoint;
        parametricInputImage->TransformContinuousIndexToPhysicalPoint( fixedCidx, fieldPoint );

        fieldPoints->SetPoint( count, fieldPoint );
        fieldPoints->SetPointData( count, vector );

        if( useWeights )
          {
          std::vector<LabelType>::const_iterator it = std::find( userLabels.begin(), userLabels.end(), mItD.Value() );
          if( it != userLabels.end() )
            {
            weights->InsertElement( count, labelWeights[it - userLabels.begin()] );
            }
          else
            {
            std::cerr << "Unspecified label " << mItD.Value() << " in specified user label weights." << std::endl;
            return EXIT_FAILURE;
            }
          }
        else
          {
          weights->InsertElement( count, weight );
          }

        count++;

        break;
        }
      ++fItD;
      ++fIt;
      }

    ++mItD;
    ++mIt;
    }

  bool enforceStationaryBoundary = true;
  if( argc > 9 )
    {
    enforceStationaryBoundary = static_cast<bool>( atoi( argv[9] ) );
    }
  if( enforceStationaryBoundary )
    {
    typename LabelImageType::IndexType startIndex2 = fixedImage->GetLargestPossibleRegion().GetIndex();

    typename LabelImageType::SizeType inputSize2 = fixedImage->GetLargestPossibleRegion().GetSize();
    for( ItF.GoToBegin(); !ItF.IsAtEnd(); ++ItF )
      {
      typename LabelImageType::IndexType index = ItF.GetIndex();

      bool isOnStationaryBoundary = false;
      for( unsigned int d = 0; d < ImageDimension; d++ )
        {
        if( index[d] == startIndex2[d] || index[d] == startIndex2[d] + static_cast<int>( inputSize2[d] ) - 1 )
          {
          isOnStationaryBoundary = true;
          break;
          }
        }

      if( isOnStationaryBoundary )
        {
        VectorType vector;

        vector.Fill( 0.0 );

        typename PointSetType::PointType fixedPoint;
        parametricInputImage->TransformIndexToPhysicalPoint( index, fixedPoint );

        fieldPoints->SetPoint( count, fixedPoint );
        fieldPoints->SetPointData( count, vector );
        weights->InsertElement( count, boundaryWeight );
        count++;
        }
      }
    }

  typename BSplineFilterType::Pointer bspliner = BSplineFilterType::New();

  unsigned int numberOfLevels = 4;
  if( argc > 7 )
    {
    numberOfLevels = atoi( argv[7] );
    }

  unsigned int splineOrder = 3;
  if( argc > 8 )
    {
    splineOrder = atoi( argv[8] );
    }


  typename BSplineFilterType::ArrayType ncps;
  ncps.Fill( 1 + splineOrder );

  if( argc > 6 )
    {
    std::vector<unsigned int> meshSize = ConvertVector<unsigned int>( std::string( argv[6] ) );

    if( meshSize.size() == 1 )
      {
      ncps.Fill( meshSize[0] + splineOrder );
      }
    else if( meshSize.size() == ImageDimension )
      {
      for( unsigned int d = 0; d < ImageDimension; d++ )
        {
        ncps[d] = meshSize[d] + splineOrder;
        }
      }
    else
      {
      std::cerr << "Invalid meshSize format." << std::endl;
      }
    }

  bspliner->SetOrigin( fixedImage->GetOrigin() );
  bspliner->SetSpacing( fixedImage->GetSpacing() );
  bspliner->SetSize( fixedImage->GetLargestPossibleRegion().GetSize() );
  bspliner->SetDirection( fixedImage->GetDirection() );
  bspliner->SetGenerateOutputImage( true );
  bspliner->SetNumberOfLevels( numberOfLevels );
  bspliner->SetSplineOrder( splineOrder );
  bspliner->SetNumberOfControlPoints( ncps );
  bspliner->SetInput( fieldPoints );
  bspliner->SetPointWeights( weights );
  bspliner->Update();

  typedef itk::ImageFileWriter<DisplacementFieldType> WriterType;
  typename WriterType::Pointer writer = WriterType::New();
  writer->SetFileName( argv[5] );
  writer->SetInput( bspliner->GetOutput() );
  writer->Update();

  return EXIT_SUCCESS;
}

// entry point for the library; parameter 'args' is equivalent to 'argv' in (argc,argv) of commandline parameters to
// 'main()'
int antsLandmarkBasedTransformInitializer( std::vector<std::string> args, std::ostream* /*out_stream = NULL */)
{
  // put the arguments coming in as 'args' into standard (argc,argv) format;
  // 'args' doesn't have the command name as first, argument, so add it manually;
  // 'args' may have adjacent arguments concatenated into one argument,
  // which the parser should handle
  args.insert( args.begin(), "antsLandmarkBasedTransformInitializer" );
  int     argc = args.size();
  char* * argv = new char *[args.size() + 1];
  for( unsigned int i = 0; i < args.size(); ++i )
    {
    // allocate space for the string plus a null character
    argv[i] = new char[args[i].length() + 1];
    std::strncpy( argv[i], args[i].c_str(), args[i].length() );
    // place the null character in the end
    argv[i][args[i].length()] = '\0';
    }
  argv[argc] = ITK_NULLPTR;
  // class to automatically cleanup argv upon destruction
  class Cleanup_argv
  {
public:
    Cleanup_argv( char* * argv_, int argc_plus_one_ ) : argv( argv_ ), argc_plus_one( argc_plus_one_ )
    {
    }

    ~Cleanup_argv()
    {
      for( unsigned int i = 0; i < argc_plus_one; ++i )
        {
        delete[] argv[i];
        }
      delete[] argv;
    }

private:
    char* *      argv;
    unsigned int argc_plus_one;
  };
  Cleanup_argv cleanup_argv( argv, argc + 1 );

  // antscout->set_stream( out_stream );

  if( argc < 6 )
    {
    std::cerr << "Usage: " << argv[0] << " dimension fixedImage movingImage transformType "
      << "outputTransform [meshSize[0]xmeshSize[1]x...=1x1x1] [numberOfLevels=4] [order=3] "
      << "[enforceStationaryBoundaries=1] [landmarkWeights] " << std::endl;
    std::cerr << std::endl << "Notes:" << std::endl;
    std::cerr << " 1) transformType = \"rigid\", \"affine\", or \"bspline\"." << std::endl;
    std::cerr << " 2) The optional arguments only apply to the bspline transform." << std::endl;
    std::cerr << " 3) The landmark weights are read from a text file where each row is either:" << std::endl;
    std::cerr << "     \"label,labelWeight\" or " << std::endl;
    std::cerr << "     \"labelWeight\" " << std::endl;
    std::cerr << "    If the latter format is used, the label weights are assumed to be arranged in ascending order by label."
      << std::endl;
    if( argc >= 2 &&
        ( std::string( argv[1] ) == std::string("--help") || std::string( argv[1] ) == std::string("-h") ) )
      {
      return EXIT_SUCCESS;
      }
    return EXIT_FAILURE;
    }

  const unsigned int dimension = static_cast<unsigned int>( atoi( argv[1] ) );

  typedef itk::Rigid2DTransform<double>           Rigid2DTransformType;
  typedef itk::VersorRigid3DTransform<double>     Rigid3DTransformType;
  typedef itk::AffineTransform<double, 2>         AffineTransform2DType;
  typedef itk::AffineTransform<double, 3>         AffineTransform3DType;

  std::string transformType( argv[4] );
  ConvertToLowerCase( transformType );

  switch( dimension )
    {
    case 2:
      {
      if( std::strcmp( transformType.c_str(), "affine" ) == 0 )
        {
        return InitializeLinearTransform<2, AffineTransform2DType>( argc, argv );
        }
      else if( std::strcmp( transformType.c_str(), "rigid" ) == 0 )
        {
        return InitializeLinearTransform<2, Rigid2DTransformType>( argc, argv );
        }
      else if( std::strcmp( transformType.c_str(), "bspline" ) == 0 )
        {
        return InitializeBSplineTransform<2>( argc, argv );
        }
      else
        {
        std::cerr << "Unrecognized transform:  " << transformType.c_str() << std::endl;
        }
      }
      break;
    case 3:
      {
      if( std::strcmp( transformType.c_str(), "affine" ) == 0 )
        {
        return InitializeLinearTransform<3, AffineTransform3DType>( argc, argv );
        }
      else if( std::strcmp( transformType.c_str(), "rigid" ) == 0 )
        {
        return InitializeLinearTransform<3, Rigid3DTransformType>( argc, argv );
        }
      else if( std::strcmp( transformType.c_str(), "bspline" ) == 0 )
        {
        return InitializeBSplineTransform<3>( argc, argv );
        }
      else
        {
        std::cerr << "Unrecognized transform:  " << transformType.c_str() << std::endl;
        }
      }
      break;
    default:
      std::cerr << "Unsupported dimension" << std::endl;
      return EXIT_FAILURE;
    }
  return EXIT_SUCCESS;
}
} // namespace ants
