/*****************************************************************************
 * $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 "MeshToImageStencil.h"
#include <Log.h>
#include <Property.h>

// Qt includes
#include <QColor>
#include <QVector3D>

// Vtk includes
// disable warning generated by clang about the surrounded headers
#include <CamiTKDisableWarnings>
#include <vtkMetaImageWriter.h>
#include <vtkPolyDataToImageStencil.h>
#include <vtkImageStencil.h>
#include <CamiTKReEnableWarnings>

#include <vtkSmartPointer.h>
#include <vtkPolyData.h>
#include <vtkImageData.h>
#include <vtkSphereSource.h>
#include <vtkPointData.h>

using namespace camitk;


// --------------- Constructor -------------------
MeshToImageStencil::MeshToImageStencil(ActionExtension* extension) : Action(extension) {
    // Setting name, description and input component
    setName("Mesh To Image Stencil");
    setDescription(tr("This action generates a new volume image and converts the mesh into a volume representation (vtkImageData) where the foreground voxels are colored and the background voxels are black."));
    setComponentClassName("MeshComponent");

    // Setting classification family and tags
    setFamily("Basic Mesh");
    addTag(tr("STE"));

    // Setting the action's parameters
//    Property* colorProperty = new Property(tr("Color"), QVariant(QColor(255, 255, 255, 255)), tr("Foreground voxels colour."), "");
//    addParameter(colorProperty);

    Property* dimensionProperty = new Property(tr("Dimension"), QVariant(QVector3D(0.0, 0.0, 0.0)), tr("The dimension of the output volumic image"), "");
    addParameter(dimensionProperty);

    Property* originProperty = new Property(tr("Origin"), QVariant(QVector3D(0.0, 0.0, 0.0)), tr("The origin frame position"), "");
    addParameter(originProperty);

    Property* spacingProperty = new Property(tr("Spacing"), QVariant(QVector3D(1.0, 1.0, 1.0)), tr("The spacing between each voxel"), "");
    addParameter(spacingProperty);

    Property* outputFileProperty = new Property(tr("Output file"), "", tr("The output filename"), "");
    addParameter(outputFileProperty);
}

// --------------- destructor -------------------
MeshToImageStencil::~MeshToImageStencil() {
    // Do not do anything yet.
    // Delete stuff if you create stuff
    // (except if you use smart pointers of course !!)
}

// --------------- apply -------------------
Action::ApplyStatus MeshToImageStencil::apply() {
    for (Component* comp : getTargets()) {
        MeshComponent* input = dynamic_cast<MeshComponent*>(comp);
        ApplyStatus status = process(input);
        // abort if one of the selecteed component processing is not successful
        if (status != SUCCESS) {
            return status;
        }
    }
    return SUCCESS;
}

// --------------- process -------------------
Action::ApplyStatus MeshToImageStencil::process(MeshComponent* comp) {
    // Get the parameters
    QVector3D dimension = getParameterValue("Dimension").value<QVector3D>();
    QVector3D origin = getParameterValue("Origin").value<QVector3D>();
    QVector3D spacing = getParameterValue("Spacing").value<QVector3D>();
    QString file = getParameterValue("Output file").value<QString>();

    // get mesh as polydata
    if (comp->getPointSet() == nullptr) {
        CAMITK_WARNING(tr("The mesh point set of \"%1\" does not have a vtkPointSet. Action aborted.").arg(comp->getName()))
        return ABORTED;
    }

    vtkSmartPointer<vtkPolyData> pData = vtkPolyData::SafeDownCast(comp->getPointSet());

    if (pData == nullptr) {
        CAMITK_WARNING(tr("The mesh point set of \"%1\" is not a vtkPolyData. Action aborted.").arg(comp->getName()))
        return ABORTED;
    }

    vtkSmartPointer<vtkImageData> colorImage = vtkSmartPointer<vtkImageData>::New();
    double bounds[6];
    pData->GetBounds(bounds);
    double space[3]; // desired volume spacing
    space[0] = spacing.x();
    space[1] = spacing.y();
    space[2] = spacing.z();
    colorImage->SetSpacing(space);

    // compute dimensions
    int dim[3];
    // if no dimension are specified used the bounds and spacing to compute a default value
    if (dimension.x() == 0.0 && dimension.y() == 0.0 && dimension.z() == 0.0) {
        for (int i = 0; i < 3; i++) {
            dim[i] = static_cast<int>(ceil((bounds[i * 2 + 1] - bounds[i * 2]) / spacing[i]));
        }
        CAMITK_WARNING(QString("Dimensions not set → set to (%1,%2,%3), deduced using bounds and spacing").arg(dim[0]).arg(dim[1]).arg(dim[2]));
    }
    else {
        dim[0] = dimension.x();
        dim[1] = dimension.y();
        dim[2] = dimension.z();
    }
    colorImage->SetDimensions(dim);
    colorImage->SetExtent(0, dim[0] - 1, 0, dim[1] - 1, 0, dim[2] - 1);

    double ori[3];
    ori[0] = origin.x();
    ori[1] = origin.y();
    ori[2] = origin.z();
    colorImage->SetOrigin(ori);

    colorImage->AllocateScalars(VTK_UNSIGNED_CHAR, 1);

    // fill the image with foreground voxels:
    //unsigned char inval = 255;
    unsigned char outval = 0;
    vtkIdType count = colorImage->GetNumberOfPoints();
    for (vtkIdType i = 0; i < count; ++i) {
        colorImage->GetPointData()->GetScalars()->SetTuple1(i, 255);
    }

    // polygonal data --> image stencil:
    vtkSmartPointer<vtkPolyDataToImageStencil> pol2stenc = vtkSmartPointer<vtkPolyDataToImageStencil>::New();
    pol2stenc->SetInputData(pData);

    pol2stenc->SetOutputOrigin(ori);
    pol2stenc->SetOutputSpacing(space);
    pol2stenc->SetOutputWholeExtent(colorImage->GetExtent());
    pol2stenc->Update();

    // cut the corresponding color image and set the background:
    vtkSmartPointer<vtkImageStencil> imgstenc = vtkSmartPointer<vtkImageStencil>::New();
    imgstenc->SetInputData(colorImage);
    imgstenc->SetStencilConnection(pol2stenc->GetOutputPort());
    imgstenc->ReverseStencilOff();
    imgstenc->SetBackgroundValue(outval);
    imgstenc->Update();

    vtkSmartPointer<vtkMetaImageWriter> writer = vtkSmartPointer<vtkMetaImageWriter>::New();
    writer->SetFileName(file.toLatin1());
    writer->SetInputData(imgstenc->GetOutput());
    writer->Write();

    return SUCCESS;
}


