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

// -- Core stuff classes
#include <Application.h>
#include <Property.h>
#include <ImageComponent.h>
#include <MeshComponent.h>
#include <TransformationManager.h>

// -- vtk filters stuff
// disable warning generated by clang about the surrounded headers
#include <CamiTKDisableWarnings>
#include <vtkPolyDataWriter.h>
#include <CamiTKReEnableWarnings>

#include <vtkImageResample.h>
#include <vtkContourFilter.h>
#include <vtkPolyDataNormals.h>
#include <vtkPolyDataConnectivityFilter.h>
#include <vtkMarchingCubes.h>
#include <vtkPointSet.h>
#include <vtkCallbackCommand.h>

using namespace camitk;


// -------------------- constructor --------------------
ImageReconstructionAction::ImageReconstructionAction(ActionExtension* extension) : Action(extension) {
    this->setName("Reconstruction");
    this->setDescription("<br /> Marching cubes is a simple algorithm for creating a triangle mesh from an implicit function (one of the form f(x, y, z) = 0).<br /> It works by iterating (\"marching\") over a uniform grid of cubes superimposed over a region of the function. If all 8 vertices of the cube are positive, or all 8 vertices are negative, the cube is entirely above or entirely below the surface and no triangles are emitted. Otherwise, the cube straddles the function and some triangles and vertices are generated. Since each vertex can either be positive or negative, there are technically 28 possible configurations, but many of these are equivalent to one another. <br/> <br/>This action uses this algorithm to build a 3D surfacic mesh of the input image.");
    this->setComponentClassName("ImageComponent");
    this->setFamily("Reconstruction");
    this->addTag("Reconstruction");

    // Action parameters
    Property* thresholdProperty = new Property(tr("Threshold"), 127.5, tr("Grey level value corresponding to the isovalue the isosurface represents."), "");
    thresholdProperty->setAttribute("minimum", 0.0);
    thresholdProperty->setAttribute("singleStep", 0.1);
    addParameter(thresholdProperty);

    // Do I keep the largest component ?
    Property* keepLargestCCProperty = new Property(tr("Keep only largest component?"), true, tr("Do we keep only the largest component?"), "");
    addParameter(keepLargestCCProperty);

    // Subsample original image ?
    Property* subsampleOrginalImageProperty = new Property(tr("Subsample original image?"), false, tr("Do we subsample input image for faster computation?"), "");
    addParameter(subsampleOrginalImageProperty);

    Property* subsampledImageX = new Property(tr("Subsampled image width (X)"), 64, tr("The width (X dimension) of the subsampled image"), "");
    subsampledImageX->setAttribute("minimum", 1);
    subsampledImageX->setAttribute("singleStep", 1);
    addParameter(subsampledImageX);

    Property* subsampledImageY = new Property(tr("Subsampled image height (Y)"), 64, tr("The height (Y dimension) of the subsampled image"), "");
    subsampledImageY->setAttribute("minimum", 1);
    subsampledImageY->setAttribute("singleStep", 1);
    addParameter(subsampledImageY);

    Property* subsampledImageZ = new Property(tr("Subsampled image depth (Z)"), 64, tr("The depth (Z dimension) of the subsampled image"), "");
    subsampledImageZ->setAttribute("minimum", 1);
    subsampledImageZ->setAttribute("singleStep", 1);
    addParameter(subsampledImageZ);
}

// -------------------- destructor --------------------
ImageReconstructionAction::~ImageReconstructionAction() {
    if (actionWidget) {
        delete actionWidget;
    }
}

// -------------------- apply --------------------
Action::ApplyStatus ImageReconstructionAction::apply() {
    // get original input ImageComponent
    this->myComponent = dynamic_cast<ImageComponent*>(getTargets().last());

    // compute model
    vtkSmartPointer<vtkPointSet> resultPointSet = this->getMarchingCubesReconstruction();

    MeshComponent* outputComp = new MeshComponent(resultPointSet, myComponent->getName() + "_mesh");

    // as the mesh is computed from the image data, the mesh is defined in the image component data frame
    outputComp->setFrame(TransformationManager::getFrameOfReferenceOwnership(myComponent->getDataFrame()));

    refreshApplication();

    return SUCCESS;
}

// -------------------- getMarchingCubesReconstruction --------------------
vtkSmartPointer<vtkPointSet> ImageReconstructionAction::getMarchingCubesReconstruction() {
    double isoValue = getParameterValue("Threshold").toDouble();
    bool keepLargestConnectedComponent = getParameterValue("Keep only largest component?").toBool();
    bool subsample = getParameterValue("Subsample original image?").toBool();
    int subSampledDimX = getParameterValue("Subsampled image width (X)").toInt();
    int subSampledDimY = getParameterValue("Subsampled image height (Y)").toInt();
    int subSampledDimZ = getParameterValue("Subsampled image depth (Z)").toInt();

    Application::showStatusBarMessage("Computing Marching Cube Reconstruction");
    Application::resetProgressBar();
    vtkSmartPointer<vtkCallbackCommand> progressCallback = vtkSmartPointer<vtkCallbackCommand>::New();
    progressCallback->SetCallback(&Application::vtkProgressFunction);

    // The filter we have chosen to use to compute 3D isosurface of the volume image
    // is vtkMarchingcubes. We could also use vtkContourFilter since it will automatically
    // create an instance of vtkMarchingCubes as it delegates to the fastest subclass
    // for a particular dataset type.
    vtkSmartPointer<vtkMarchingCubes> mcSurface = vtkSmartPointer<vtkMarchingCubes>::New();

    // For medical volumes, marching cubes generates a large number of triangles.
    // To be practical, we'll use a reduced resolution dataset.
    // For example, we take original 256^3 data and reduce it to 64^2 slices
    // by averaging neighboring pixels twice in the slice plane.
    // We call the resulting dataset quarter since it has 1/4 the resolution of the original data.
    vtkSmartPointer<vtkImageResample> resampleFilter = nullptr;

    if (subsample) {
        resampleFilter = vtkSmartPointer<vtkImageResample>::New();
        int* previousDimensions = myComponent->getImageData()->GetDimensions();

        double x_magnification = 0.0;
        double y_magnification = 0.0;
        double z_magnification = 0.0;
        x_magnification = subSampledDimX / (double) previousDimensions[0];
        y_magnification = subSampledDimY / (double) previousDimensions[1];

        resampleFilter->SetAxisMagnificationFactor(0, x_magnification);
        resampleFilter->SetAxisMagnificationFactor(1, y_magnification);

        if (myComponent->getImageData()->GetDataDimension() == 3) {
            z_magnification = subSampledDimZ / (double) previousDimensions[2];
            resampleFilter->SetAxisMagnificationFactor(2, z_magnification);
        }

        // setup the initial data for the pipeline
        resampleFilter->SetInputData(myComponent->getImageData());
        resampleFilter->AddObserver(vtkCommand::ProgressEvent, progressCallback);

        // setup next step in the pipeline
        mcSurface->SetInputConnection(resampleFilter->GetOutputPort());
    }
    else {
        // setup the initial data for the pipeline
        mcSurface->SetInputData(myComponent->getImageData());
    }

    mcSurface->AddObserver(vtkCommand::ProgressEvent, progressCallback);
    mcSurface->SetValue(0, isoValue);

    // The class vtkPolyDataNormals is used to generate
    // nice surface normals for the data. vtkMarchingCubes can also generate normals,
    // but sometimes better results are achieved when the normals are directly from the
    // surface (vtkPolyDataNormals) versus from the data (vtkMarchingCubes).
    vtkSmartPointer<vtkPolyDataNormals> normalsFilter = vtkSmartPointer<vtkPolyDataNormals>::New();

    //Specify the angle that defines a sharp edge. If the difference in angle
    // across neighboring polygons is greater than this value,
    // the shared edge is considered "sharp".
    normalsFilter->SetFeatureAngle(60.0);

    // Keep largest connected component
    vtkSmartPointer<vtkPolyDataConnectivityFilter> connectivityFilter = nullptr;

    if (keepLargestConnectedComponent) {
        connectivityFilter = vtkPolyDataConnectivityFilter::New();
        connectivityFilter->SetInputConnection(mcSurface->GetOutputPort());
        connectivityFilter->SetExtractionModeToLargestRegion();
        // setup next step in the pipeline
        normalsFilter->SetInputConnection(connectivityFilter->GetOutputPort());
    }
    else {
        // setup next step in the pipeline
        normalsFilter->SetInputConnection(mcSurface->GetOutputPort());
    }

    normalsFilter->AddObserver(vtkCommand::ProgressEvent, progressCallback);
    normalsFilter->Update();

    vtkSmartPointer<vtkPointSet> resultPointSet = normalsFilter->GetOutput();

    // One can make efficient data visualization using vtkTriangleFilter
    // and vtkStripper. But first, let us try without...

    Application::showStatusBarMessage("");
    Application::resetProgressBar();

    return resultPointSet;
}

