//
// This file is part of MorphoGraphX - http://www.MorphoGraphX.org
// Copyright (C) 2012-2015 Richard S. Smith and collaborators.
//
// If you use MorphoGraphX in your work, please cite:
//   http://dx.doi.org/10.7554/eLife.05864
//
// MorphoGraphX is free software, and is licensed under under the terms of the 
// GNU General (GPL) Public License version 2.0, http://www.gnu.org/licenses.
// 
#include "ITKProgress.hpp"

#include "ITKFilter.hpp"

#include <itkSmoothingRecursiveGaussianImageFilter.h>
#include <itkCurvatureFlowImageFilter.h>
#include <itkRescaleIntensityImageFilter.h>
#include <itkCurvatureAnisotropicDiffusionImageFilter.h>
#include <itkGradientMagnitudeRecursiveGaussianImageFilter.h>
#include <itkSigmoidImageFilter.h>

namespace mgx
{
  bool ITKSmoothingRecursiveGaussianImageFilter::run(const Store* input, Store* output, float radius)
  {
    UImageConverter::Pointer converter = UImageConverter::New();
  
    converter->SetStore(input);
  
    typedef itk::SmoothingRecursiveGaussianImageFilter<UImageType, UImageType> FilterType;
  
    FilterType::Pointer filter = FilterType::New();
    filter->SetSigma(radius);
    filter->SetNormalizeAcrossScale(false);
    filter->SetInput(converter->GetOutput());
  
    ITKProgress progress("Smoothing Recursive Gaussian Image Filter");
    progress.setFilter(filter);
  
    filter->Update();
  
    UImageType::ConstPointer result = filter->GetOutput();
  
    if(!UImageConverter::TransferImage(output->stack(), output, result)) {
      setErrorMessage("ITK Process tried to change the size of the stack");
      return false;
    }
  
    output->changed();
    return true;
  }
  REGISTER_PROCESS(ITKSmoothingRecursiveGaussianImageFilter);
  
  bool ITKCurvatureFlowImageFilter::run(const Store* input, Store* output, float timestep, int steps)
  {
    UImageConverter::Pointer converter = UImageConverter::New();
  
    converter->SetStore(input);
  
    typedef itk::CurvatureFlowImageFilter<UImageType, FImageType> FilterType;
    FilterType::Pointer filter = FilterType::New();
    filter->SetInput(converter->GetOutput());
    filter->SetTimeStep(timestep);
    filter->SetNumberOfIterations(steps);
  
    typedef itk::RescaleIntensityImageFilter<FImageType, UImageType> RescalerType;
    RescalerType::Pointer rescaler = RescalerType::New();
    rescaler->SetInput(filter->GetOutput());
    rescaler->SetOutputMinimum(0);
    rescaler->SetOutputMaximum(65535);
  
    ITKProgress progress("Curvature Flow Image Filter");
    progress.setFilter(filter);
  
    rescaler->Update();
  
    UImageType::ConstPointer result = rescaler->GetOutput();
  
    if(!UImageConverter::TransferImage(output->stack(), output, result)) {
      setErrorMessage("ITK Process tried to change the size of the stack");
      return false;
    }
  
    output->changed();
    return true;
  }
  REGISTER_PROCESS(ITKCurvatureFlowImageFilter);
  
  bool ITKCurvatureAnisotropicDiffusionImageFilter::run(const Store* input, Store* output, float timestep,
                                                               int iterations, float conductance)
  {
    FImageConverter::Pointer converter = FImageConverter::New();
  
    converter->SetStore(input);
  
    typedef itk::CurvatureAnisotropicDiffusionImageFilter<FImageType, FImageType> FilterType;
    FilterType::Pointer filter = FilterType::New();
    filter->SetInput(converter->GetOutput());
    filter->SetTimeStep(timestep);
    filter->SetNumberOfIterations(iterations);
    filter->SetConductanceParameter(conductance);
  
    ITKProgress progress("Curvature Anisotropic Diffusion Image Filter");
    progress.setFilter(filter);
  
    filter->Update();
  
    FImageType::ConstPointer result = filter->GetOutput();
  
    if(!FImageConverter::TransferImage(output->stack(), output, result)) {
      setErrorMessage("ITK Process tried to change the size of the stack");
      return false;
    }
  
    output->changed();
    return true;
  }
  REGISTER_PROCESS(ITKCurvatureAnisotropicDiffusionImageFilter);
  
  bool ITKGradientMagnitudeRecursiveGaussianImageFilter::run(const Store* input, Store* output, float sigma)
  {
    UImageConverter::Pointer converter = UImageConverter::New();
  
    converter->SetStore(input);
  
    typedef itk::GradientMagnitudeRecursiveGaussianImageFilter<UImageType, UImageType> FilterType;
    FilterType::Pointer filter = FilterType::New();
    filter->SetInput(converter->GetOutput());
    filter->SetSigma(sigma);
  
    typedef itk::RescaleIntensityImageFilter<UImageType, UImageType> RescalerType;
    RescalerType::Pointer rescaler = RescalerType::New();
    rescaler->SetInput(filter->GetOutput());
    rescaler->SetOutputMinimum(0);
    rescaler->SetOutputMaximum(65535);
  
    ITKProgress progress("Gradient Magnitude Recursive Gaussian Image Filter");
    progress.setFilter(filter);
  
    rescaler->Update();
  
    UImageType::ConstPointer result = rescaler->GetOutput();
  
    if(!UImageConverter::TransferImage(output->stack(), output, result)) {
      setErrorMessage("ITK Process tried to change the size of the stack");
      return false;
    }
  
    output->changed();
    return true;
  }
  REGISTER_PROCESS(ITKGradientMagnitudeRecursiveGaussianImageFilter);
  
  bool ITKSigmoidImageFilter::run(const Store* input, Store* output, float alpha, float beta)
  {
    UImageConverter::Pointer converter = UImageConverter::New();
  
    converter->SetStore(input);
  
    typedef itk::SigmoidImageFilter<UImageType, UImageType> FilterType;
    FilterType::Pointer filter = FilterType::New();
    filter->SetInput(converter->GetOutput());
    filter->SetAlpha(alpha);
    filter->SetBeta(beta);
  
    ITKProgress progress("Sigmoid Image Filter");
    progress.setFilter(filter);
  
    filter->Update();
  
    UImageType::ConstPointer result = filter->GetOutput();
  
    if(!UImageConverter::TransferImage(output->stack(), output, result)) {
      setErrorMessage("ITK Process tried to change the size of the stack");
      return false;
    }
  
    output->changed();
    return true;
  }
  REGISTER_PROCESS(ITKSigmoidImageFilter);
}
