/*=========================================================================
 *
 *  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 <iostream>
#include <fstream>

#include "itkOpenCVVideoIO.h"
#include "itkImportImageFilter.h"
#include "itkImageFileWriter.h"
#include "itkOpenCVVideoIOFactory.h"

// ITK type alias
using PixelType = itk::RGBPixel<char>;
using ImportFilterType = itk::ImportImageFilter<PixelType, 2>;
using ImageType = itk::Image<PixelType, 2>;
using WriterType = itk::ImageFileWriter<ImageType>;
using FrameOffsetType = itk::SizeValueType;

//
// Utility function to get an ITK image from an void* buffer
//
ImageType::Pointer
itkImageFromBuffer(itk::OpenCVVideoIO::Pointer opencvIO, void * buffer, size_t bufferSize)
{
  // Set up for incoming image
  ImageType::RegionType region;
  ImageType::SizeType   size;
  ImageType::IndexType  start;

  size[0] = opencvIO->GetDimensions(0);
  size[1] = opencvIO->GetDimensions(1);
  start.Fill(0);
  region.SetIndex(start);
  region.SetSize(size);
  ImageType::PointType   origin;
  ImageType::SpacingType space;
  origin.Fill(0.0);
  space.Fill(1.0); // May need fixing

  // Use itkImportImageFilter to create an ITK image
  auto importFilter = ImportFilterType::New();
  importFilter->SetRegion(region);
  importFilter->SetOrigin(origin);
  importFilter->SetSpacing(space);
  importFilter->SetImportPointer(reinterpret_cast<PixelType *>(buffer), bufferSize, false);
  importFilter->Update();
  ImageType::Pointer frame = importFilter->GetOutput();

  return frame;
}

//
// Utility function for comparing output of Read buffer to OpenCV image
//
// Note: opencvIO should already have called ReadImageInformation
//
bool
readCorrectly(itk::OpenCVVideoIO::Pointer opencvIO, CvCapture * capture, FrameOffsetType frameNumber)
{
  bool ret = true;
  // Check meta data
  for (unsigned int i = 0; i < ImageType::ImageDimension; ++i)
  {
    if (opencvIO->GetSpacing(i) != 1.0)
    {
      std::cerr << "Frame Spacing not set correctly" << std::endl;
      ret = false;
    }
    if (opencvIO->GetOrigin(i) != 0.0)
    {
      std::cerr << "Frame Origin not set correctly" << std::endl;
      ret = false;
    }
    if (opencvIO->GetDirection(i) != opencvIO->GetDefaultDirection(i))
    {
      std::cerr << "Frame Direction not set correctly" << std::endl;
      ret = false;
    }
  }

  // Set up the buffer for the frame data
  itk::SizeValueType bufferSize = opencvIO->GetImageSizeInBytes();
  PixelType *        buffer = new PixelType[bufferSize];

  // Read the frame data
  opencvIO->Read(static_cast<void *>(buffer));

  // Open the frame directly with OpenCV
  IplImage * cvFrameBGR = cvQueryFrame(capture);
  IplImage * cvFrameRGB = cvCreateImage(
    cvSize(opencvIO->GetDimensions(0), opencvIO->GetDimensions(1)), IPL_DEPTH_8U, opencvIO->GetNumberOfComponents());
  cvCvtColor(cvFrameBGR, cvFrameRGB, CV_BGR2RGB);

  // Make sure buffers are same sized
  if (cvFrameRGB->imageSize != static_cast<int>(bufferSize))
  {
    std::cerr << "Frame buffer sizes don't match. Got: " << bufferSize << ", Expected: " << cvFrameRGB->imageSize
              << std::endl;
    ret = false;
  }

  // Compare buffer contents
  if (memcmp(reinterpret_cast<void *>(buffer), reinterpret_cast<void *>(cvFrameRGB->imageData), bufferSize))
  {
    std::cerr << "Frame buffers don't match for frame " << frameNumber << std::endl;
    ret = false;
  }

  delete[] buffer;
  // Return
  cvReleaseImage(&cvFrameRGB);
  return ret;
}

//
// Utility function to compare two videos frame by frame
//
bool
videosMatch(char * file1, char * file2)
{
  itk::OpenCVVideoIO::Pointer io1 = itk::OpenCVVideoIO::New();
  itk::OpenCVVideoIO::Pointer io2 = itk::OpenCVVideoIO::New();

  io1->SetFileName(file1);
  io2->SetFileName(file2);

  // Make sure files can be read
  if (!io1->CanReadFile(file1) || !io2->CanReadFile(file2))
  {
    std::cerr << "Cannot read specified files" << std::endl;
    return false;
  }

  // Read in the file information for both
  io1->ReadImageInformation();
  io2->ReadImageInformation();

  // Make sure image info matches
  double e = 0.0001;
  if (io1->GetFrameTotal() != io2->GetFrameTotal() || io1->GetDimensions(0) != io2->GetDimensions(0) ||
      io1->GetDimensions(1) != io2->GetDimensions(1) || io1->GetFramesPerSecond() + e < io2->GetFramesPerSecond() ||
      io1->GetFramesPerSecond() - e > io2->GetFramesPerSecond() ||
      io1->GetNumberOfComponents() != io2->GetNumberOfComponents())
  {

    std::cerr << "Frame information doesn't match" << std::endl;
    std::cerr << "  FrameTotal: " << io1->GetFrameTotal() << ", " << io2->GetFrameTotal() << std::endl;
    std::cerr << "  Width: " << io1->GetDimensions(0) << ", " << io2->GetDimensions(0) << std::endl;
    std::cerr << "  Height: " << io1->GetDimensions(1) << ", " << io2->GetDimensions(1) << std::endl;
    std::cerr << "  FpS: " << io1->GetFramesPerSecond() << ", " << io2->GetFramesPerSecond() << std::endl;
    std::cerr << "  NChannels: " << io1->GetNumberOfComponents() << ", " << io2->GetNumberOfComponents() << std::endl;

    return false;
  }

  // Loop through each frame and compare the buffer for exact match
  itk::SizeValueType bufferSize = io1->GetImageSizeInBytes();
  PixelType *        buffer1 = new PixelType[bufferSize];
  PixelType *        buffer2 = new PixelType[bufferSize];
  for (unsigned int i = 0; i < io1->GetFrameTotal(); ++i)
  {
    io1->Read(reinterpret_cast<void *>(buffer1));
    io2->Read(reinterpret_cast<void *>(buffer2));
    if (memcmp(reinterpret_cast<void *>(buffer1), reinterpret_cast<void *>(buffer2), bufferSize))
    {
      std::cerr << "Frame buffers don't match for frame " << i << std::endl;
      return false;
    }
  }
  delete[] buffer1;
  delete[] buffer2;
  // Close the readers
  io1->FinishReadingOrWriting();
  io2->FinishReadingOrWriting();

  // return result
  return true;
}

//
// This tests all of the functionality of the OpenCVVideoIO
//
// Usage: [Video Input] [Non-Video Input] [Video Output] [Width] [Height]
//            [Num Frames] [FpS]

int
test_OpenCVVideoIO(char *          input,
                   char *          nonVideoInput,
                   char *          output,
                   char *          cameraOutput,
                   unsigned int    inWidth,
                   unsigned int    inHeight,
                   FrameOffsetType inNumFrames,
                   double          inFpS)
{

  int ret = EXIT_SUCCESS;

  // Create the VideoIO
  itk::OpenCVVideoIO::Pointer opencvIO = itk::OpenCVVideoIO::New();

  //
  // CanReadFile
  //
  std::cout << "OpenCVVideoIO::CanReadFile..." << std::endl;

  // Test CanReadFile on good file
  if (!opencvIO->CanReadFile(input))
  {
    std::cerr << "Could not read " << input << std::endl;
    ret = EXIT_FAILURE;
  }

  // Test CanReadFile on non-existant file
  std::string nonExistantFile = "Bad/Path/To/Nothing";
  if (opencvIO->CanReadFile(nonExistantFile.c_str()))
  {
    std::cerr << "Should have failed to open \"" << nonExistantFile << '"' << std::endl;
    ret = EXIT_FAILURE;
  }

  // Test CanReadFile on non-video file
  if (opencvIO->CanReadFile(nonVideoInput))
  {
    std::cerr << "Should have failed to open \"" << nonVideoInput << '"' << std::endl;
    ret = EXIT_FAILURE;
  }

  //
  // ReadImageInformation
  //
  std::cout << "OpenCVVideoIO::ReadImageInformation..." << std::endl;

  opencvIO->SetFileName(input);
  opencvIO->ReadImageInformation();
  bool              infoSet = true;
  std::stringstream paramMessage;
  if (opencvIO->GetDimensions(0) != inWidth)
  {
    infoSet = false;
    paramMessage << "Width mismatch: (expected) " << inWidth << " != (got) " << opencvIO->GetDimensions(0) << std::endl;
  }
  if (opencvIO->GetDimensions(1) != inHeight)
  {
    infoSet = false;
    paramMessage << "Height mismatch: (expected) " << inHeight << " != (got) " << opencvIO->GetDimensions(1)
                 << std::endl;
  }
  double epsilon = 0.0001;
  if (opencvIO->GetFramesPerSecond() < inFpS - epsilon || opencvIO->GetFramesPerSecond() > inFpS + epsilon)
  {
    infoSet = false;
    paramMessage << "FpS mismatch: (expected) " << inFpS << " != (got) " << opencvIO->GetFramesPerSecond() << std::endl;
  }
  if (opencvIO->GetFrameTotal() != inNumFrames)
  {
    infoSet = false;
    paramMessage << "FrameTotal mismatch: (expected) " << inNumFrames << " != (got) " << opencvIO->GetFrameTotal()
                 << std::endl;
  }

  if (!infoSet)
  {
    std::cerr << paramMessage.str();
    ret = EXIT_FAILURE;
  }

  //
  // Read
  //
  std::cout << "OpenCVVideoIO::Read..." << std::endl;
  std::cout << "Comparing all " << opencvIO->GetFrameTotal() << " frames" << std::endl;

  // Set up OpenCV capture
  CvCapture * capture = cvCaptureFromFile(opencvIO->GetFileName());
  // Loop through all frames
  for (FrameOffsetType i = 0; i * opencvIO->GetIFrameInterval() < opencvIO->GetFrameTotal(); ++i)
  {
    if (!readCorrectly(opencvIO, capture, i * opencvIO->GetIFrameInterval()))
    {
      std::cerr << "Failed to read frame " << i * opencvIO->GetIFrameInterval() << " correctly" << std::endl;
      ret = EXIT_FAILURE;
      break;
    }
  }

  // Release capture
  cvReleaseCapture(&capture);

  //
  // SetNextFrameToRead
  //
  std::cout << "OpenCVVideoIO::SetNextFrameToRead" << std::endl;

  // Set up the buffer for the frame data so Read can be called
  itk::SizeValueType bufferSize = opencvIO->GetImageSizeInBytes();
  PixelType *        buffer = new PixelType[bufferSize];

  // try seeking to an I-Frame
  // seekFrame is 0-based index of the frame to be captured next
  FrameOffsetType seekFrame = opencvIO->GetIFrameInterval() - 1;
  if (!opencvIO->SetNextFrameToRead(seekFrame))
  {
    std::cerr << "Failed to seek to second I-Frame..." << std::endl;
    ret = EXIT_FAILURE;
  }

  // Read the frame data which updates the current frame correctly
  opencvIO->Read(static_cast<void *>(buffer));

  // GetCurrentFrame() returns 0-based index of the frame to be captured next
  if (opencvIO->GetCurrentFrame() - 1 != seekFrame)
  {
    std::cerr << "Seek to I-Frame didn't end up in the right place" << std::endl;
    ret = EXIT_FAILURE;
  }

  // If there are I-Frame intervals, check behavior
  if (opencvIO->GetIFrameInterval() > 1)
  {

    // Try seeking in-between I-Frames
    seekFrame = opencvIO->GetIFrameInterval() / 2;
    if (!opencvIO->SetNextFrameToRead(seekFrame))
    {
      std::cerr << "Failed to seek between I-Frames" << std::endl;
      ret = EXIT_FAILURE;
    }
    opencvIO->Read(static_cast<void *>(buffer));
    if (opencvIO->GetCurrentFrame() != opencvIO->GetIFrameInterval())
    {
      std::cerr << "Seek between I-Frames didn't end up in the right place" << std::endl;
      ret = EXIT_FAILURE;
    }

    delete[] buffer;
    // try seeking past last I-Frame
    seekFrame = opencvIO->GetLastIFrame() + 1;
    if (opencvIO->SetNextFrameToRead(seekFrame))
    {
      std::cerr << "Did no fail when seeking past the last I-Frame" << std::endl;
      ret = EXIT_FAILURE;
    }
  }

  // Save the current parameters
  double       fps = opencvIO->GetFramesPerSecond();
  unsigned int width = opencvIO->GetDimensions(0);
  unsigned int height = opencvIO->GetDimensions(1);
  const char * fourCC = "MP42";
  unsigned int nChannels = opencvIO->GetNumberOfComponents();

  // Reset the VideoIO
  opencvIO->FinishReadingOrWriting();

  //
  // Test reading from camera -- If webcam 0 can be opened, it will, otherwise
  // this will be skipped
  //

  // Check to see if camera is available
  if (opencvIO->CanReadCamera(0))
  {

    std::cout << "OpenCVVideoIO::Read (from camera)..." << std::endl;

    // Set the reader to use the camera
    opencvIO->SetReadFromCamera();

    // Get information from the camera
    try
    {
      opencvIO->ReadImageInformation();
    }
    catch (const itk::ExceptionObject & e)
    {
      std::cerr << "Could not read information from the camera" << std::endl;
      std::cerr << "If a camera is present, this test requires it to be on with OpenCV 2" << std::endl;
      std::cerr << e << std::endl;
      return EXIT_FAILURE;
    }

    // Set up buffer for camera
    itk::SizeValueType camBufferSize = opencvIO->GetImageSizeInBytes();
    PixelType *        camBuffer = new PixelType[camBufferSize];

    // Read from the camera
    try
    {
      opencvIO->Read(reinterpret_cast<void *>(camBuffer));
    }
    catch (const itk::ExceptionObject & e)
    {
      std::cerr << "Could not read from the camera" << std::endl;
      std::cerr << e << std::endl;
      ret = EXIT_FAILURE;
    }

    // Get an ITK image from the camera's frame
    ImageType::Pointer cameraFrame = itkImageFromBuffer(opencvIO, camBuffer, camBufferSize);
    delete[] camBuffer;

    // Write out the ITK image -- DEBUG
    auto writer = WriterType::New();
    writer->SetFileName(cameraOutput);
    writer->SetInput(cameraFrame);
    writer->Update();

    // Overwrite the file right away so we're not saving pictures of the tester!
    std::ofstream fs;
    fs.open(cameraOutput);
    fs << "EMPTY... deleted picture from webcam\n";
    fs.close();

    // Finish reading
    opencvIO->FinishReadingOrWriting();
  }

  //
  // Test Writing
  //

  //
  // SetWriterParameters
  //
  std::cout << "OpenCVVIdeoIO::SetWriterParameters..." << std::endl;

  // Reset the saved parameters
  std::vector<itk::SizeValueType> size;
  size.push_back(width);
  size.push_back(height);
  opencvIO->SetWriterParameters(fps, size, fourCC, nChannels, itk::IOComponentEnum::UCHAR);

  // Make sure they set correctly
  if (opencvIO->GetFramesPerSecond() != fps || opencvIO->GetDimensions(0) != width ||
      opencvIO->GetDimensions(1) != height || opencvIO->GetNumberOfComponents() != nChannels)
  {
    std::cerr << "Didn't set writer parmeters correctly" << std::endl;
    ret = EXIT_FAILURE;
  }

  //
  // CanWriteFile
  //
  std::cout << "OpenCVVideoIO::CanWriteFile..." << std::endl;

  // Test CanWriteFile on good filename
  if (!opencvIO->CanWriteFile(output))
  {
    std::cerr << "CanWriteFile didn't return true correctly" << std::endl;
    ret = EXIT_FAILURE;
  }

  // Test CanWriteFile on bad filename
  if (opencvIO->CanWriteFile("asdfasdfasdf"))
  {
    std::cerr << "CanWriteFile should have failed for bad filename" << std::endl;
    ret = EXIT_FAILURE;
  }

  //
  // Write
  //
  std::cout << "OpenCVVIdeoIO::Write..." << std::endl;

  // Set output filename
  opencvIO->SetFileName(output);

  // Set up a second VideoIO to read while we're writing
  itk::OpenCVVideoIO::Pointer opencvIO2 = itk::OpenCVVideoIO::New();
  opencvIO2->SetFileName(input);
  opencvIO2->ReadImageInformation();
  // Loop through all frames to read with opencvIO2 and write with opencvIO
  for (unsigned int i = 0; i * opencvIO2->GetIFrameInterval() < inNumFrames; ++i)
  {
    // Set up a buffer to read to
    itk::SizeValueType bufferSizeT = opencvIO2->GetImageSizeInBytes();
    PixelType *        bufferT = new PixelType[bufferSizeT];

    // Read into the buffer
    opencvIO2->Read(static_cast<void *>(bufferT));

    // Write out the frame from the buffer
    opencvIO->Write(static_cast<void *>(bufferT));
    delete[] bufferT;
  }

  // Finish writing
  opencvIO2->FinishReadingOrWriting();
  opencvIO->FinishReadingOrWriting();

  // DEBUG -- Don't do this for now, need a better comparison method
  // Compare input and output videos to make sure they are identical
  // if (!videosMatch(input, output))
  //  {
  //  std::cerr << "Written video does not match input video" << std::endl;
  //  ret = EXIT_FAILURE;
  //  }

  // DEBUG
  // std::cout << "PIM1 = " << CV_FOURCC('P','I','M','1') << std::endl;
  // std::cout << "MJPG = " << CV_FOURCC('M','J','P','G') << std::endl;
  // std::cout << "MP42 = " << CV_FOURCC('M', 'P', '4', '2') << std::endl;
  // std::cout << "DIV3 = " << CV_FOURCC('D', 'I', 'V', '3') << std::endl;
  // std::cout << "DIVX = " << CV_FOURCC('D', 'I', 'V', 'X') << std::endl;
  // std::cout << "U263 = " << CV_FOURCC('U', '2', '6', '3') << std::endl;
  // std::cout << "I263 = " << CV_FOURCC('I', '2', '6', '3') << std::endl;
  // std::cout << "FLV1 = " << CV_FOURCC('F', 'L', 'V', '1') << std::endl;

  std::cout << "Done !" << std::endl;
  return ret;
}

int
itkOpenCVVideoIOTest(int argc, char * argv[])
{
  if (argc != 9)
  {
    std::cerr << "Usage: [Video Input] [Non-Video Input] [Video Output] [Webcam Output] "
                 "[Width] [Height] [Num Frames] [FpS]"
              << std::endl;
    return EXIT_FAILURE;
  }

  itk::ObjectFactoryBase::RegisterFactory(itk::OpenCVVideoIOFactory::New());

  return test_OpenCVVideoIO(
    argv[1], argv[2], argv[3], argv[4], std::stoi(argv[5]), std::stoi(argv[6]), std::stoi(argv[7]), std::stod(argv[8]));
}