#include <nanobind/nanobind.h>

#include <nanobind/stl/vector.h>

#include <nanobind/stl/string.h>

#include <nanobind/stl/tuple.h>

#include <nanobind/stl/list.h>

#include <nanobind/ndarray.h>

#include <nanobind/stl/shared_ptr.h>

#include "itkMacro.h"

#include "itkImage.h"

#include "itkVectorImage.h"

#include "itkVector.h"

#include "itkImageRegionIteratorWithIndex.h"

#include "vnl/vnl_vector_ref.h"

#include "itkTransform.h"

#include "itkAffineTransform.h"

#include "itkTransformFileReader.h"

#include "antscore/antsUtilities.h"

#include "itkAffineTransform.h"

#include "itkCenteredAffineTransform.h"

#include "itkEuler2DTransform.h"

#include "itkEuler3DTransform.h"

#include "itkRigid2DTransform.h"

#include "itkRigid3DTransform.h"

#include "itkCenteredRigid2DTransform.h"

#include "itkCenteredEuler3DTransform.h"

#include "itkSimilarity2DTransform.h"

#include "itkCenteredSimilarity2DTransform.h"

#include "itkSimilarity3DTransform.h"

#include "itkQuaternionRigidTransform.h"

#include "itkTranslationTransform.h"

#include "itkResampleImageFilter.h"

#include "itkTransformFileReader.h"

#include "itkCompositeTransform.h"

#include "itkMatrixOffsetTransformBase.h"

#include "itkDisplacementFieldTransform.h"

#include "itkConstantBoundaryCondition.h"

#include "itkBSplineInterpolateImageFunction.h"

#include "itkLinearInterpolateImageFunction.h"

#include "itkGaussianInterpolateImageFunction.h"

#include "itkInterpolateImageFunction.h"

#include "itkNearestNeighborInterpolateImageFunction.h"

#include "itkWindowedSincInterpolateImageFunction.h"

#include "itkLabelImageGaussianInterpolateImageFunction.h"

#include "itkTransformFileWriter.h"

#include "itkMacro.h"

#include "itkImage.h"

#include "itkVectorImage.h"

#include "itkVector.h"

#include "itkImageRegionIteratorWithIndex.h"

#include "vnl/vnl_vector_ref.h"

#include "itkTransform.h"

#include "itkAffineTransform.h"

#include "readTransform.h"

#include "register_transforms.h"

namespace nb = nanobind;

using namespace nb::literals;

unsigned int getTransformDimensionFromFile( std::string filename )

{

    register_transforms();

    typedef itk::TransformFileReader TransformReaderType1;

    typedef typename TransformReaderType1::Pointer TransformReaderType;

    TransformReaderType reader = itk::TransformFileReader::New();

    reader->SetFileName( filename.c_str() );

    reader->Update();

    const TransformReaderType1::TransformListType * transforms = reader->GetTransformList();

    const TransformReaderType1::TransformPointer tx = *(transforms->begin());

    return tx->GetInputSpaceDimension();

}

std::string getTransformNameFromFile( std::string filename )

{

    register_transforms();

    typedef itk::TransformFileReader TransformReaderType1;

    typedef typename TransformReaderType1::Pointer TransformReaderType;

    TransformReaderType reader = itk::TransformFileReader::New();

    reader->SetFileName( filename.c_str() );

    reader->Update();

    const TransformReaderType1::TransformListType * transforms = reader->GetTransformList();

    const TransformReaderType1::TransformPointer tx = *(transforms->begin());

    return std::string( tx->GetNameOfClass() );

}

template <typename PrecisionType, unsigned int Dimension>

AntsTransform<itk::Transform<PrecisionType, Dimension, Dimension>> newAntsTransform( std::string precision, unsigned int dimension, std::string type)

{

    //auto transformPointer = TransformType::New();

  // Initialize transform by type

  if ( type == "AffineTransform" )

    {

    typedef itk::AffineTransform<PrecisionType,Dimension> TransformType;

    auto transformPointer = TransformType::New();

    typedef itk::Transform<PrecisionType,Dimension,Dimension> TransformBaseType;

    typedef typename TransformBaseType::Pointer               TransformBasePointerType;

    TransformBasePointerType basePointer = dynamic_cast<TransformBaseType *>( transformPointer.GetPointer() );

    AntsTransform<TransformBaseType> outTransform = { basePointer };

    return outTransform;

    }

  else if ( type == "CenteredAffineTransform" )

    {

    typedef itk::CenteredAffineTransform<PrecisionType,Dimension> TransformType;

    auto transformPointer = TransformType::New();

    typedef itk::Transform<PrecisionType,Dimension,Dimension> TransformBaseType;

    typedef typename TransformBaseType::Pointer               TransformBasePointerType;

    TransformBasePointerType basePointer = dynamic_cast<TransformBaseType *>( transformPointer.GetPointer() );

    AntsTransform<TransformBaseType> outTransform = { basePointer };

    return outTransform;

    }

  else if ( type == "Euler3DTransform" )

    {

    typedef itk::Euler3DTransform<PrecisionType> TransformType;

    auto transformPointer = TransformType::New();

    typedef itk::Transform<PrecisionType,Dimension,Dimension> TransformBaseType;

    typedef typename TransformBaseType::Pointer               TransformBasePointerType;

    TransformBasePointerType basePointer = dynamic_cast<TransformBaseType *>( transformPointer.GetPointer() );

    AntsTransform<TransformBaseType> outTransform = { basePointer };

    return outTransform;

    }

  else if ( type == "Euler2DTransform" )

    {

    typedef itk::Euler2DTransform<PrecisionType> TransformType;

    auto transformPointer = TransformType::New();

    typedef itk::Transform<PrecisionType,Dimension,Dimension> TransformBaseType;

    typedef typename TransformBaseType::Pointer               TransformBasePointerType;

    TransformBasePointerType basePointer = dynamic_cast<TransformBaseType *>( transformPointer.GetPointer() );

    AntsTransform<TransformBaseType> outTransform = { basePointer };

    return outTransform;

    }

  else if ( type == "QuaternionRigidTransform" )

    {

    typedef itk::QuaternionRigidTransform<PrecisionType> TransformType;

    auto transformPointer = TransformType::New();

    typedef itk::Transform<PrecisionType,Dimension,Dimension> TransformBaseType;

    typedef typename TransformBaseType::Pointer               TransformBasePointerType;

    TransformBasePointerType basePointer = dynamic_cast<TransformBaseType *>( transformPointer.GetPointer() );

    AntsTransform<TransformBaseType> outTransform = { basePointer };

    return outTransform;

    }

  else if ( type == "Rigid2DTransform" )

    {

    typedef itk::Rigid2DTransform<PrecisionType> TransformType;

    auto transformPointer = TransformType::New();

    typedef itk::Transform<PrecisionType,Dimension,Dimension> TransformBaseType;

    typedef typename TransformBaseType::Pointer               TransformBasePointerType;

    TransformBasePointerType basePointer = dynamic_cast<TransformBaseType *>( transformPointer.GetPointer() );

    AntsTransform<TransformBaseType> outTransform = { basePointer };

    return outTransform;

    }

  else if ( type == "Rigid3DTransform" )

    {

    typedef itk::Rigid3DTransform<PrecisionType> TransformType;

    auto transformPointer = TransformType::New();

    typedef itk::Transform<PrecisionType,Dimension,Dimension> TransformBaseType;

    typedef typename TransformBaseType::Pointer               TransformBasePointerType;

    TransformBasePointerType basePointer = dynamic_cast<TransformBaseType *>( transformPointer.GetPointer() );

    AntsTransform<TransformBaseType> outTransform = { basePointer };

    return outTransform;

    }

  else if ( type == "CenteredEuler3DTransform" )

    {

    typedef itk::CenteredEuler3DTransform<PrecisionType> TransformType;

    auto transformPointer = TransformType::New();

    typedef itk::Transform<PrecisionType,Dimension,Dimension> TransformBaseType;

    typedef typename TransformBaseType::Pointer               TransformBasePointerType;

    TransformBasePointerType basePointer = dynamic_cast<TransformBaseType *>( transformPointer.GetPointer() );

    AntsTransform<TransformBaseType> outTransform = { basePointer };

    return outTransform;

    }

  else if ( type == "CenteredRigid2DTransform" )

    {

    typedef itk::CenteredRigid2DTransform<PrecisionType> TransformType;

    auto transformPointer = TransformType::New();

    typedef itk::Transform<PrecisionType,Dimension,Dimension> TransformBaseType;

    typedef typename TransformBaseType::Pointer               TransformBasePointerType;

    TransformBasePointerType basePointer = dynamic_cast<TransformBaseType *>( transformPointer.GetPointer() );

    AntsTransform<TransformBaseType> outTransform = { basePointer };

    return outTransform;

    }

  else if ( type == "Similarity3DTransform" )

    {

    typedef itk::Similarity3DTransform<PrecisionType> TransformType;

    auto transformPointer = TransformType::New();

    typedef itk::Transform<PrecisionType,Dimension,Dimension> TransformBaseType;

    typedef typename TransformBaseType::Pointer               TransformBasePointerType;

    TransformBasePointerType basePointer = dynamic_cast<TransformBaseType *>( transformPointer.GetPointer() );

    AntsTransform<TransformBaseType> outTransform = { basePointer };

    return outTransform;

    }

  else if ( type == "Similarity2DTransform" )

    {

    typedef itk::Similarity2DTransform<PrecisionType> TransformType;

    auto transformPointer = TransformType::New();

    typedef itk::Transform<PrecisionType,Dimension,Dimension> TransformBaseType;

    typedef typename TransformBaseType::Pointer               TransformBasePointerType;

    TransformBasePointerType basePointer = dynamic_cast<TransformBaseType *>( transformPointer.GetPointer() );

    AntsTransform<TransformBaseType> outTransform = { basePointer };

    return outTransform;

    }

  else if ( type == "CenteredSimilarity2DTransform" )

    {

    typedef itk::CenteredSimilarity2DTransform<PrecisionType> TransformType;

    auto transformPointer = TransformType::New();

    typedef itk::Transform<PrecisionType,Dimension,Dimension> TransformBaseType;

    typedef typename TransformBaseType::Pointer               TransformBasePointerType;

    TransformBasePointerType basePointer = dynamic_cast<TransformBaseType *>( transformPointer.GetPointer() );

    AntsTransform<TransformBaseType> outTransform = { basePointer };

    return outTransform;

    }

    typedef itk::AffineTransform<PrecisionType,Dimension> TransformType;

    auto transformPointer = TransformType::New();

    typedef itk::Transform<PrecisionType,Dimension,Dimension> TransformBaseType;

    typedef typename TransformBaseType::Pointer               TransformBasePointerType;

    TransformBasePointerType basePointer = dynamic_cast<TransformBaseType *>( transformPointer.GetPointer() );

    AntsTransform<TransformBaseType> outTransform = { basePointer };

    return outTransform;

}

void local_readTransform(nb::module_ &m)

{

    m.def("newAntsTransformF2", &newAntsTransform<float, 2>);

    m.def("newAntsTransformF3", &newAntsTransform<float, 3>);

    m.def("newAntsTransformD2", &newAntsTransform<double,2>);

    m.def("newAntsTransformD3", &newAntsTransform<double,3>);

    m.def("getTransformDimensionFromFile", &getTransformDimensionFromFile);

    m.def("getTransformNameFromFile", &getTransformNameFromFile);

}