/*****************************************************************************
 * $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$
 ****************************************************************************/
// 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;

ICPRegistration::ICPRegistration(camitk::ActionExtension* ext) : Action(ext) {

    this->setName("ICP Registration");
    this->setDescription(QString(tr("Iterative Closest Point algorithm bewteen 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->setComponent("MeshComponent");
    this->setFamily("Mesh Processing");
    this->addTag(tr("registration"));
    distanceMeasureType = RMS;

    Property* numberOfIterationsProperty = new Property(tr("Number of iterations"), 20, tr("The number of iteration of the ICP algorithm."), "");
    addParameter(numberOfIterationsProperty);

    Property* distanceMesureTypeProperty = new Property(tr("Distance mesure type"), RMS, tr("The distance mesure type use by the ICP algorithm."), "");
    distanceMesureTypeProperty->setEnumTypeName("DistanceMeasureType", this);
    addParameter(distanceMesureTypeProperty);
}

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;

    }

    CAMITK_TRACE(tr("ICP Registration from mesh \"%1\" to mesh \"%2\".").arg(getTargets().first()->getName(), getTargets().last()->getName()))

    // set up ICP
    vtkSmartPointer<vtkIterativeClosestPointTransform> icp = vtkSmartPointer<vtkIterativeClosestPointTransform>::New();
    icp->SetSource(sourceMesh->getPointSet());
    icp->SetTarget(targetMesh->getPointSet());
    icp->GetLandmarkTransform()->SetModeToRigidBody();
    icp->SetMaximumNumberOfIterations(property("Number of iterations").toInt());
    switch (property("Distance mesure type").toInt()) {
        case RMS :
            icp->SetMeanDistanceModeToRMS();
            break;
        case ABS :
            icp->SetMeanDistanceModeToAbsoluteValue();
            break;
        default:
            icp->SetMeanDistanceModeToRMS();
            break;
    }
    icp->AddObserver(vtkCommand::ProgressEvent, progressCallback);
    icp->Modified();
    icp->Update();

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

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

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