/*****************************************************************************
 * $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 "ConnectedComponents.h"

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

#include <itkConnectedComponentImageFilter.h>
#include <itkRelabelComponentImageFilter.h>
#include <itkBinaryBallStructuringElement.h>

using namespace camitk;

// --------------- constructor -------------------
ConnectedComponents::ConnectedComponents(ActionExtension* extension) : Action(extension) {
    // Setting name, description and input component
    setName("Connected Components Classification");
    setDescription("Labels connected components of a binary image and order the labes with respect to the size of the connected components (i.e., the larges connected component has the label 1, the one a little smaller has the label 2, and so on...)  ");
    setComponentClassName("ImageComponent");

    // Setting classification family and tags
    this->setFamily("ITK Filter");
    this->addTag("Neighborhood Filter");
    this->addTag("Connected");

    Property* numberOfCCProperty = new Property(tr("Number of Connected Components"), "?", tr("The computed number of connected components computed by this filter. \nThis property is read only and only updated after applying the action on a component."), "");
    numberOfCCProperty->setReadOnly(true);
    addParameter(numberOfCCProperty);
}

// --------------- getWidget -------------------
QWidget* ConnectedComponents::getWidget() {
    setParameterValue("Number of Connected Components", "?");

    //-- return the default widget
    return Action::getWidget();
}

// --------------- apply -------------------
Action::ApplyStatus ConnectedComponents::apply() {
    for (Component* comp : getTargets()) {
        ImageComponent* input = dynamic_cast<ImageComponent*>(comp);
        process(input);
        CAMITK_INFO(tr("Number of Connected Components for %1: %2").arg(input->getName()).arg(getParameterValue("Number of Connected Components").toDouble()))
        // update the value of Computed Threshold in the GUI (no update if the widget is not shown, e.g., in a headless pipelines or test)
        ActionWidget* myActionWidget = dynamic_cast<ActionWidget*>(actionWidget);
        if (myActionWidget != nullptr) {
            myActionWidget->update();
        }
        refreshApplication();
    }
    return SUCCESS;
}

void ConnectedComponents::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() + "_cc");

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

    refreshApplication();
}

#include "ConnectedComponents.impl"

// ITK filter implementation
template <class InputPixelType, class OutputPixelType, const int dim>
vtkSmartPointer<vtkImageData> ConnectedComponents::itkProcess(vtkSmartPointer<vtkImageData> img) {
    vtkSmartPointer<vtkImageData> filteredImage = 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();

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

    // Filter to label connected components
    typedef itk::ConnectedComponentImageFilter<InputImageType, OutputImageType> FilterType;
    typename FilterType::Pointer filter = FilterType::New();
    // Filter to re-order connected components by size
    typedef itk::RelabelComponentImageFilter<OutputImageType, OutputImageType> RelabelFilterType;
    typename RelabelFilterType::Pointer relabelFilter = RelabelFilterType::New();

    // To update CamiTK progress bar while filtering, add an ITK observer to the filters.
    ItkProgressObserver::Pointer observer = ItkProgressObserver::New();
    observer->SetCoef(50.0);
    observer->SetStartValue(0.0);
    ItkProgressObserver::Pointer observerRelabel = ItkProgressObserver::New();
    observerRelabel->SetCoef(50.0);
    observerRelabel->SetStartValue(50.0);

    // --------------------- Plug filters and parameters ---------------------------
    // From VTK to ITK
    vtkToItkFilter->SetInput(img);
    // For the filter itself
    filter->SetInput(vtkToItkFilter->GetOutput());
    filter->AddObserver(itk::ProgressEvent(), observer);

    relabelFilter->AddObserver(itk::ProgressEvent(), observerRelabel);
    relabelFilter->SetInput(filter->GetOutput());

    // From ITK to VTK
    itkToVtkFilter->SetInput(relabelFilter->GetOutput());

    // --------------------- Actually execute all filters parts --------------------
    itkToVtkFilter->Update();

    // --------------------- Create and return a copy (the filters will be deleted)--
    resultImage = itkToVtkFilter->GetOutput();

    int extent[6];
    resultImage->GetExtent(extent);
    filteredImage->SetExtent(extent);
    filteredImage->DeepCopy(resultImage);

    observer->Reset();
    observerRelabel->Reset();

    setParameterValue("Number of Connected Components", (int) relabelFilter->GetNumberOfObjects());

    return filteredImage;
}

