/*=========================================================================
 *
 *  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 <itkElastixRegistrationMethod.h>
#include <elxParameterObject.h>
#include <Core/elxVersionMacros.h>
#include <itkImage.h>
#include <itkSize.h>
#include <cassert>

int
main()
{
  try
  {
    enum
    {
      imageSizeX = 10,
      imageSizeY = 10
    };

    std::cout << "Image registation example: estimating the translation between two images of " << imageSizeX << "x"
              << imageSizeY << " pixels using Impact metric\n"
              << " - Linked to elastix library version " ELASTIX_VERSION_STRING "\n - Web site: https://elastix.dev"
                 "\n - Source code repository: https://github.com/SuperElastix/elastix"
                 "\n - Discussion forum: https://github.com/SuperElastix/elastix/discussions\n\n";

    using ImageType = itk::Image<float>;

    const auto printImage = [](const ImageType & image) {
      for (int y{}; y < imageSizeY; ++y)
      {
        std::cout << '\t';

        for (int x{}; x < imageSizeX; ++x)
        {
          std::cout << image.GetPixel({ x, y }) << ' ';
        }
        std::cout << '\n';
      }
      std::cout << '\n';
    };

    const auto fixedImage = ImageType::New();
    fixedImage->SetRegions(itk::Size<>{ imageSizeX, imageSizeY });
    fixedImage->AllocateInitialized();
    fixedImage->SetPixel({ 1, 1 }, 3);
    fixedImage->SetPixel({ 2, 1 }, 4);
    fixedImage->SetPixel({ 1, 2 }, 5);
    std::cout << "Fixed image:\n";

    printImage(*fixedImage);

    const auto movingImage = ImageType::New();
    movingImage->SetRegions(itk::Size<>{ imageSizeX, imageSizeY });
    movingImage->AllocateInitialized();
    movingImage->SetPixel({ 3, 2 }, 3);
    movingImage->SetPixel({ 4, 2 }, 4);
    movingImage->SetPixel({ 3, 3 }, 5);
    std::cout << "Moving image:\n";
    printImage(*movingImage);

    const auto parameterObject = elx::ParameterObject::New();
    parameterObject->SetParameterMap(
      elx::ParameterObject::ParameterMapType{ { "ImageSampler", { "Full" } },
                                              { "MaximumNumberOfIterations", { "44" } },
                                              { "Metric", { "Impact" } },
                                              { "ImpactModelsPath0", { "./Conv2D_ChannelReplicator.pt" } },
                                              { "ImpactDimension0", { "2" } },
                                              { "ImpactNumberOfChannels0", { "1" } },
                                              { "ImpactPatchSize0", { "3", "3" } },
                                              { "ImpactVoxelSize0", { "1", "1" } },
                                              { "ImpactLayersMask0", { "1" } },
                                              { "ImpactSubsetFeatures0", { "3" } },
                                              { "ImpactLayersWeight0", { "1" } },
                                              { "ImpactPCA0", { "0" } },
                                              { "ImpactDistance0", { "L2" } },
                                              { "ImpactModelsPath1", { "./Conv2D_ChannelReplicator.pt" } },
                                              { "ImpactDimension1", { "2" } },
                                              { "ImpactNumberOfChannels1", { "1" } },
                                              { "ImpactPatchSize1", { "3", "3" } },
                                              { "ImpactVoxelSize1", { "1", "1" } },
                                              { "ImpactLayersMask1", { "1" } },
                                              { "ImpactSubsetFeatures1", { "3" } },
                                              { "ImpactLayersWeight1", { "1" } },
                                              { "ImpactPCA1", { "0" } },
                                              { "ImpactDistance1", { "L2" } },
                                              { "ImpactMode", { "Jacobian" } },
                                              { "ImpactGPU", { "-1" } },
                                              { "NumberOfResolutions", { "2" } },
                                              { "Optimizer", { "AdaptiveStochasticGradientDescent" } },
                                              { "Transform", { "TranslationTransform" } } });

    const auto filter = itk::ElastixRegistrationMethod<ImageType, ImageType>::New();
    filter->SetFixedImage(fixedImage);
    filter->SetMovingImage(movingImage);
    filter->SetParameterObject(parameterObject);
    filter->LogToConsoleOn();
    filter->Update();

    const auto * const transformParameterObject = filter->GetTransformParameterObject();
    assert(transformParameterObject != nullptr);

    const auto & transformParameterMaps = transformParameterObject->GetParameterMaps();
    assert(!transformParameterMaps.empty());

    const auto & transformParameterMap = transformParameterMaps.front();
    const auto   foundTransformParameter = transformParameterMap.find("TransformParameters");
    assert(foundTransformParameter != transformParameterMap.cend());

    std::cout << "Estimated translation:\n\t";

    for (const auto & estimatedValue : foundTransformParameter->second)
    {
      std::cout << estimatedValue << ' ';
    }
    std::cout << std::endl;
  }
  catch (const std::exception & stdException)
  {
    // An exception may occur, for instance, when a components used here was not installed when elastix was built.
    std::cerr << "\nCaught exception:\n" << stdException.what() << '\n';
  }

} // end main()
