/*=========================================================================
 *
 *  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 "itkHDF5ImageIOFactory.h"
#include "itkIOTestHelper.h"
#include "itkPipelineMonitorImageFilter.h"
#include "itkStreamingImageFilter.h"
#include "itkGenerateImageSource.h"
#include "itkImageRegionIteratorWithIndex.h"
#include "itkImageDuplicator.h"
#include "itkMath.h"

namespace itk
{
/**
 * \class DemoImageSource
 *
 * \brief Streamable process that will generate image regions from the write requests
 *
 * We do not directly allocate the Image, because, as an Image is a data and is not streamable,
 * the writer would write the image in one pass, without streaming.
 *
 * Instead, we need to set a streamable source process as the writer input.
 * This source process, 'DemoImageSource', will allocate a region of the image
 * (and set pixels values) on the fly, based on the information
 * received from the writer requests.
 */
template <class TOutputImage>
class DemoImageSource : public GenerateImageSource<TOutputImage>
{
public:
  ITK_DISALLOW_COPY_AND_MOVE(DemoImageSource);

  /** Standard class type aliases. */
  using Self = DemoImageSource;
  using Superclass = DemoImageSource<TOutputImage>;
  using Pointer = SmartPointer<Self>;

  /** Method for creation through the object factory. */
  itkNewMacro(Self);

  /** \see LightObject::GetNameOfClass() */
  itkOverrideGetNameOfClassMacro(DemoImageSource);

  /** Set the value to fill the image. */
  itkSetMacro(Value, typename TOutputImage::PixelType);

protected:
  DemoImageSource()
    : m_Value(typename TOutputImage::PixelType{})
  {}
  ~DemoImageSource() override = default;

  /** Does the real work. */
  void
  GenerateData() override
  {
    TOutputImage * out = this->GetOutput();
    out->SetBufferedRegion(out->GetRequestedRegion());
    out->Allocate();
    itk::ImageRegionIteratorWithIndex<TOutputImage> it(out, out->GetRequestedRegion());
    for (it.GoToBegin(); !it.IsAtEnd(); ++it)
    {
      typename TOutputImage::IndexType       idx = it.GetIndex();
      const typename TOutputImage::PixelType pixel(idx[2] * 100 + idx[1] * 10 + idx[0]);
      it.Set(pixel);
    }
  }

private:
  typename TOutputImage::PixelType m_Value;
};

} // namespace itk

template <typename TPixel>
int
HDF5ReadWriteTest2(const char * fileName)
{
  { // START an isolation block for the streaming readers and writers.
    // For streaming to work, the itk::ImageFileReader / itk::ImageFileWriter
    // pipeline and the itk::StreamingImageFilter will maintain handles to open HDF5 files
    // until their destructors run.

    using ImageType = itk::Image<TPixel, 3>;

    // Create a source object (in this case a constant image).
    typename ImageType::SizeType size;
    size[2] = 5;
    size[1] = 5;
    size[0] = 5;
    const typename itk::DemoImageSource<ImageType>::Pointer imageSource = itk::DemoImageSource<ImageType>::New();
    imageSource->SetValue(static_cast<TPixel>(23)); // Not used.
    imageSource->SetSize(size);

    // Write image with streaming.
    using WriterType = itk::ImageFileWriter<ImageType>;
    auto writer = WriterType::New();
    using MonitorFilterType = itk::PipelineMonitorImageFilter<ImageType>;
    auto writerMonitor = MonitorFilterType::New();
    writerMonitor->SetInput(imageSource->GetOutput());
    writer->SetFileName(fileName);
    writer->SetInput(writerMonitor->GetOutput());
    writer->SetNumberOfStreamDivisions(5);
    try
    {
      writer->Write();
    }
    catch (const itk::ExceptionObject & err)
    {
      std::cout << "itkHDF5ImageIOTest" << std::endl << "Exception Object caught: " << std::endl << err << std::endl;
      return EXIT_FAILURE;
    }

    // Check streaming regions.
    if (!writerMonitor->VerifyInputFilterExecutedStreaming(5))
    {
      return EXIT_FAILURE;
    }
    typename ImageType::RegionType expectedRegion;
    expectedRegion.SetIndex(0, 0);
    expectedRegion.SetIndex(1, 0);
    expectedRegion.SetSize(0, 5);
    expectedRegion.SetSize(1, 5);
    expectedRegion.SetSize(2, 1);
    typename MonitorFilterType::RegionVectorType writerRegionVector = writerMonitor->GetUpdatedBufferedRegions();
    for (typename ImageType::RegionType::IndexValueType iRegion = 0; iRegion < 5; ++iRegion)
    {
      expectedRegion.SetIndex(2, iRegion);
      if (writerRegionVector[iRegion] != expectedRegion)
      {
        std::cout << "Written image region number " << iRegion << " :" << writerRegionVector[iRegion]
                  << " doesn't match expected one: " << expectedRegion << std::endl;
        return EXIT_FAILURE;
      }
    }

    // Force writer close.
    writer = typename WriterType::Pointer();

    // Read image with streaming.
    using ReaderType = itk::ImageFileReader<ImageType>;
    auto reader = ReaderType::New();
    reader->SetFileName(fileName);
    reader->SetUseStreaming(true);
    auto readerMonitor = MonitorFilterType::New();
    readerMonitor->SetInput(reader->GetOutput());
    using StreamingFilter = itk::StreamingImageFilter<ImageType, ImageType>;
    auto streamer = StreamingFilter::New();
    streamer->SetInput(readerMonitor->GetOutput());
    streamer->SetNumberOfStreamDivisions(5);
    typename ImageType::Pointer image;
    try
    {
      streamer->Update();
    }
    catch (const itk::ExceptionObject & err)
    {
      std::cout << "itkHDF5ImageIOTest" << std::endl << "Exception Object caught: " << std::endl << err << std::endl;
      return EXIT_FAILURE;
    }
    image = streamer->GetOutput();

    // Check the largest possible and buffered regions.
    expectedRegion.SetIndex(0, 0);
    expectedRegion.SetIndex(1, 0);
    expectedRegion.SetIndex(2, 0);
    expectedRegion.SetSize(0, 5);
    expectedRegion.SetSize(1, 5);
    expectedRegion.SetSize(2, 5);
    if (image->GetLargestPossibleRegion() != expectedRegion)
    {
      std::cout << "Read image largest possible region: " << image->GetLargestPossibleRegion()
                << " doesn't match expected one: " << expectedRegion << std::endl;
    }
    if (image->GetBufferedRegion() != expectedRegion)
    {
      std::cout << "Read image buffered region: " << image->GetBufferedRegion()
                << " doesn't match expected one: " << expectedRegion << std::endl;
    }

    // Check image pixel values.
    itk::ImageRegionIterator<ImageType> it(image, image->GetLargestPossibleRegion());
    typename ImageType::IndexType       idx;
    for (it.GoToBegin(); !it.IsAtEnd(); ++it)
    {
      idx = it.ComputeIndex();
      const TPixel origValue(idx[2] * 100 + idx[1] * 10 + idx[0]);
      if (itk::Math::NotAlmostEquals(it.Get(), origValue))
      {
        std::cout << "Original Pixel (" << origValue << ") doesn't match read-in Pixel (" << it.Get() << ')'
                  << std::endl;
        return EXIT_FAILURE;
      }
    }

    // Check number of streaming regions.
    if (!readerMonitor->VerifyInputFilterExecutedStreaming(5))
    {
      return EXIT_FAILURE;
    }

    // Check streaming regions.
    typename MonitorFilterType::RegionVectorType readerRegionVector = readerMonitor->GetUpdatedBufferedRegions();
    expectedRegion.SetIndex(0, 0);
    expectedRegion.SetIndex(1, 0);
    expectedRegion.SetSize(0, 5);
    expectedRegion.SetSize(1, 5);
    expectedRegion.SetSize(2, 1);
    for (typename ImageType::RegionType::IndexValueType iRegion = 0; iRegion < 5; ++iRegion)
    {
      expectedRegion.SetIndex(2, iRegion);
      if (readerRegionVector[iRegion] != expectedRegion)
      {
        std::cout << "Read image region number " << iRegion << " :" << readerRegionVector[iRegion]
                  << " doesn't match expected one: " << expectedRegion << std::endl;
        return EXIT_FAILURE;
      }
    }
  } // END the streaming isolation block to force destructors and close files
  // Clean working directory.
  itk::IOTestHelper::Remove(fileName); // Must ensure that the file is closed before removing the file.
  return EXIT_SUCCESS;
}

int
itkHDF5ImageIOStreamingReadWriteTest(int argc, char * argv[])
{
  std::string prefix("");
  if (argc > 1)
  {
    prefix = *++argv;
    --argc;
    itksys::SystemTools::ChangeDirectory(prefix);
  }
  itk::ObjectFactoryBase::RegisterFactory(itk::HDF5ImageIOFactory::New());

  int result(0);
  result += HDF5ReadWriteTest2<unsigned char>("StreamingUCharImage.hdf5");
  result += HDF5ReadWriteTest2<float>("StreamingFloatImage.hdf5");
  result += HDF5ReadWriteTest2<itk::RGBPixel<unsigned char>>("StreamingRGBImage.hdf5");
  return result != 0;
}
