/*****************************************************************************
 * $CAMITK_LICENCE_BEGIN$
 *
 * CamiTK - Computer Assisted Medical Intervention ToolKit
 * (c) 2001-2018 Univ. Grenoble Alpes, CNRS, TIMC-IMAG UMR 5525 (GMCAO)
 *
 * Visit http://camitk.imag.fr for more information
 *
 * This file is part of CamiTK.
 *
 * CamiTK is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License version 3
 * only, as published by the Free Software Foundation.
 *
 * CamiTK is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License version 3 for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * version 3 along with CamiTK.  If not, see <http://www.gnu.org/licenses/>.
 *
 * $CAMITK_LICENCE_END$
 ****************************************************************************/
#include "AnisotropicDiffusion.h"

#include <Application.h>

#include <ItkProgressObserver.h>
#include <itkImageToVTKImageFilter.h>
#include <itkVTKImageToImageFilter.h>

#include <itkCastImageFilter.h>
#include <itkRescaleIntensityImageFilter.h>
#include <itkGradientAnisotropicDiffusionImageFilter.h>
#include <itkCurvatureAnisotropicDiffusionImageFilter.h>
#include <QVariant>

// CamiTK includes
#include <Property.h>

using namespace camitk;


// --------------- constructor -------------------
AnisotropicDiffusion::AnisotropicDiffusion(ActionExtension* extension) : Action(extension) {
    // Setting name, description and input component
    setName("Anisotropic Diffusion");
    setDescription("<p>Anisotropic diffusion methods are formulated to reduce noise (or unwanted detail) in images while preserving specific image features. For many applications, there is an assumption that light-dark transitions (edges) are interesting. Standard isotropic diffusion methods move and blur light-dark boundaries. Anisotropic diffusion methods are formulated to specifically preserve edges.</p> \
				   <p>The numberOfIterations parameter specifies the number of iterations (time-step updates) that the solver will perform to produce a solution image. The appropriate number of iterations is dependent on the application and the image being processed. As a general rule, the more iterations performed, the more diffused the image will become.</p> \
				   <p>The conductance parameter controls the sensitivity of the conductance term in the basic anisotropic diffusion equation. It affect the conductance term in different ways depending on the particular variation on the basic equation. As a general rule, the lower the value, the more strongly the diffusion equation preserves image features (such as high gradients or curvature). A high value for conductance will cause the filter to diffuse image features more readily. Typical values range from 0.5 to 2.0 for data like the Visible Human color data, but the correct value for your application is wholly dependent on the results you want from a specific data set and the number or iterations you perform.</p> \
				   <p>The <i>Gradient</i> anisotropic diffusion implements an N-dimensional version of the classic Perona-Malik anisotropic diffusion equation for scal-valued images.</p> \
                   <p>The <i>Curvature</i> anisotropic diffusion performs anisotropic diffusion on an image using a modified curvature diffusion equation (MCDE). MCDE does not exhibit the edge enhancing properties of classic anisotropic diffusion, which can under certain conditions undergo a <i>negative</i> diffusion, which enhances the contrast of edges. Equations of the form MCDE always undergo positive diffusion, with the conductance term only varying the stregth of that diffusion.</p> \
				   ");
    setComponent("ImageComponent");

    // Setting classification family and tags
    this->setFamily("ITK Filter");
    this->addTag("Edge Preserving");
    this->addTag("Smoothing");
    this->addTag("Blur");
    this->addTag("Perona and Malik");

    // Setting parameters default values by using properties
    addParameter(new Property(tr("Keep original image voxel type"), true,
                              tr("Keep the original image voxel type ?"), ""));
    Property* nbIterationProperty = new Property(tr("Number of iterations"), 5,
            tr("The more iterations, the more smoothing. \nEach iteration takes the same amount of time. \nIf it takes 10 seconds for one iteration, then it will take 100 seconds for 10 iterations. \nNote that the conductance controls how much each iteration smooths across edges. "), "");
    nbIterationProperty->setAttribute("minimum", 1);
    nbIterationProperty->setAttribute("maximum", 100);
    nbIterationProperty->setAttribute("singleStep", 1);
    addParameter(nbIterationProperty);

    Property* conductanceProperty = new Property(tr("Conductance"), 1.0,
            tr("Conductance controls the sensitivity of the conductance term. \nAs a general rule, the lower the value, the more strongly the filter preserves edges. \nA high value will cause diffusion (smoothing) across edges. \nNote that the number of iterations controls how much smoothing is done within regions bounded by edges."), "");
    conductanceProperty->setAttribute("minimum", 0.10);
    conductanceProperty->setAttribute("maximum", 10.00);
    conductanceProperty->setAttribute("singleStep", 0.05);
    addParameter(conductanceProperty);

    Property* implementationProperty = new Property("Diffusion type", AnisotropicDiffusion::GRADIENT,
            "The type of diffusion to use.", "");
    implementationProperty->setEnumTypeName("AnisoDiffType", this);
    addParameter(implementationProperty);

}

// --------------- destructor -------------------
AnisotropicDiffusion::~AnisotropicDiffusion() {
    // do not delete the widget has it might have been used in the ActionViewer (i.e. the ownership might have been taken by the stacked widget)
}

AnisotropicDiffusion::AnisoDiffType AnisotropicDiffusion::getDiffusionType() {
    return (AnisotropicDiffusion::AnisoDiffType) property("Diffusion type").toInt();
}


// --------------- apply -------------------
Action::ApplyStatus AnisotropicDiffusion::apply() {
    foreach (Component* comp, getTargets()) {
        ImageComponent* input = dynamic_cast<ImageComponent*>(comp);
        process(input);
    }
    return SUCCESS;
}

// --------------- process -------------------
void AnisotropicDiffusion::process(ImageComponent* comp) {
    // Get the parameters
    this->keepOrgVoxelType = property("Keep original image voxel type").toBool();
    this->numberOfIterations = property("Number of iterations").toInt();
    this->conductance = property("Conductance").toDouble();

    // ITK filter implementation using templates
    vtkSmartPointer<vtkImageData> inputImage = comp->getImageData();
    vtkSmartPointer<vtkImageData> outputImage = implementProcess(inputImage);
    ImageComponent* outputComp = new ImageComponent(outputImage, comp->getName() + "_anisoDiff");

    // consider frame policy on new image created
    Action::applyTargetPosition(comp, outputComp);

    Application::refresh();
}

#include "AnisotropicDiffusion.impl"

// ITK filter implementation
template <class InputPixelType, class OutputPixelType, const int dim>
vtkSmartPointer<vtkImageData> AnisotropicDiffusion::itkProcess(vtkSmartPointer<vtkImageData> img) {
    vtkSmartPointer<vtkImageData> outputImage = vtkSmartPointer<vtkImageData>::New();

    switch (this->getDiffusionType()) {
        default:
        case AnisotropicDiffusion::GRADIENT:
            outputImage = itkProcessGradientAnisotropicDiffusion<InputPixelType, OutputPixelType, dim>(img);
            break;
        case AnisotropicDiffusion::CURVATURE:
            outputImage = itkProcessCurvatureAnisotropicDiffusion<InputPixelType, OutputPixelType, dim>(img);
            break;
    }

    return outputImage;

}

template <class InputPixelType, class OutputPixelType, const int dim>
vtkSmartPointer<vtkImageData> AnisotropicDiffusion::itkProcessGradientAnisotropicDiffusion(vtkSmartPointer<vtkImageData> img) {
    vtkSmartPointer<vtkImageData> result = vtkSmartPointer<vtkImageData>::New();
    vtkSmartPointer<vtkImageData> resultImage;

    // --------------------- Filters declaration and creation ----------------------
    // Define ITK input and output image types with respect to the instantiation
    //    types of the tamplate.
    typedef itk::Image< InputPixelType,  dim > InputImageType;
    typedef itk::Image< OutputPixelType, dim > OutputImageType;

    // Convert the image from CamiTK in VTK format to ITK format to use ITK filters.
    typedef itk::VTKImageToImageFilter<InputImageType> vtkToItkFilterType;
    typename vtkToItkFilterType::Pointer vtkToItkFilter = vtkToItkFilterType::New();

    // To use anisotropic diffusion on real type data
    typedef itk::CastImageFilter<InputImageType, OutputImageType> CastFilterType;
    typename CastFilterType::Pointer toDoubleFilter = CastFilterType::New();

    // Anisotropic filter
    typedef itk::GradientAnisotropicDiffusionImageFilter<OutputImageType, OutputImageType> FilterType;
    typename FilterType::Pointer filter = FilterType::New();

    // To go back to the original image type.
    typedef itk::RescaleIntensityImageFilter<OutputImageType, InputImageType> ToOrgFilterType;
    typename ToOrgFilterType::Pointer toOrgFilter = ToOrgFilterType::New();

    // In the same way, once the image is filtered, we need to convert it again to
    // VTK format to give it to CamiTK.
    typedef itk::ImageToVTKImageFilter<OutputImageType> ItkToVtkFloatFilterType;
    typename ItkToVtkFloatFilterType::Pointer itkToVtkFloatFilter = ItkToVtkFloatFilterType::New();

    typedef itk::ImageToVTKImageFilter<InputImageType> ItkToVtkOrgFilterType;
    typename ItkToVtkOrgFilterType::Pointer itkToVtkOrgFilter = ItkToVtkOrgFilterType::New();

// ------------------------- WRITE YOUR CODE HERE ----------------------------------

    // To update CamiTK progress bar while filtering, add an ITK observer to the filters.
    ItkProgressObserver::Pointer observer = ItkProgressObserver::New();
    // ITK observers generally give values between 0 and 1, and CamiTK progress bar
    //    wants values between 0 and 100...
    observer->SetCoef(100.0);

    // --------------------- Plug filters and parameters ---------------------------
    // From VTK to ITK
    vtkToItkFilter->SetInput(img);
    vtkToItkFilter->SetInput(img);
    vtkToItkFilter->AddObserver(itk::ProgressEvent(), observer);
    vtkToItkFilter->Update();
    observer->Reset();

    toDoubleFilter->SetInput(vtkToItkFilter->GetOutput());
    toDoubleFilter->AddObserver(itk::ProgressEvent(), observer);
    toDoubleFilter->Update();
    observer->Reset();

    filter->SetInput(toDoubleFilter->GetOutput());
    filter->SetNumberOfIterations(numberOfIterations);
    filter->SetConductanceParameter(conductance);
    filter->AddObserver(itk::ProgressEvent(), observer);
    filter->Update();
    observer->Reset();

    if (keepOrgVoxelType) {
        toOrgFilter->SetInput(filter->GetOutput());
        toOrgFilter->AddObserver(itk::ProgressEvent(), observer);
        toOrgFilter->Update();
        observer->Reset();

        itkToVtkOrgFilter->SetInput(toOrgFilter->GetOutput());
        itkToVtkOrgFilter->AddObserver(itk::ProgressEvent(), observer);
        itkToVtkOrgFilter->Update();
        observer->Reset();

        resultImage = itkToVtkOrgFilter->GetOutput();
    }
    else {
        itkToVtkFloatFilter->SetInput(filter->GetOutput());
        itkToVtkFloatFilter->AddObserver(itk::ProgressEvent(), observer);
        itkToVtkFloatFilter->Update();
        observer->Reset();

        resultImage = itkToVtkFloatFilter->GetOutput();
    }

    // --------------------- Create and return a copy (the filters will be deleted)--
    int extent[6];
    resultImage->GetExtent(extent);
    result->SetExtent(extent);
    result->DeepCopy(resultImage);

    // Set CamiTK progress bar back to zero (the processing filter is over)
    observer->Reset();

    return result;
}

template <class InputPixelType, class OutputPixelType, const int dim>
vtkSmartPointer<vtkImageData> AnisotropicDiffusion::itkProcessCurvatureAnisotropicDiffusion(vtkSmartPointer<vtkImageData> img) {
    vtkSmartPointer<vtkImageData> result = vtkSmartPointer<vtkImageData>::New();
    vtkSmartPointer<vtkImageData> resultImage;

    // --------------------- Filters declaration and creation ----------------------
    // Define ITK input and output image types with respect to the instantiation
    //    types of the tamplate.
    typedef itk::Image< InputPixelType,  dim > InputImageType;
    typedef itk::Image< OutputPixelType, dim > OutputImageType;

    // Convert the image from CamiTK in VTK format to ITK format to use ITK filters.
    typedef itk::VTKImageToImageFilter<InputImageType> vtkToItkFilterType;
    typename vtkToItkFilterType::Pointer vtkToItkFilter = vtkToItkFilterType::New();

    // To use anisotropic diffusion on real type data
    typedef itk::CastImageFilter<InputImageType, OutputImageType> CastFilterType;
    typename CastFilterType::Pointer toDoubleFilter = CastFilterType::New();

    // Anisotropic filter
    typedef itk::CurvatureAnisotropicDiffusionImageFilter<OutputImageType, OutputImageType> FilterType;
    typename FilterType::Pointer filter = FilterType::New();

    // To go back to the original image type.
    typedef itk::RescaleIntensityImageFilter<OutputImageType, InputImageType> ToOrgFilterType;
    typename ToOrgFilterType::Pointer toOrgFilter = ToOrgFilterType::New();

    // In the same way, once the image is filtered, we need to convert it again to
    // VTK format to give it to CamiTK.
    typedef itk::ImageToVTKImageFilter<OutputImageType> ItkToVtkFloatFilterType;
    typename ItkToVtkFloatFilterType::Pointer itkToVtkFloatFilter = ItkToVtkFloatFilterType::New();

    typedef itk::ImageToVTKImageFilter<InputImageType> ItkToVtkOrgFilterType;
    typename ItkToVtkOrgFilterType::Pointer itkToVtkOrgFilter = ItkToVtkOrgFilterType::New();

// ------------------------- WRITE YOUR CODE HERE ----------------------------------

    // To update CamiTK progress bar while filtering, add an ITK observer to the filters.
    ItkProgressObserver::Pointer observer = ItkProgressObserver::New();
    // ITK observers generally give values between 0 and 1, and CamiTK progress bar
    //    wants values between 0 and 100...
    observer->SetCoef(100.0);

    // --------------------- Plug filters and parameters ---------------------------
    // From VTK to ITK
    vtkToItkFilter->SetInput(img);
    vtkToItkFilter->SetInput(img);
    vtkToItkFilter->AddObserver(itk::ProgressEvent(), observer);
    vtkToItkFilter->Update();
    observer->Reset();

    toDoubleFilter->SetInput(vtkToItkFilter->GetOutput());
    toDoubleFilter->AddObserver(itk::ProgressEvent(), observer);
    toDoubleFilter->Update();
    observer->Reset();

    filter->SetInput(toDoubleFilter->GetOutput());
    filter->SetNumberOfIterations(numberOfIterations);
    filter->SetConductanceParameter(conductance);
    filter->AddObserver(itk::ProgressEvent(), observer);
    filter->Update();
    observer->Reset();

    if (keepOrgVoxelType) {
        toOrgFilter->SetInput(filter->GetOutput());
        toOrgFilter->AddObserver(itk::ProgressEvent(), observer);
        toOrgFilter->Update();
        observer->Reset();

        itkToVtkOrgFilter->SetInput(toOrgFilter->GetOutput());
        itkToVtkOrgFilter->AddObserver(itk::ProgressEvent(), observer);
        itkToVtkOrgFilter->Update();
        observer->Reset();

        resultImage = itkToVtkOrgFilter->GetOutput();
    }
    else {
        itkToVtkFloatFilter->SetInput(filter->GetOutput());
        itkToVtkFloatFilter->AddObserver(itk::ProgressEvent(), observer);
        itkToVtkFloatFilter->Update();
        observer->Reset();

        resultImage = itkToVtkFloatFilter->GetOutput();
    }

    // --------------------- Create and return a copy (the filters will be deleted)--
    int extent[6];
    resultImage->GetExtent(extent);
    result->SetExtent(extent);
    result->DeepCopy(resultImage);

    // Set CamiTK progress bar back to zero (the processing filter is over)
    observer->Reset();

    return result;
}