/*****************************************************************************
 * $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$
 ****************************************************************************/
// CamiTK includes
#include "ICPRegistration.h"
#include <Application.h>
#include <MeshComponent.h>
#include <Log.h>
#include <Property.h>

// Vtk includes
#include <vtkIterativeClosestPointTransform.h>
#include <vtkTransformPolyDataFilter.h>
#include <vtkLandmarkTransform.h>
#include <vtkCallbackCommand.h>

using namespace camitk;

// -------------------- constructor --------------------
ICPRegistration::ICPRegistration(camitk::ActionExtension* ext) : Action(ext) {

    this->setName("ICP Registration");
    this->setDescription(QString(tr("Iterative Closest Point algorithm between two mesh.<br/>")) +
                         QString(tr("At least two mesh components must be selected :")) +
                         QString("<ul>") +
                         QString(tr("<li> The first one is the source mesh (the one to be registered) </li>")) +
                         QString(tr("<li> The second one is the target mesh </li>")) +
                         QString("</ul>"));
    this->setComponentClassName("MeshComponent");
    this->setFamily("Registration");
    this->addTag(tr("Mesh Processing"));
    this->addTag(tr("CPI"));

    addParameter(new Property(tr("Maximum Number Of Iterations"), 50, tr("The maximum number of iterations of the ICP algorithm."), ""));

    addParameter(new Property(tr("Maximum Number Of Landmarks"), 200, tr("The maximum number of landmarks sampled in your dataset. If your dataset is dense, then you will typically not need all the points to compute the ICP transform. The default is 200."), ""));

    Property* distanceMesureTypeProperty = new Property(tr("Distance Mesure Type"), RMS, tr("The distance mesure type used by the ICP algorithm. Specify the mean distance mode. The RMS mode is the square root of the average of the sum of squares of the closest point distances. The Absolute Value mode is the mean of the sum of absolute values of the closest point distances."), "");
    distanceMesureTypeProperty->setEnumTypeName("DistanceMeasureType", this);
    addParameter(distanceMesureTypeProperty);

    Property* transformTypeProperty = new Property(tr("Transform Type"), Similarity, tr("The landmark transformation type used by the ICP agorithm. Set the number of degrees of freedom to constrain the solution to: Rigidbody (rotation and translation only), Similarity (rotation, translation and isotropic scaling) or Affine (collinearity is preserved, ratios of distances along a line are preserved."), "");
    transformTypeProperty->setEnumTypeName("LandmarkTransformType", this);
    addParameter(transformTypeProperty);

    addParameter(new Property(tr("Match Centroid First"), false, tr("Starts the process by translating source centroid to target centroid."), ""));

    addParameter(new Property(tr("Check Mean Distance"), false, tr("Force the algorithm to check the mean distance between two iterations. This is slower but generally more accurate."), ""));

    Property* maxMeanDistanceProperty = new Property(tr("Maximum Mean Distance"), 0.01, tr("The maximum mean distance between two iterations. If the mean distance is lower than this, the convergence stops."), "");
    maxMeanDistanceProperty->setAttribute("decimals", 5);
    maxMeanDistanceProperty->setAttribute("minimum", 0);
    maxMeanDistanceProperty->setAttribute("singleStep", 0.001);
    addParameter(maxMeanDistanceProperty);
}

// -------------------- apply --------------------
camitk::Action::ApplyStatus ICPRegistration::apply() {
    // set waiting cursor and status bar
    QApplication::setOverrideCursor(QCursor(Qt::WaitCursor));
    Application::showStatusBarMessage("Registration...");
    Application::resetProgressBar();
    vtkSmartPointer<vtkCallbackCommand> progressCallback = vtkSmartPointer<vtkCallbackCommand>::New();
    progressCallback->SetCallback(&Application::vtkProgressFunction);

    // at least 2 targets must be selected
    if (getTargets().size() < 2) {
        QApplication::restoreOverrideCursor();
        CAMITK_WARNING(tr("At least 2 mesh components must be selected. Action aborted."))
        return ABORTED;
    }

    // use the first target as source mesh and the last as target mesh
    MeshComponent* sourceMesh = dynamic_cast<MeshComponent*>(getTargets().first());
    MeshComponent* targetMesh = dynamic_cast<MeshComponent*>(getTargets().last());

    // check if targets are MeshComponents
    if ((targetMesh == nullptr) || (sourceMesh == nullptr)) {
        QApplication::restoreOverrideCursor();
        CAMITK_WARNING(tr("Mesh \"%1\" and/or \"%2\" is not a MeshComponent. Action aborted.").arg(getTargets().first()->getName(), getTargets().last()->getName()))
        return ABORTED;

    }

    // ICP Registration from first target to last targets

    // set up ICP
    vtkSmartPointer<vtkIterativeClosestPointTransform> icp = vtkSmartPointer<vtkIterativeClosestPointTransform>::New();
    icp->SetSource(sourceMesh->getPointSet());
    icp->SetTarget(targetMesh->getPointSet());
    icp->SetMaximumNumberOfIterations(getParameterValue("Maximum Number Of Iterations").toInt());
    icp->SetMaximumNumberOfLandmarks(getParameterValue("Maximum Number Of Landmarks").toInt());
    switch (getParameterValue("Distance Mesure Type").toInt()) {
        case ABS :
            icp->SetMeanDistanceModeToAbsoluteValue();
            break;
        case RMS :
        default:
            icp->SetMeanDistanceModeToRMS();
            break;
    }
    switch (getParameterValue("Transform Type").toInt()) {
        case Rigidbody:
            icp->GetLandmarkTransform()->SetModeToRigidBody();
            break;
        case Affine:
            icp->GetLandmarkTransform()->SetModeToAffine();
            break;
        case Similarity:
        default:
            icp->GetLandmarkTransform()->SetModeToSimilarity();
            break;
    }
    icp->SetStartByMatchingCentroids(getParameterValue("Match Centroid First").toBool());
    icp->SetCheckMeanDistance(getParameterValue("Check Mean Distance").toBool());
    icp->SetMaximumMeanDistance(getParameterValue("Maximum Mean Distance").toDouble());
    icp->AddObserver(vtkCommand::ProgressEvent, progressCallback);
    icp->Modified();
    icp->Update();

    // transform the source points by the ICP solution
    vtkSmartPointer<vtkTransformPolyDataFilter> icpTransformFilter =
        vtkSmartPointer<vtkTransformPolyDataFilter>::New();
    icpTransformFilter->SetInputData(sourceMesh->getPointSet());
    icpTransformFilter->SetTransform(icp);
    icpTransformFilter->AddObserver(vtkCommand::ProgressEvent, progressCallback);
    icpTransformFilter->Update();

    // get the algorithm output information
    vtkSmartPointer<vtkMatrix4x4> registrationMatrix = icp->GetMatrix();
    QStringList stringMatrix;
    for (int i = 0; i < 4; i++) {
        for (int j = 0; j < 4; j++) {
            stringMatrix << QString::number(icp->GetMatrix()->GetElement(i, j));
        }
    }
    int numberOfIterations = icp->GetNumberOfIterations();
    double meanDistance = icp->GetMeanDistance();
    CAMITK_TRACE("ICP algorithm information: " + QString::number(numberOfIterations) + " iterations, mean distance = " + QString::number(meanDistance) + ", transform: [" + stringMatrix.join(", ") + "]")

    // create a new MeshComponent for the result
    new MeshComponent(icpTransformFilter->GetOutput(), sourceMesh->getName() + "_registered_mesh");

    // refresh restore the normal cursor and progress bar
    refreshApplication();
    Application::resetProgressBar();
    Application::showStatusBarMessage("");
    QApplication::restoreOverrideCursor();
    return SUCCESS;
}
