/*****************************************************************************
 * $CAMITK_LICENCE_BEGIN$
 *
 * CamiTK - Computer Assisted Medical Intervention ToolKit
 * (c) 2001-2025 Univ. Grenoble Alpes, CNRS, Grenoble INP - UGA, TIMC, 38000 Grenoble, France
 *
 * 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 "RegionGrowing.h"

// includes from CamiTK
#include <Log.h>
#include <ItkProgressObserver.h>
#include <itkImageToVTKImageFilter.h>
#include <itkVTKImageToImageFilter.h>

// includes from Itk
#include <itkConnectedThresholdImageFilter.h>
#include <itkCurvatureFlowImageFilter.h>
#include <itkCastImageFilter.h>

// local includes
#include "RegionGrowingWidget.h"

using namespace camitk;


// --------------- constructor -------------------
RegionGrowing::RegionGrowing(ActionExtension* extension) : Action(extension) {
    // Setting name, description and input component
    setName("Region Growing Segmentation");
    setDescription("	This filter segments volume images using region growing approach. It starts from a seed point that is considered to be inside the object to be segmented. The pixels neighboring this seed point are evaluated to determine if they should also be considered part of the object. If so, they are added to the region and the process continues as long as new pixels are added to the region. 	For connected threshold filter, the criterion is based on the intensity of the neighboring pixels. They are considered as part of the object if theyr value is between a lower and an upper threshold.  ");
    setComponentClassName("ImageComponent");

    setEmbedded(true);

    // Setting classification family and tags
    this->setFamily("ITK Segmentation");
    this->addTag("Threshold");
    this->addTag("Manal");
    this->addTag("Classification");
    this->addTag("Growing Region");
    this->addTag("Seed Point");
    this->addTag("Region Growing");

    // Setting the widget containing the parameters
    theWidget = nullptr;

}

// --------------- destructor -------------------
RegionGrowing::~RegionGrowing() {
    // 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)
}

// --------------- getWidget --------------
QWidget* RegionGrowing::getWidget() {

    RegionGrowingWidget* rgWidget = dynamic_cast<RegionGrowingWidget*>(theWidget);

    //-- create the widget if needed
    if (!rgWidget) {
        theWidget = new RegionGrowingWidget(this);
        rgWidget = dynamic_cast<RegionGrowingWidget*>(theWidget);
    }

    //-- update the widget with a PickedPixelMap param
    rgWidget->updateComponent(dynamic_cast<ImageComponent*>(getTargets().last()));


    return theWidget;
}

// --------------- apply -------------------
Action::ApplyStatus RegionGrowing::apply() {
    // check the widget
    RegionGrowingWidget* rgWidget = dynamic_cast<RegionGrowingWidget*>(theWidget);
    // Get the image
    ImageComponent* input = dynamic_cast<ImageComponent*>(getTargets().last());

    // this call works only with a GUI (i.e. if theWidget exists)
    if ((input == nullptr) || (rgWidget == nullptr)) {
        CAMITK_WARNING(tr("Cannot be called without a GUI (input data are required to be set manually. Action Aborted."))
        return ABORTED;
    }

    // Get the parameters
    filterBefore = rgWidget->isSmoothingChecked();
    nbIterations  = rgWidget->getNumberOfIterations();
    timeStep   = rgWidget->getTimeStep();
    lowThreshold  = rgWidget->getLowThreshold();
    highThreshold = rgWidget->getHighThreshold();
    seedPoints = rgWidget->getSeedPoints(input);

    process(input);

    return SUCCESS;
}

Action::ApplyStatus RegionGrowing::apply(int lowThreshold, int highThreshold, QList<QVector3D>* seedPoints,
        bool filterBefore, int nbIterations, double timeStep) {
    // Get the image
    ImageComponent* input =  dynamic_cast<ImageComponent*>(getTargets().last());

    // Get the parameters
    this->lowThreshold = lowThreshold;
    this->highThreshold = highThreshold;
    this->seedPoints = seedPoints;
    this->filterBefore = filterBefore;
    this->nbIterations = nbIterations;
    this->timeStep = timeStep;

    process(input);

    return SUCCESS;
}

void RegionGrowing::process(ImageComponent* comp) {
    // ITK filter implementation using templates
    vtkSmartPointer<vtkImageData> inputImage = comp->getImageData();
    vtkSmartPointer<vtkImageData> outputImage = implementProcess(inputImage);

    ImageComponent* outputComp = new ImageComponent(outputImage, comp->getName() + "_segmented");

    // same frame as input
    outputComp->setFrameFrom(comp);

    refreshApplication();

}

#include "RegionGrowing.impl"

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

    // --------------------- 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< double,          dim > DoubleImageType;
    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();

    // Convert the input image as a DoubleImageType to be able to perform filtering
    typedef itk::CastImageFilter<InputImageType, DoubleImageType> CastFilterInType;
    typename CastFilterInType::Pointer castFilter = CastFilterInType::New();

    // Smoothing filter
    typedef itk::CurvatureFlowImageFilter<DoubleImageType, DoubleImageType> CurvatureFlowImageFilterType;
    typename CurvatureFlowImageFilterType::Pointer smoothingFilter = CurvatureFlowImageFilterType::New();

    // R�gion Growing filter
    typedef itk::ConnectedThresholdImageFilter<DoubleImageType, OutputImageType> ConnectedFilterType;
    typename ConnectedFilterType::Pointer connectedThreshold = ConnectedFilterType::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> itkToVtkFilterType;
    typename itkToVtkFilterType::Pointer itkToVtkFilter = itkToVtkFilterType::New();

    // Connecting the pipeline ...

    // 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);

    // From VTK to ITK
    vtkToItkFilter->SetInput(img);

    // From input type to double type
    castFilter->SetInput(vtkToItkFilter->GetOutput());
    castFilter->AddObserver(itk::ProgressEvent(), observer);
    castFilter->Update();
    observer->Reset();

    // Smoothing
    if (filterBefore) {
        smoothingFilter->SetInput(castFilter->GetOutput());
        smoothingFilter->SetNumberOfIterations(nbIterations);
        smoothingFilter->SetTimeStep(timeStep);
        smoothingFilter->AddObserver(itk::ProgressEvent(), observer);
        smoothingFilter->Update();
        observer->Reset();

        connectedThreshold->SetInput(smoothingFilter->GetOutput());
    }
    else {
        connectedThreshold->SetInput(castFilter->GetOutput());
    }

    connectedThreshold->AddObserver(itk::ProgressEvent(), observer);
    connectedThreshold->SetLower(lowThreshold);
    connectedThreshold->SetUpper(highThreshold);
    connectedThreshold->SetReplaceValue(255);
    // set the initial seeds
    for (int i = 0; i < seedPoints->size(); i++) {
        typename InputImageType::IndexType seed;
        seed[0] = seedPoints->at(i).x();
        seed[1] = seedPoints->at(i).y();
        if (dim == 3) {
            seed[2] = seedPoints->at(i).z();
        }
        connectedThreshold->AddSeed(seed);
    }
    connectedThreshold->Update();
    observer->Reset();

    itkToVtkFilter->SetInput(connectedThreshold->GetOutput());
    itkToVtkFilter->Update();
    observer->Reset();

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

    return result;
}

