/*=========================================================================
 *
 *  Copyright UMC Utrecht and contributors
 *
 *  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 "elxParameterObject.h"
#include "elxConversion.h"

#include "itkParameterFileParser.h"

#include <itkNumberToString.h>
#include "itksys/SystemTools.hxx"

#include <fstream>
#include <iostream>
#include <cmath>
#include <iomanip> // For quoted.

namespace elastix
{

/**
 * ********************* SetParameterMap *********************
 */

void
ParameterObject::SetParameterMap(const ParameterMapType & parameterMap)
{
  ParameterMapVectorType parameterMapVector(1, parameterMap);
  this->SetParameterMaps(parameterMapVector);
}

/**
 * ********************* SetParameterMap *********************
 */

void
ParameterObject::SetParameterMap(const unsigned int index, const ParameterMapType & parameterMap)
{
  GetMutableParameterMap(index) = parameterMap;
}


/**
 * ********************* SetParameterMaps *********************
 */

void
ParameterObject::SetParameterMaps(const ParameterMapVectorType & parameterMaps)
{
  if (m_ParameterMaps != parameterMaps)
  {
    m_ParameterMaps = parameterMaps;
    this->Modified();
  }
}


/**
 * ********************* AddParameterMap *********************
 */

void
ParameterObject::AddParameterMap(const ParameterMapType & parameterMap)
{
  m_ParameterMaps.push_back(parameterMap);
  this->Modified();
}


/**
 * ********************* GetMutableParameterMap *********************
 */

ParameterObject::ParameterMapType &
ParameterObject::GetMutableParameterMap(const unsigned int index)
{
  if (const auto numberOfParameterMaps = m_ParameterMaps.size(); index >= m_ParameterMaps.size())
  {
    itkExceptionMacro("The index (" << index << ") is out of range! It must be less than the number of parameter maps ("
                                    << numberOfParameterMaps << ")");
  }
  return m_ParameterMaps[index];
}


/**
 * ********************* GetParameterMap *********************
 */

const ParameterObject::ParameterMapType &
ParameterObject::GetParameterMap(const unsigned int index) const
{
  // Let the const version call the non-const version, as suggested by Scott Meyers (Effective C++).
  return const_cast<Self &>(*this).GetMutableParameterMap(index);
}


/**
 * ********************* SetParameter *********************
 */

void
ParameterObject::SetParameter(const unsigned int index, const ParameterKeyType & key, const ParameterValueType & value)
{
  GetMutableParameterMap(index)[key] = { value };
}


/**
 * ********************* SetParameter *********************
 */

void
ParameterObject::SetParameter(const unsigned int               index,
                              const ParameterKeyType &         key,
                              const ParameterValueVectorType & value)
{
  GetMutableParameterMap(index)[key] = value;
}


/**
 * ********************* SetParameter *********************
 */

void
ParameterObject::SetParameter(const ParameterKeyType & key, const ParameterValueType & value)
{
  for (auto & parameterMap : m_ParameterMaps)
  {
    parameterMap[key] = { value };
  }
}


/**
 * ********************* SetParameter *********************
 */

void
ParameterObject::SetParameter(const ParameterKeyType & key, const ParameterValueVectorType & value)
{
  for (auto & parameterMap : m_ParameterMaps)
  {
    parameterMap[key] = value;
  }
}


/**
 * ********************* GetParameter *********************
 */

const ParameterObject::ParameterValueVectorType &
ParameterObject::GetParameter(const unsigned int index, const ParameterKeyType & key) const
{
  const auto & parameterMap = GetParameterMap(index);

  if (const auto found = parameterMap.find(key); found != parameterMap.end())
  {
    return found->second;
  }

  itkExceptionMacro("There is no parameter named " << std::quoted(key) << " in parameter map #" << index << "!");
}


/**
 * ********************* HasParameter *********************
 */

bool
ParameterObject::HasParameter(const unsigned int index, const ParameterKeyType & key) const
{
  if (index >= m_ParameterMaps.size())
  {
    return false;
  }
  return m_ParameterMaps[index].count(key) > 0;
}


/**
 * ********************* HasParameter *********************
 */

bool
ParameterObject::HasParameter(const ParameterKeyType & key) const
{
  for (const auto & parameterMap : m_ParameterMaps)
  {
    if (parameterMap.find(key) != parameterMap.end())
    {
      return true;
    }
  }
  return false;
}


/**
 * ********************* RemoveParameter *********************
 */

void
ParameterObject::RemoveParameter(const unsigned int index, const ParameterKeyType & key)
{
  GetMutableParameterMap(index).erase(key);
}


/**
 * ********************* RemoveParameter *********************
 */

void
ParameterObject::RemoveParameter(const ParameterKeyType & key)
{
  for (auto & parameterMap : m_ParameterMaps)
  {
    parameterMap.erase(key);
  }
}


/**
 * ********************* ReadParameterFile *********************
 */

void
ParameterObject::ReadParameterFile(const ParameterFileNameType & parameterFileName)
{
  this->SetParameterMaps({ itk::ParameterFileParser::ReadParameterMap(parameterFileName) });
}


/**
 * ********************* ReadParameterFiles *********************
 */

void
ParameterObject::ReadParameterFiles(const ParameterFileNameVectorType & parameterFileNameVector)
{
  if (parameterFileNameVector.empty())
  {
    itkExceptionMacro("Parameter filename container is empty.");
  }

  m_ParameterMaps.clear();

  for (const auto & parameterFileName : parameterFileNameVector)
  {
    if (!itksys::SystemTools::FileExists(parameterFileName))
    {
      itkExceptionMacro("Parameter file \"" << parameterFileName << "\" does not exist.");
    }

    this->AddParameterFile(parameterFileName);
  }
}


/**
 * ********************* AddParameterFile *********************
 */

void
ParameterObject::AddParameterFile(const ParameterFileNameType & parameterFileName)
{
  m_ParameterMaps.push_back(itk::ParameterFileParser::ReadParameterMap(parameterFileName));
}


/**
 * ********************* WriteParameterFile *********************
 */


void
ParameterObject::WriteParameterFile(const ParameterMapType &      parameterMap,
                                    const ParameterFileNameType & parameterFileName)
{
  std::ofstream parameterFile;
  parameterFile.exceptions(std::ofstream::failbit | std::ofstream::badbit);
  parameterFile << std::fixed;

  try
  {
    parameterFile.open(parameterFileName, std::ofstream::out);
  }
  catch (const std::ios_base::failure & e)
  {
    itkGenericExceptionMacro("Error opening parameter file: " << e.what());
  }

  try
  {
    const auto format = itksys::SystemTools::GetFilenameLastExtension(parameterFileName) ==
                            Conversion::CreateParameterMapFileNameExtension(ParameterMapStringFormat::Toml)
                          ? ParameterMapStringFormat::Toml
                          : ParameterMapStringFormat::LegacyTxt;

    parameterFile << Conversion::ParameterMapToString(parameterMap, format);
  }
  catch (const std::ios_base::failure & e)
  {
    itkGenericExceptionMacro("Error writing to paramter file: " << e.what());
  }

  try
  {
    parameterFile.close();
  }
  catch (const std::ios_base::failure & e)
  {
    itkGenericExceptionMacro("Error closing parameter file:" << e.what());
  }
}


void
ParameterObject::WriteParameterFile(const ParameterFileNameType & parameterFileName) const
{
  if (m_ParameterMaps.empty())
  {
    itkExceptionMacro("Error writing parameter map to disk: The parameter object is empty.");
  }

  if (m_ParameterMaps.size() > 1)
  {
    itkExceptionMacro("Error writing to disk: The number of parameter maps ("
                      << m_ParameterMaps.size()
                      << ") does not match the number of provided filenames (1). Please call WriteParameterFiles "
                         "instead, and provide a vector of filenames.");
  }

  this->WriteParameterFile(m_ParameterMaps[0], parameterFileName);
}


/**
 * ********************* WriteParameterFiles *********************
 */


void
ParameterObject::WriteParameterFiles() const
{
  ParameterFileNameVectorType parameterFileNameVector;
  for (unsigned int i = 0; i < m_ParameterMaps.size(); ++i)
  {
    parameterFileNameVector.push_back("ParametersFile." + std::to_string(i) + ".txt");
  }

  Self::WriteParameterFiles(m_ParameterMaps, parameterFileNameVector);
}


void
ParameterObject::WriteParameterFiles(const ParameterMapVectorType &      parameterMapVector,
                                     const ParameterFileNameVectorType & parameterFileNameVector)
{
  if (parameterMapVector.size() != parameterFileNameVector.size())
  {
    itkGenericExceptionMacro("Error writing to disk: The number of parameter maps ("
                             << parameterMapVector.size() << ") does not match the number of provided filenames ("
                             << parameterFileNameVector.size() << ").");
  }

  if (!parameterMapVector.empty())
  {
    Self::WriteParameterFile(parameterMapVector.front(), parameterFileNameVector.front());
  }

  // Add initial transform parameter file names. Do not touch the first one,
  // since it may have one already
  for (unsigned int i = 1; i < parameterMapVector.size(); ++i)
  {
    ParameterMapType parameterMap = parameterMapVector[i];
    if (parameterMap.find("TransformParameters") != parameterMap.end())
    {
      parameterMap["InitialTransformParameterFileName"][0] = parameterFileNameVector[i - 1];
    }

    Self::WriteParameterFile(parameterMap, parameterFileNameVector[i]);
  }
}


void
ParameterObject::WriteParameterFiles(const ParameterFileNameVectorType & parameterFileNameVector) const
{
  Self::WriteParameterFiles(m_ParameterMaps, parameterFileNameVector);
}


/**
 * ********************* GetDefaultParameterMap *********************
 */

ParameterObject::ParameterMapType
ParameterObject::GetDefaultParameterMap(const std::string & transformName,
                                        const unsigned int  numberOfResolutions,
                                        const double        finalGridSpacingInPhysicalUnits)
{
  if (transformName == "nonrigid")
  {
    itkGenericExceptionMacro("The transform name "
                             << std::quoted(transformName)
                             << " is no longer supported, after elastix version 5.2.0. It was originally equivalent to "
                                "\"bspline\" (which is the preferred transform name).");
  }

  // Parameters that depend on size and number of resolutions
  ParameterMapType parameterMap{};

  // Common Components
  parameterMap["FixedImagePyramid"] = { "FixedSmoothingImagePyramid" };
  parameterMap["MovingImagePyramid"] = { "MovingSmoothingImagePyramid" };
  parameterMap["Interpolator"] = { "LinearInterpolator" };
  parameterMap["Optimizer"] = { "AdaptiveStochasticGradientDescent" };
  parameterMap["Resampler"] = { "DefaultResampler" };
  parameterMap["ResampleInterpolator"] = { "FinalBSplineInterpolator" };
  parameterMap["FinalBSplineInterpolationOrder"] = { "3" };
  parameterMap["NumberOfResolutions"] = { std::to_string(numberOfResolutions) };
  parameterMap["WriteIterationInfo"] = { "false" };
  parameterMap["MaximumNumberOfIterations"] = { "256" };

  // Image Sampler
  parameterMap["ImageSampler"] = { "RandomCoordinate" };
  parameterMap["NumberOfSpatialSamples"] = { "2048" };
  parameterMap["CheckNumberOfSamples"] = { "true" };
  parameterMap["MaximumNumberOfSamplingAttempts"] = { "8" };
  parameterMap["NewSamplesEveryIteration"] = { "true" };

  // Optimizer
  parameterMap["NumberOfSamplesForExactGradient"] = { "4096" };
  parameterMap["DefaultPixelValue"] = { "0" };
  parameterMap["AutomaticParameterEstimation"] = { "true" };

  // Output
  parameterMap["WriteResultImage"] = { "true" };
  parameterMap["ResultImageFormat"] = { "nii" };

  // transformNames
  if (transformName == "translation")
  {
    parameterMap["Registration"] = { "MultiResolutionRegistration" };
    parameterMap["Transform"] = { "TranslationTransform" };
    parameterMap["Metric"] = { "AdvancedMattesMutualInformation" };
    parameterMap["AutomaticTransformInitialization"] = { "true" };
  }
  else if (transformName == "rigid")
  {
    parameterMap["Registration"] = { "MultiResolutionRegistration" };
    parameterMap["Transform"] = { "EulerTransform" };
    parameterMap["Metric"] = { "AdvancedMattesMutualInformation" };
    parameterMap["AutomaticScalesEstimation"] = { "true" };
  }
  else if (transformName == "affine")
  {
    parameterMap["Registration"] = { "MultiResolutionRegistration" };
    parameterMap["Transform"] = { "AffineTransform" };
    parameterMap["Metric"] = { "AdvancedMattesMutualInformation" };
    parameterMap["AutomaticScalesEstimation"] = { "true" };
  }
  else if (transformName == "bspline")
  {
    parameterMap["Registration"] = { "MultiMetricMultiResolutionRegistration" };
    parameterMap["Transform"] = { "BSplineTransform" };
    parameterMap["Metric"] = { "AdvancedMattesMutualInformation", "TransformBendingEnergyPenalty" };
    parameterMap["Metric0Weight"] = { "1.0" };
    parameterMap["Metric1Weight"] = { "1.0" };
  }
  else if (transformName == "spline")
  {
    parameterMap["Registration"] = { "MultiResolutionRegistration" };
    parameterMap["Transform"] = { "SplineKernelTransform" };
    parameterMap["Metric"] = { "AdvancedMattesMutualInformation" };
  }
  else if (transformName == "groupwise")
  {
    parameterMap["Registration"] = { "MultiResolutionRegistration" };
    parameterMap["Transform"] = { "BSplineStackTransform" };
    parameterMap["Metric"] = { "VarianceOverLastDimensionMetric" };
    parameterMap["Interpolator"] = { "ReducedDimensionBSplineInterpolator" };
    parameterMap["ResampleInterpolator"] = { "FinalReducedDimensionBSplineInterpolator" };
  }
  else
  {
    itkGenericExceptionMacro("No default parameter map \"" << transformName << "\".");
  }

  // B-spline transform settings
  if (transformName == "bspline" || transformName == "groupwise")
  {
    parameterMap["GridSpacingSchedule"] = [numberOfResolutions] {
      ParameterValueVectorType gridSpacingSchedule{};
      gridSpacingSchedule.reserve(numberOfResolutions);

      for (unsigned int resolution = 0; resolution < numberOfResolutions; ++resolution)
      {
        gridSpacingSchedule.push_back(
          itk::ConvertNumberToString(std::pow(2.0, (numberOfResolutions - resolution - 1) / 2.0)));
      }
      return gridSpacingSchedule;
    }();

    parameterMap["FinalGridSpacingInPhysicalUnits"] = { std::to_string(finalGridSpacingInPhysicalUnits) };
  }

  return parameterMap;
}


/**
 * ********************* PrintSelf *********************
 */

void
ParameterObject::PrintSelf(std::ostream & os, itk::Indent indent) const
{
  Superclass::PrintSelf(os, indent);

  for (unsigned int i = 0; i < m_ParameterMaps.size(); ++i)
  {
    os << "ParameterMap " << i << ": " << std::endl;

    for (const auto & parameter : m_ParameterMaps[i])
    {
      os << "  (" << parameter.first;

      for (const std::string & value : parameter.second)
      {
        if (Conversion::IsNumber(value))
        {
          os << ' ' << value;
        }
        else
        {
          os << " \"" << value << '"';
        }
      }

      os << ')' << std::endl;
    }
  }
}


} // namespace elastix
