#if defined(_MSC_VER) 
  #pragma warning ( disable : 4786 )
#endif

#include "itkWellComposedImageFilter.h"

#include "itkImageFileReader.h"
#include "itkImageFileWriter.h"
#include "itkImageRegionIterator.h"

int itkWellComposedImageFilterTest2D( int argc, char **argv )
{
  typedef int                      PixelType;
  typedef itk::Image<PixelType, 2> Image2DType;

  typedef itk::ImageFileReader<Image2DType> Reader2DType;
  Reader2DType::Pointer reader2D = Reader2DType::New();
  reader2D->SetFileName( argv[1] );
  reader2D->Update();

  try
    {
    typedef itk::WellComposedImageFilter<Image2DType> Filter2DType;
    Filter2DType::Pointer filter2D = Filter2DType::New();
    filter2D->SetInput( reader2D->GetOutput() );
    filter2D->SetFullInvariance( true );
    filter2D->SetTotalNumberOfLabels( atoi( argv[3] ) );

    typedef itk::ImageFileWriter<Image2DType> Writer2DType;
    Writer2DType::Pointer writer2D = Writer2DType::New();
    writer2D->SetInput( filter2D->GetOutput() );
    writer2D->SetFileName( argv[2] );
    writer2D->Update();
    }
  catch (itk::ExceptionObject & err)
    { 
    std::cout << "ExceptionObject caught !" << std::endl; 
    std::cout << err << std::endl; 
    return EXIT_FAILURE;
    }
}

int itkWellComposedImageFilterTest3D( int argc, char **argv )
{
  typedef int                      PixelType;
  typedef itk::Image<PixelType, 3> Image3DType;

  typedef itk::ImageFileReader<Image3DType> Reader3DType;
  Reader3DType::Pointer reader3D = Reader3DType::New();
  reader3D->SetFileName( argv[1] );
  reader3D->Update();

  try
    {
    typedef itk::WellComposedImageFilter<Image3DType> Filter3DType;
    Filter3DType::Pointer filter3D = Filter3DType::New();
    filter3D->SetInput( reader3D->GetOutput() );
    filter3D->SetTotalNumberOfLabels( atoi( argv[3] ) );

    typedef itk::ImageFileWriter<Image3DType> Writer3DType;
    Writer3DType::Pointer writer3D = Writer3DType::New();
    writer3D->SetInput( filter3D->GetOutput() );
    writer3D->SetFileName( argv[2] );
    writer3D->Update();
    }
  catch (itk::ExceptionObject & err)
    { 
    std::cout << "ExceptionObject caught !" << std::endl; 
    std::cout << err << std::endl; 
    return EXIT_FAILURE;
    }

  return EXIT_SUCCESS;
}

int main( int argc, char **argv )
{
  if ( argc != 5 )
    {
    std::cout << "Usage: " << argv[0] << " inputImage outputImage numberOfLabels (2 or 3)-D" << std::endl;
    return EXIT_FAILURE;
    }

  if ( atoi( argv[4] ) == 2)
    {
    return itkWellComposedImageFilterTest2D( argc, argv );
    }
  else
    {
    return itkWellComposedImageFilterTest3D( argc, argv );
    }
}