/*=========================================================================
 *
 *  Copyright NumFOCUS
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *         https://www.apache.org/licenses/LICENSE-2.0.txt
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
 *=========================================================================*/

#include <fstream>
#include "itkImageFileReader.h"
#include "itkImageFileWriter.h"
#include "itkTestingComparisonImageFilter.h"
#include "itkExtractImageFilter.h"
#include "itkPipelineMonitorImageFilter.h"
#include "itkTestingMacros.h"


constexpr unsigned int VDimension{ 3 };
using PixelType = unsigned char;
using ImageType = itk::Image<PixelType, VDimension>;
using ImagePointer = ImageType::Pointer;
using SpacingType = ImageType::SpacingType;

namespace
{


bool
SameImage(ImagePointer testImage, ImagePointer baselineImage)
{
  constexpr PixelType     intensityTolerance{ 5 }; // need this for compression
  constexpr int           radiusTolerance{ 0 };
  constexpr unsigned long numberOfPixelTolerance{ 0 };

  // NOTE ALEX: it looks like this filter does not take the spacing
  // into account, to check later.
  using DiffType = itk::Testing::ComparisonImageFilter<ImageType, ImageType>;
  auto diff = DiffType::New();
  diff->SetValidInput(baselineImage);
  diff->SetTestInput(testImage);
  diff->SetDifferenceThreshold(intensityTolerance);
  diff->SetToleranceRadius(radiusTolerance);
  diff->UpdateLargestPossibleRegion();

  const unsigned long status = diff->GetNumberOfPixelsWithDifferences();

  if (status > numberOfPixelTolerance)
  {
    return false;
  }

  SpacingType testImageSpacing = testImage->GetSpacing();
  SpacingType baselineImageSpacing = baselineImage->GetSpacing();
  // compare spacing
  for (unsigned int i = 0; i < VDimension; ++i)
  {
    if (itk::Math::NotAlmostEquals(testImageSpacing[i], baselineImageSpacing[i]))
    {
      return false;
    }
  }

  return true;
}

// NOTE ALEX: why this function is not a wrapper of the above?
bool
SameImage(std::string testImageFileName, ImagePointer baselineImage)
{
  using ReaderType = itk::ImageFileReader<ImageType>;
  auto readerTestImage = ReaderType::New();
  readerTestImage->SetFileName(testImageFileName);

  // NOTE ALEX: here we suppose the reading went well
  // we should surround the GetOutput() with a try/catch
  return SameImage(readerTestImage->GetOutput(), baselineImage);
}

bool
ActualTest(const std::string & inputFileName,
           const std::string & outputFileNameBase,
           const std::string & outputFileNameExtension,
           bool                streamWriting,
           bool                pasteWriting,
           bool                compressWriting,
           int                 expectException = -1)
{

  std::cout << "Writing Combination: ";
  std::cout << streamWriting << ' ';
  std::cout << pasteWriting << ' ' << compressWriting << std::endl;

  std::ostringstream outputFileNameStream;
  outputFileNameStream << outputFileNameBase << streamWriting;
  outputFileNameStream << pasteWriting << compressWriting;
  outputFileNameStream << '.' << outputFileNameExtension;
  const std::string outputFileName = outputFileNameStream.str();

  std::cout << "Writing to File: " << outputFileName << std::endl;
  constexpr unsigned int numberOfPieces{ 10 };

  // We remove the output file
  // NOTE ALEX: should we check it exists first?
  itksys::SystemTools::RemoveFile(outputFileName);

  using PixelType = unsigned char;
  using ImageType = itk::Image<PixelType, 3>;

  using ReaderType = itk::ImageFileReader<ImageType>;
  using WriterType = itk::ImageFileWriter<ImageType>;

  auto reader = ReaderType::New();
  reader->SetFileName(inputFileName.c_str());
  reader->SetUseStreaming(true);

  // read the region info
  reader->UpdateOutputInformation();

  ImageType::RegionType largestRegion;
  largestRegion = reader->GetOutput()->GetLargestPossibleRegion().GetSize();

  ImageType::IndexType pasteIndex;
  pasteIndex[0] = largestRegion.GetIndex()[0] + largestRegion.GetSize()[0] / 3;
  pasteIndex[1] = largestRegion.GetIndex()[1] + largestRegion.GetSize()[1] / 3;
  pasteIndex[2] = largestRegion.GetIndex()[2] + largestRegion.GetSize()[2] / 3;
  ImageType::SizeType pasteSize;
  pasteSize[0] = largestRegion.GetSize()[0] / 3;
  pasteSize[1] = largestRegion.GetSize()[1] / 3;
  pasteSize[2] = 1;
  const ImageType::RegionType pasteRegion(pasteIndex, pasteSize);

  // TODO: drew, check and save the spacing of the input image here

  // ??
  using MonitorFilter = itk::PipelineMonitorImageFilter<ImageType>;
  auto monitor = MonitorFilter::New();
  monitor->SetInput(reader->GetOutput());

  // Setup the writer
  auto writer = WriterType::New();
  writer->SetFileName(outputFileName);
  writer->SetInput(monitor->GetOutput());

  // create a valid region from the largest
  itk::ImageIORegion ioregion(3);
  itk::ImageIORegionAdaptor<3>::Convert(pasteRegion, ioregion, largestRegion.GetIndex());

  if (streamWriting)
  {
    writer->SetNumberOfStreamDivisions(numberOfPieces);
  }

  if (pasteWriting)
  {
    writer->SetIORegion(ioregion);
  }

  writer->SetUseCompression(compressWriting);

  try
  {

    writer->Update();
  }
  catch (const itk::ExceptionObject & err)
  {
    if (expectException == -1 || expectException == 1)
    {
      std::cout << "Expected ExceptionObject caught !" << std::endl;
      std::cout << err << std::endl;
      std::cout << "TEST PASSED" << std::endl;
      return EXIT_SUCCESS;
    }

    std::cout << "UnExpected ExceptionObject caught !" << std::endl;
    std::cout << err << std::endl;
    std::cout << "TEST FAILED" << std::endl;
    return EXIT_FAILURE;
  }

  if (expectException == 1)
  {
    std::cout << "Did not get expected exception!" << std::endl;
    std::cout << "TEST FAILED" << std::endl;
    return EXIT_FAILURE;
  }


  // if we didn't have an exception then we should have produced the
  // correct image - This is the TEST !!
  if (pasteWriting)
  {
    using ExtractImageFilterType = itk::ExtractImageFilter<ImageType, ImageType>;
    auto extractBaselineImage = ExtractImageFilterType::New();
    extractBaselineImage->SetDirectionCollapseToSubmatrix();
    extractBaselineImage->SetInput(reader->GetOutput());
    extractBaselineImage->SetExtractionRegion(pasteRegion);

    auto readerTestImage = ReaderType::New();
    readerTestImage->SetFileName(outputFileName);
    auto extractTestImage = ExtractImageFilterType::New();
    extractTestImage->SetDirectionCollapseToSubmatrix();
    extractTestImage->SetInput(readerTestImage->GetOutput());
    extractTestImage->SetExtractionRegion(pasteRegion);

    if (!SameImage(extractTestImage->GetOutput(), extractBaselineImage->GetOutput()))
    {
      std::cout << "Paste regions of images differ" << std::endl;
      std::cout << "TEST FAILED" << std::endl;
      return EXIT_FAILURE;
    }
  }
  else if (!SameImage(outputFileName, reader->GetOutput()))
  {
    std::cout << "Images differ" << std::endl;
    std::cout << "TEST FAILED" << std::endl;
    return EXIT_FAILURE;
  }

  std::cout << "TEST PASSED" << std::endl;
  return EXIT_SUCCESS;
}

} // namespace


int
itkImageFileWriterStreamingPastingCompressingTest1(int argc, char * argv[])
{
  if (argc < 3)
  {
    std::cerr << "Missing parameters." << std::endl;
    std::cerr << "Usage: " << itkNameOfTestExecutableMacro(argv);
    std::cerr << " input outputBase outputExtension [expect exception (0|1)] ..." << std::endl;
    return EXIT_FAILURE;
  }

  int           expectException[8];
  constexpr int expectedExceptionOffset{ 4 };
  for (int i = 0; i < 8; ++i)
  {
    if (argc > i + expectedExceptionOffset)
    {
      expectException[i] = 0;
      if (std::stoi(argv[i + expectedExceptionOffset]) == 1)
      {
        expectException[i] = 1;
      }
    }
    else
    {
      expectException[i] = -1;
    }
  }

  int i = 0;

  int retValue = ActualTest(argv[1], argv[2], argv[3], false, false, false, expectException[i++]);
  retValue = (retValue == EXIT_FAILURE)
               ? EXIT_FAILURE
               : ActualTest(argv[1], argv[2], argv[3], false, false, true, expectException[i++]);
  retValue = (retValue == EXIT_FAILURE)
               ? EXIT_FAILURE
               : ActualTest(argv[1], argv[2], argv[3], false, true, false, expectException[i++]);
  retValue = (retValue == EXIT_FAILURE)
               ? EXIT_FAILURE
               : ActualTest(argv[1], argv[2], argv[3], false, true, true, expectException[i++]);
  retValue = (retValue == EXIT_FAILURE)
               ? EXIT_FAILURE
               : ActualTest(argv[1], argv[2], argv[3], true, false, false, expectException[i++]);
  retValue = (retValue == EXIT_FAILURE)
               ? EXIT_FAILURE
               : ActualTest(argv[1], argv[2], argv[3], true, false, true, expectException[i++]);
  retValue = (retValue == EXIT_FAILURE)
               ? EXIT_FAILURE
               : ActualTest(argv[1], argv[2], argv[3], true, true, false, expectException[i++]);
  retValue = (retValue == EXIT_FAILURE) ? EXIT_FAILURE
                                        : ActualTest(argv[1], argv[2], argv[3], true, true, true, expectException[i++]);


  return retValue;
}
