/*=========================================================================
 *
 *  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
 *
 *         http://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 <SimpleITK.h>
#include <memory>
#include <itksys/SystemTools.hxx>

#include "sitkImageCompare.h"

namespace sitk = itk::simple;

void
ImageCompare::NormalizeAndSave(const sitk::Image & input, const std::string & filename)
{
  sitk::Image image = input;

  if (image.GetPixelIDValue() == sitk::sitkLabelUInt8)
  {
    image = sitk::Cast(image, sitk::sitkUInt8);
  }
  else if (image.GetPixelIDValue() == sitk::sitkLabelUInt16)
  {
    image = sitk::Cast(image, sitk::sitkUInt16);
  }
  else if (image.GetPixelIDValue() == sitk::sitkLabelUInt32)
  {
    image = sitk::Cast(image, sitk::sitkUInt32);
  }
  else if (image.GetPixelIDValue() == sitk::sitkLabelUInt64)
  {
    image = sitk::Cast(image, sitk::sitkUInt64);
  }
  else if (image.GetPixelIDValue() == sitk::sitkComplexFloat64 || image.GetPixelIDValue() == sitk::sitkComplexFloat32)
  {
    image = sitk::ComplexToModulus(image);
  }
  else if (image.GetNumberOfComponentsPerPixel() != 1)
  {
    // just use the magnitude for vector images
    image = sitk::VectorMagnitude(image);
  }

  // Extract the center slice of our image
  image = extractSlice(image);

  sitk::Image out = sitk::RescaleIntensity(image, 0, 255);

  out = sitk::Cast(out, sitk::sitkUInt8);
  sitk::WriteImage(out, filename);
}


ImageCompare::ImageCompare()
{
  mTolerance = 0.0;
  mMessage = "";
}

float
ImageCompare::testImages(const itk::simple::Image & testImage,
                         const itk::simple::Image & baselineImage,
                         bool                       reportErrors,
                         const std::string &        baselineImageFilename)
{

  const std::string OutputDir = dataFinder.GetOutputDirectory();

  const std::string shortFilename = itksys::SystemTools::GetFilenameName(baselineImageFilename);

  // verify they have the same size
  if (baselineImage.GetHeight() != testImage.GetHeight() || baselineImage.GetWidth() != testImage.GetWidth() ||
      baselineImage.GetDepth() != testImage.GetDepth())
  {
    mMessage = "ImageCompare: Image dimensions are different";
    return -1;
  }


  // Compute image difference squared
  sitk::Image diffSquared(0, 0, itk::simple::sitkUInt8);
  try
  {

    if (baselineImage.GetPixelID() == sitk::sitkComplexFloat32 ||
        baselineImage.GetPixelID() == sitk::sitkComplexFloat64)
    {

      const sitk::Image diff = sitk::Subtract(testImage, baselineImage);
      // for complex number we multiply the image by it's complex
      // conjugate, this will produce only a real value result
      const sitk::Image conj = sitk::RealAndImaginaryToComplex(sitk::ComplexToReal(diff),
                                                               sitk::Multiply(sitk::ComplexToImaginary(diff), -1.0));
      diffSquared = sitk::ComplexToReal(sitk::Multiply(diff, conj));
    }
    else if (baselineImage.GetNumberOfComponentsPerPixel() > 1)
    {
      const sitk::Image diff = sitk::Subtract(sitk::Cast(testImage, sitk::sitkVectorFloat32),
                                              sitk::Cast(baselineImage, sitk::sitkVectorFloat32));

      // for vector image just do a sum of the components
      diffSquared = sitk::Pow(sitk::VectorIndexSelectionCast(diff, 0), 2.0);
      for (unsigned int i = 1; i < diff.GetNumberOfComponentsPerPixel(); ++i)
      {
        const sitk::Image temp = sitk::Pow(sitk::VectorIndexSelectionCast(diff, i), 2.0);
        diffSquared = sitk::Add(temp, diffSquared);
      }

      diffSquared = sitk::Divide(diffSquared, diff.GetNumberOfComponentsPerPixel());
    }
    else
    {
      sitk::Image diff =
        sitk::Subtract(sitk::Cast(testImage, sitk::sitkFloat32), sitk::Cast(baselineImage, sitk::sitkFloat32));
      diffSquared = sitk::Multiply(diff, diff);
    }
  }
  catch (std::exception & e)
  {
    mMessage = "ImageCompare: Failed to subtract image " + baselineImageFilename + " because: " + e.what();
    return -1;
  }


  sitk::StatisticsImageFilter stats;
  stats.Execute(diffSquared);
  const double rms = std::sqrt(stats.GetMean());

  if (!reportErrors)
  {
    // The measurement errors should be reported for both success and errors
    // to facilitate setting tight tolerances of tests.
    std::cout << "<DartMeasurement name=\"RMSEDifference " << shortFilename << "\" type=\"numeric/float\">" << rms
              << "</DartMeasurement>" << std::endl;
  }
  else
  {
    std::ostringstream msg;
    msg << "ImageCompare: image Root Mean Square (RMS) difference was " << rms << " which exceeds the tolerance of "
        << mTolerance;
    msg << "\n";
    mMessage = msg.str();

    std::cout << "<DartMeasurement name=\"RMSEDifference\" type=\"numeric/float\">" << rms << "</DartMeasurement>"
              << std::endl;
    std::cout << "<DartMeasurement name=\"Tolerance\" type=\"numeric/float\">" << mTolerance << "</DartMeasurement>"
              << std::endl;

    std::string volumeName = OutputDir + "/" + shortFilename + ".nrrd";
    sitk::ImageFileWriter().SetFileName(volumeName).Execute(testImage);

    // Save pngs
    std::string ExpectedImageFilename = OutputDir + "/" + shortFilename + "_Expected.png";
    std::string ActualImageFilename = OutputDir + "/" + shortFilename + "_Actual.png";
    std::string DifferenceImageFilename = OutputDir + "/" + shortFilename + "_Difference.png";

    try
    {
      NormalizeAndSave(baselineImage, ExpectedImageFilename);
      NormalizeAndSave(testImage, ActualImageFilename);
      NormalizeAndSave(sitk::Sqrt(diffSquared), DifferenceImageFilename);

      // Let ctest know about it
      std::cout << "<DartMeasurementFile name=\"ExpectedImage\" type=\"image/png\">";
      std::cout << ExpectedImageFilename << "</DartMeasurementFile>" << std::endl;
      std::cout << "<DartMeasurementFile name=\"ActualImage\" type=\"image/png\">";
      std::cout << ActualImageFilename << "</DartMeasurementFile>" << std::endl;
      std::cout << "<DartMeasurementFile name=\"DifferenceImage\" type=\"image/png\">";
      std::cout << DifferenceImageFilename << "</DartMeasurementFile>" << std::endl;

      stats.Execute(sitk::Sqrt(diffSquared));

      std::cout << "<DartMeasurement name=\"DifferenceImage Minimum\" type=\"numeric/double\">";
      std::cout << stats.GetMinimum() << "</DartMeasurement>" << std::endl;

      std::cout << "<DartMeasurement name=\"DifferenceImage Maximum\" type=\"numeric/double\">";
      std::cout << stats.GetMaximum() << "</DartMeasurement>" << std::endl;

      std::cout << "<DartMeasurement name=\"DifferenceImage Mean\" type=\"numeric/double\">";
      std::cout << stats.GetMean() << "</DartMeasurement>" << std::endl;
    }
    catch (std::exception & e)
    {
      std::cerr << "Exception encountered while trying to normalize and save images for dashboard!" << std::endl;
      std::cerr << e.what() << std::endl;
    }
    catch (...)
    {
      std::cerr << "Unexpected error while trying to normalize and save images for dashboard!" << std::endl;
    }
  }

  return (rms > fabs(mTolerance)) ? rms : 0.0;
}

sitk::Image
ImageCompare::extractSlice(const sitk::Image & image)
{
  // Extract the center slice of our test image
  if (image.GetDimension() == 3)
  {
    std::vector<int>          idx(3, 0);
    std::vector<unsigned int> sz = image.GetSize();

    // set to just the center slice
    idx[2] = (int)(image.GetDepth() / 2);
    sz[2] = 1;
    return sitk::RegionOfInterest(image, sz, idx);
  }

  return image;
}

bool
ImageCompare::compare(const sitk::Image & image, std::string inTestCase, std::string inTag)
{
  sitk::Image centerSlice(0, 0, sitk::sitkUInt8);
  std::string testCase = inTestCase;
  std::string tag = inTag;
  std::string testName = ::testing::UnitTest::GetInstance()->current_test_info()->name();

  if (testCase.empty())
  {
    testCase = ::testing::UnitTest::GetInstance()->current_test_info()->test_case_name();
  }

  // Does the baseline exist?
  std::string extension = ".nrrd";
  std::string OutputDir = dataFinder.GetOutputDirectory();

  std::string name = testCase.append("_").append(testName);

  if (!tag.empty())
  {
    name.append("_").append(tag);
  }

  centerSlice = extractSlice(image);

  const std::string baselineFileName = dataFinder.GetFile("Baseline/" + name + extension);


  if (!itksys::SystemTools::FileExists(baselineFileName.c_str(), true))
  {
    // Baseline does not exist, write out what we've been given
    std::string newBaselineDir = OutputDir + "/Newbaseline/";
    itksys::SystemTools::MakeDirectory(newBaselineDir.c_str());
    std::cout << "Making directory " << newBaselineDir << std::endl;
    std::string newBaseline = newBaselineDir + name + extension;
    sitk::ImageFileWriter().SetFileName(newBaseline).Execute(centerSlice);
    mMessage = "Baseline does not exist, wrote " + newBaseline + "\ncp " + newBaseline + " " + baselineFileName;
    return false;
  }

  // Generate all possible baseline filenames
  std::vector<std::string> baselineFileNames;

  {
    int x = 0;

    baselineFileNames.push_back(baselineFileName);


    while (++x)
    {
      std::ostringstream filename_stream;
      filename_stream << "Baseline/" << name << "." << x << extension;

      std::string filename = dataFinder.GetFile(filename_stream.str());

      if (!itksys::SystemTools::FileExists(filename, true))
      {
        break;
      }
      baselineFileNames.push_back(filename);
    }
  }

  std::vector<std::string>::const_iterator iterName;

  std::string bestBaselineName = *baselineFileNames.begin();
  float       bestRMS = std::numeric_limits<float>::max();

  for (iterName = baselineFileNames.begin(); iterName != baselineFileNames.end(); ++iterName)
  {

    sitk::Image baseline(0, 0, sitk::sitkUInt8);

    try
    {
      baseline = sitk::ImageFileReader().SetFileName(*iterName).Execute();
    }
    catch (std::exception & e)
    {
      mMessage = "ImageCompare: Failed to load image " + *iterName + " because: " + e.what();
      return false;
    }

    float RMS = testImages(centerSlice, baseline, false, *iterName);

    if (RMS >= 0.0 && RMS < bestRMS)
    {
      bestBaselineName = *iterName;
      bestRMS = RMS;
    }
  }


  if (bestRMS > fabs(mTolerance))
  {
    sitk::Image baseline = sitk::ImageFileReader().SetFileName(bestBaselineName).Execute();
    testImages(centerSlice, baseline, true, bestBaselineName);
    return false;
  }
  else
  {
    return true;
  }
}
