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

// CamiTK includes
#include <Property.h>
#include <ImageComponent.h>
#include <MeshComponent.h>
#include <Application.h>
#include <TransformationManager.h>
#include <ActionWidget.h>
#include <Log.h>

#include <QColor>

#include <qmessagebox.h>
#include <vtkDataArrayRange.h>
#include <vtkArrayDispatch.h>
#include <vtkPointData.h>
#include <vtkDiscreteFlyingEdges3D.h>
#include <vtkDiscreteMarchingCubes.h>
#include <vtkImageAccumulate.h>
#include <vtkWindowedSincPolyDataFilter.h>
#include <vtkThreshold.h>
#include <vtkMaskFields.h>
#include <vtkGeometryFilter.h>


using namespace camitk;

// -------------------- init --------------------
void MeshesFromImage::init() {

}

// -------------------- nextColor --------------------
QColor MeshesFromImage::nextColor() {
    static int index = -1;
    static const QList<QColor> colors = QList<QColor>({
        // -- R set3 palette
        // from the Qt doc:
        // lighter = multiplies the value (V) component by the given factor
        QColor("#8DD3C7"),
        QColor("#FFFFB3"),
        QColor("#BEBADA"),
        QColor("#FB8072"),
        QColor("#80B1D3"),
        QColor("#FDB462"),
        QColor("#B3DE69"),
        QColor("#FCCDE5"),
        QColor("#D9D9D9"),
        QColor("#BC80BD"),
        QColor("#CCEBC5"),
        QColor("#FFED6F"),
        // -- initial colors from D3 category20
        // QColor("#aec7e8"), QColor("#ffbb78"), QColor("#98df8a"), QColor("#ff9896"), QColor("#c5b0d5"), QColor("#c49c94"), QColor("#f7b6d2"), QColor("#c7c7c7"), QColor("#dbdb8d"), QColor("#9edae5")
        // -- R Accent palette
        QColor("#7FC97F"),
        QColor("#BEAED4"),
        QColor("#FDC086"),
        QColor("#FFFF99"),
        QColor("#386CB0"),
        QColor("#F0027F"),
        QColor("#BF5B17"),
        QColor("#666666"),
        QColor("#B3E2CD")
    });
    ++index;
    return colors[index % 21];
}

// -------------------- parseLabelValues --------------------
QList<int> MeshesFromImage::parseLabelValues(QString labels) {
    QList<int> labelValues;
    if (labels == "") {
        return labelValues;
    }
    QRegularExpression validChars("^[0-9,-]+$");
    if (!validChars.match(labels).hasMatch()) {
        CAMITK_WARNING(tr("The parameter \"Selection Of Label Values\" must only contain digits, commas and hyphens.\n Current value: \"%1\"").arg(labels));
        return labelValues;
    }

    // Split the string by commas
    QStringList parts = labels.split(',');

    for (const QString& part : parts) {
        // Check if the part contains a range
        if (part.contains('-')) {
            // Split the range into from and to
            QStringList range = part.split('-');
            if (range.size() == 2) {
                bool fromValueOk, toValueOk;
                int from = range[0].toInt(&fromValueOk);
                int to = range[1].toInt(&toValueOk);
                if (fromValueOk && toValueOk && from < to) { // the range is valid
                    // Add all numbers in the range to the result list
                    for (int i = from; i <= to; ++i) {
                        labelValues.append(i);
                    }
                }
            }
        }
        else {
            // If it's not a range, just add the number to the result list
            bool numberOk;
            int number = part.toInt(&numberOk);
            if (numberOk) {
                labelValues.append(number);
            }
        }
    }

    return labelValues;
}

// -------------------- process --------------------
Action::ApplyStatus MeshesFromImage::process() {
    // Set waiting cursor
    QApplication::setOverrideCursor(QCursor(Qt::WaitCursor));

    const int maximumNumberOfLabels = 150;

    int method = getParameterValue("Method").value<int>(); // enum as an int
    QString labelSelectionString = getParameterValueAsString("Selection Of Label Values");
    QList<int> selectedLabels = parseLabelValues(labelSelectionString);
    // If we failed to parse the parameter
    if (labelSelectionString != "" && selectedLabels.isEmpty()) {
        CAMITK_WARNING("Cannot parse Selection Of Label Values: " + labelSelectionString)
        return Action::ERROR;
    }

    // Process each target
    for (Component* selected : getTargets()) {
        ImageComponent* input = dynamic_cast <ImageComponent*>(selected);

        //-- Compute the maximum and the minimum possible value for the selected images
        double min = VTK_DOUBLE_MAX;
        double max = VTK_DOUBLE_MIN;

        vtkSmartPointer<vtkImageData> inputImage = input->getImageData();
        double* imgRange = inputImage->GetScalarRange();

        if (min > imgRange[0]) {
            min = imgRange[0];
        }

        if (max < imgRange[1]) {
            max = imgRange[1];
        }

        // Check that we have a label image (no more than maximumNumberOfLabels values)
        vtkSmartPointer<vtkDataArray> scalars = input->getImageData()->GetPointData()->GetScalars();
        std::unordered_set<double> uniqueValues;
        // Loop through all the scalar values
        for (vtkIdType i = 0; i < scalars->GetNumberOfTuples(); ++i) {
            double value = scalars->GetComponent(i, 0);
            uniqueValues.insert(value);
        }
        if (uniqueValues.size() >= maximumNumberOfLabels) {
            if (QMessageBox::warning(nullptr, QString("More than %1 levels detected!").arg(maximumNumberOfLabels), QString("This image has %1 different values, it is probably not a label image\nAre you sure you want to proceed?").arg(uniqueValues.size()), QMessageBox::Yes | QMessageBox::No) == QMessageBox::No) {
                return Action::ABORTED;
            }
        }

        // From https://examples.vtk.org/site/Cxx/Medical/GenerateModelsFromLabels/
        vtkNew<vtkImageAccumulate> histogram;
        vtkNew<vtkDiscreteMarchingCubes> discreteCubes;
        vtkNew<vtkDiscreteFlyingEdges3D> discreteEdges;
        vtkNew<vtkWindowedSincPolyDataFilter> smoother;
        vtkNew<vtkThreshold> selector;
        vtkNew<vtkMaskFields> scalarsOff;
        vtkNew<vtkGeometryFilter> geometry;

        unsigned int startLabel = min + 1; // min is background, do not use it
        unsigned int endLabel = max;
        bool smoothing = getParameterValue("Smooth Meshes").toBool();
        unsigned int smoothingIterations = getParameterValue("Iterations").toInt();
        double passBand = getParameterValue("Passband").toDouble();
        double featureAngle = getParameterValue("Feature Angle").toDouble();


        // 1 - Compute histogram of the the labels
        histogram->SetInputData(input->getImageData());
        histogram->SetComponentExtent(0, endLabel, 0, 0, 0, 0); // Labels are supposed to be in ]0, 65535]
        histogram->SetComponentOrigin(0, 0, 0);
        histogram->SetComponentSpacing(1, 1, 1);
        histogram->Update();

        // 2 - marching cubes / Flying edges
        if (method == 0) {
            // Marching cubes
            discreteCubes->SetInputData(input->getImageData());
            discreteCubes->GenerateValues(endLabel - startLabel + 1, startLabel, endLabel);
            smoother->SetInputConnection(discreteCubes->GetOutputPort());
        }
        else {
            // Flying edges
            discreteEdges->SetInputData(input->getImageData());
            discreteEdges->GenerateValues(endLabel - startLabel + 1, startLabel, endLabel);
            smoother->SetInputConnection(discreteEdges->GetOutputPort());
        }

        // 3 - smooth filter parameters
        if (smoothing) {
            smoother->SetNumberOfIterations(smoothingIterations);
            smoother->BoundarySmoothingOff();
            smoother->FeatureEdgeSmoothingOff();
            smoother->SetFeatureAngle(featureAngle);
            smoother->SetPassBand(passBand);
            smoother->NonManifoldSmoothingOn();
            smoother->NormalizeCoordinatesOn();
        }

        // 4 - Separate models
        if (smoothing) {
            selector->SetInputConnection(smoother->GetOutputPort());
        }
        else {
            if (method == 0) {
                selector->SetInputConnection(discreteCubes->GetOutputPort());
            }
            else {
                selector->SetInputConnection(discreteEdges->GetOutputPort());
            }
        }
        if (method == 0) {
            // Marching cube uses cells
            selector->SetInputArrayToProcess(0, 0, 0,
                                             vtkDataObject::FIELD_ASSOCIATION_CELLS,
                                             vtkDataSetAttributes::SCALARS);
        }
        else {
            // Flying edges uses points
            selector->SetInputArrayToProcess(0, 0, 0,
                                             vtkDataObject::FIELD_ASSOCIATION_POINTS,
                                             vtkDataSetAttributes::SCALARS);
        }

        // 5 - Strip the scalars from the output
        scalarsOff->SetInputConnection(selector->GetOutputPort());
        scalarsOff->CopyAttributeOff(vtkMaskFields::POINT_DATA,
                                     vtkDataSetAttributes::SCALARS);
        scalarsOff->CopyAttributeOff(vtkMaskFields::CELL_DATA,
                                     vtkDataSetAttributes::SCALARS);

        geometry->SetInputConnection(scalarsOff->GetOutputPort());
        CAMITK_TRACE("VTK Connections done");

        // Create a mesh for each label
        for (unsigned int i = startLabel; i <= endLabel; i++) {
            // if the label has no voxel, skip it
            double frequency = histogram->GetOutput()->GetPointData()->GetScalars()->GetTuple1(i);
            // Ignore empty labels and labels not chosen by the user
            if (frequency > 0.0 && (labelSelectionString == "" || selectedLabels.contains(i))) {
                // select the cells for a given label
                selector->SetLowerThreshold(i);
                selector->SetUpperThreshold(i);

                // Recompute the mesh
                geometry->Update();

                // Create the MeshComponent
                MeshComponent* outputComp = new MeshComponent(geometry->GetOutput(),
                        Application::getUniqueComponentName(input->getName() + QString(" (level %1)").arg(i)));
                outputComp->setFrame(TransformationManager::getFrameOfReferenceOwnership(input->getDataFrame()));

                // To avoid having dozens of unsaved meshes, just mark them as unmodified
                outputComp->setModified(false);

                // Add an automatic color for meshes
                QColor color = MeshesFromImage::nextColor();
                outputComp->setColor(color.redF(), color.greenF(), color.blueF());
            }
        }
    }

    // Remove waiting cursor
    Application::restoreOverrideCursor();

    refreshApplication();
    return SUCCESS;
}

// -------------------- targetDefined --------------------
void MeshesFromImage::targetDefined() {
}

// -------------------- parameterChanged --------------------
void MeshesFromImage::parameterChanged(QString parameterName) {
    // Here do something when the given parameter has changed
}

