//=========================================================================
//  Copyright (c) Kitware, Inc.
//  All rights reserved.
//  See LICENSE.txt 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 notice for more information.
//=========================================================================

#include <iostream>

#include "itkImage.h"
#include "itkImageFileReader.h"
#include "itkImageIOBase.h"
#include "itkImageRegionIterator.h"
#include "itkNiftiImageIO.h"

template<class TImage>
int ReadImage(const char* fileName, typename TImage::Pointer image)
{
  using ImageType = TImage;
  using ImageReaderType = itk::ImageFileReader<ImageType>;

  typename ImageReaderType::Pointer reader = ImageReaderType::New();
  reader->SetFileName(fileName);

  try
  {
    reader->Update();
  }
  catch (itk::ExceptionObject& e)
  {
    std::cerr << e.what() << std::endl;
    return EXIT_FAILURE;
  }

  image->Graft(reader->GetOutput());

  return EXIT_SUCCESS;
}

template<unsigned int VDimension>
int ReadScalarImage(const char* inputFileName,
  const itk::ImageIOBase::IOComponentType componentType)
{
  switch (componentType)
  {
    default:
    case itk::ImageIOBase::UNKNOWNCOMPONENTTYPE:
      std::cerr << "Unknown and unsupported component type!" << std::endl;
      return EXIT_FAILURE;

    case itk::ImageIOBase::UCHAR:
    {
      using PixelType = unsigned char;
      using ImageType = itk::Image<PixelType, VDimension>;

      typename ImageType::Pointer image = ImageType::New();

      if (ReadImage<ImageType>(inputFileName, image) == EXIT_FAILURE)
      {
        return EXIT_FAILURE;
      }

      std::cout << image << std::endl;
      break;
    }

    case itk::ImageIOBase::CHAR:
    {
      using PixelType = char;
      using ImageType = itk::Image<PixelType, VDimension>;

      typename ImageType::Pointer image = ImageType::New();

      if (ReadImage<ImageType>(inputFileName, image) == EXIT_FAILURE)
      {
        return EXIT_FAILURE;
      }

      std::cout << image << std::endl;
      break;
    }

    case itk::ImageIOBase::USHORT:
    {
      using PixelType = unsigned short;
      using ImageType = itk::Image<PixelType, VDimension>;

      typename ImageType::Pointer image = ImageType::New();

      if (ReadImage<ImageType>(inputFileName, image) == EXIT_FAILURE)
      {
        return EXIT_FAILURE;
      }

      std::cout << image << std::endl;
      break;
    }

    case itk::ImageIOBase::SHORT:
    {
      using PixelType = short;
      using ImageType = itk::Image<PixelType, VDimension>;

      typename ImageType::Pointer image = ImageType::New();

      if (ReadImage<ImageType>(inputFileName, image) == EXIT_FAILURE)
      {
        return EXIT_FAILURE;
      }

      std::cout << image << std::endl;
      break;
    }

    case itk::ImageIOBase::UINT:
    {
      using PixelType = unsigned int;
      using ImageType = itk::Image<PixelType, VDimension>;

      typename ImageType::Pointer image = ImageType::New();

      if (ReadImage<ImageType>(inputFileName, image) == EXIT_FAILURE)
      {
        return EXIT_FAILURE;
      }

      std::cout << image << std::endl;
      break;
    }

    case itk::ImageIOBase::INT:
    {
      using PixelType = int;
      using ImageType = itk::Image<PixelType, VDimension>;

      typename ImageType::Pointer image = ImageType::New();

      if (ReadImage<ImageType>(inputFileName, image) == EXIT_FAILURE)
      {
        return EXIT_FAILURE;
      }

      std::cout << image << std::endl;
      break;
    }

    case itk::ImageIOBase::ULONG:
    {
      using PixelType = unsigned long;
      using ImageType = itk::Image<PixelType, VDimension>;

      typename ImageType::Pointer image = ImageType::New();

      if (ReadImage<ImageType>(inputFileName, image) == EXIT_FAILURE)
      {
        return EXIT_FAILURE;
      }

      std::cout << image << std::endl;
      break;
    }

    case itk::ImageIOBase::LONG:
    {
      using PixelType = long;
      using ImageType = itk::Image<PixelType, VDimension>;

      typename ImageType::Pointer image = ImageType::New();

      if (ReadImage<ImageType>(inputFileName, image) == EXIT_FAILURE)
      {
        return EXIT_FAILURE;
      }

      std::cout << image << std::endl;
      break;
    }

    case itk::ImageIOBase::FLOAT:
    {
      using PixelType = float;
      using ImageType = itk::Image<PixelType, VDimension>;

      typename ImageType::Pointer image = ImageType::New();

      if (ReadImage<ImageType>(inputFileName, image) == EXIT_FAILURE)
      {
        return EXIT_FAILURE;
      }

      std::cout << image << std::endl;
      break;
    }

    case itk::ImageIOBase::DOUBLE:
    {
      using PixelType = double;
      using ImageType = itk::Image<PixelType, VDimension>;

      typename ImageType::Pointer image = ImageType::New();

      if (ReadImage<ImageType>(inputFileName, image) == EXIT_FAILURE)
      {
        return EXIT_FAILURE;
      }

      std::cout << image << std::endl;
      break;
    }
  }
  return EXIT_SUCCESS;
}

int UnitTestImport(int argc, char* argv[])
{
  (void)argc;
  (void)argv;

  std::string image_file =
    AEVA_DATA_DIR "/itk/1.3.12.2.1107.5.2.19.45406.2014120211045428131243076.0.0.0_small.nii";
  if (argc > 1)
  {
    image_file = argv[1];
  }
  std::cout << "File is \"" << image_file << "\"\n";

  itk::ImageIOBase::Pointer imageIO = itk::ImageIOFactory::CreateImageIO(
    image_file.c_str(), itk::ImageIOFactory::FileModeEnum::ReadMode);
  if (imageIO == nullptr)
  {
    std::cerr << "Unsupported format for " << image_file << "\n";
    return EXIT_FAILURE;
  }
  imageIO->SetFileName(image_file.c_str());
  imageIO->ReadImageInformation();
  const unsigned int imageDimension = imageIO->GetNumberOfDimensions();
  const auto componentType = imageIO->GetComponentType();

  switch (imageIO->GetPixelType())
  {
    case itk::ImageIOBase::SCALAR:
    {
      if (imageDimension == 2)
      {
        return ReadScalarImage<2>(image_file.c_str(), componentType);
      }
      if (imageDimension == 3)
      {
        return ReadScalarImage<3>(image_file.c_str(), componentType);
      }
      if (imageDimension == 4)
      {
        return ReadScalarImage<4>(image_file.c_str(), componentType);
      }
    }

    default:
      std::cerr << "not implemented yet!" << std::endl;
      return EXIT_FAILURE;
  }
  return EXIT_SUCCESS;
}
