////////////////////////////////////////////////////////////////////////////////////////
//
//  Copyright 2026 OVITO GmbH, Germany
//  Copyright 2020 Peter Mahler Larsen
//
//  This file is part of OVITO (Open Visualization Tool).
//
//  OVITO is free software; you can redistribute it and/or modify it either under the
//  terms of the GNU General Public License version 3 as published by the Free Software
//  Foundation (the "GPL") or, at your option, under the terms of the MIT License.
//  If you do not alter this notice, a recipient may use your version of this
//  file under either the GPL or the MIT License.
//
//  You should have received a copy of the GPL along with this program in a
//  file LICENSE.GPL.txt.  You should have received a copy of the MIT License along
//  with this program in a file LICENSE.MIT.txt
//
//  This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY KIND,
//  either express or implied. See the GPL or the MIT License for the specific language
//  governing rights and limitations.
//
////////////////////////////////////////////////////////////////////////////////////////

#include <ovito/crystalanalysis/CrystalAnalysis.h>
#include <ovito/particles/objects/BondsVis.h>
#include <ovito/particles/objects/Particles.h>
#include <ovito/stdobj/simcell/SimulationCell.h>
#include <ovito/stdobj/properties/Property.h>
#include <ovito/stdobj/table/DataTable.h>
#include <ovito/core/utilities/units/UnitsManager.h>
#include <ovito/core/dataset/pipeline/ModificationNode.h>
#include <ovito/core/dataset/DataSet.h>
#include <ovito/core/app/Application.h>
#include "GrainSegmentationModifier.h"
#include "GrainSegmentationEngine.h"

#include <ptm/ptm_functions.h>

namespace Ovito {

IMPLEMENT_CREATABLE_OVITO_CLASS(GrainSegmentationModifier);
OVITO_CLASSINFO(GrainSegmentationModifier, "DisplayName", "Grain segmentation");
OVITO_CLASSINFO(GrainSegmentationModifier, "Description", "Segment grains in polycrystalline structures.");
OVITO_CLASSINFO(GrainSegmentationModifier, "ModifierCategory", "Analysis");
DEFINE_PROPERTY_FIELD(GrainSegmentationModifier, mergeAlgorithm);
DEFINE_PROPERTY_FIELD(GrainSegmentationModifier, handleCoherentInterfaces);
DEFINE_PROPERTY_FIELD(GrainSegmentationModifier, mergingThreshold);
DEFINE_PROPERTY_FIELD(GrainSegmentationModifier, minGrainAtomCount);
DEFINE_PROPERTY_FIELD(GrainSegmentationModifier, orphanAdoption);
DEFINE_PROPERTY_FIELD(GrainSegmentationModifier, outputBonds);
DEFINE_PROPERTY_FIELD(GrainSegmentationModifier, colorParticlesByGrain);
DEFINE_REFERENCE_FIELD(GrainSegmentationModifier, bondsVis);
SET_PROPERTY_FIELD_LABEL(GrainSegmentationModifier, mergeAlgorithm, "Algorithm");
SET_PROPERTY_FIELD_LABEL(GrainSegmentationModifier, handleCoherentInterfaces, "Handle coherent interfaces/stacking faults");
SET_PROPERTY_FIELD_LABEL(GrainSegmentationModifier, mergingThreshold, "Merge threshold");
SET_PROPERTY_FIELD_LABEL(GrainSegmentationModifier, minGrainAtomCount, "Minimum grain size (# of atoms)");
SET_PROPERTY_FIELD_LABEL(GrainSegmentationModifier, orphanAdoption, "Adopt orphan atoms");
SET_PROPERTY_FIELD_LABEL(GrainSegmentationModifier, outputBonds, "Output bonds");
SET_PROPERTY_FIELD_LABEL(GrainSegmentationModifier, colorParticlesByGrain, "Color particles by grain");
SET_PROPERTY_FIELD_LABEL(GrainSegmentationModifier, bondsVis, "Bonds display");
SET_PROPERTY_FIELD_UNITS_AND_MINIMUM(GrainSegmentationModifier, minGrainAtomCount, IntegerParameterUnit, 1);

/******************************************************************************
* Constructor.
******************************************************************************/
void GrainSegmentationModifier::initializeObject(ObjectInitializationFlags flags)
{
    Modifier::initializeObject(flags);

    if(!flags.testFlag(ObjectInitializationFlag::DontInitializeObject)) {
        // Create the vis element for rendering the bonds generated by the modifier.
        setBondsVis(OORef<BondsVis>::create(flags));
    }
}

/******************************************************************************
* Asks the modifier whether it can be applied to the given input data.
******************************************************************************/
bool GrainSegmentationModifier::OOMetaClass::isApplicableTo(const DataCollection& input) const
{
    return input.containsObject<Particles>();
}

/******************************************************************************
* Is called when the value of a property of this object has changed.
******************************************************************************/
void GrainSegmentationModifier::propertyChanged(const PropertyFieldDescriptor* field)
{
    if(field == PROPERTY_FIELD(mergingThreshold) || field == PROPERTY_FIELD(minGrainAtomCount) || field == PROPERTY_FIELD(colorParticlesByGrain) || field == PROPERTY_FIELD(orphanAdoption)) {
        // Immediately update viewports if parameters are changed by the user that don't require a full recalculation.
        notifyDependents(ReferenceEvent::InteractiveStateAvailable);
    }
    Modifier::propertyChanged(field);
}

/******************************************************************************
 * Is called by the pipeline system before a new modifier evaluation begins.
 ******************************************************************************/
void GrainSegmentationModifier::preevaluateModifier(const ModifierEvaluationRequest& request, PipelineEvaluationResult::EvaluationTypes& evaluationTypes, TimeInterval& validityInterval) const
{
    // Indicate that we will do different computations depending on whether the pipeline is evaluated in interactive mode or not.
    if(request.interactiveMode())
        evaluationTypes = PipelineEvaluationResult::EvaluationType::Interactive;
    else
        evaluationTypes = PipelineEvaluationResult::EvaluationType::Noninteractive;
}

/******************************************************************************
* Modifies the input data.
******************************************************************************/
Future<PipelineFlowState> GrainSegmentationModifier::evaluateModifier(const ModifierEvaluationRequest& request, PipelineFlowState&& state)
{
    // In interactive mode, do not perform a real computation. Instead, reuse an old result from the cached state if available.
    if(request.interactiveMode()) {
        if(PipelineFlowState cachedState = request.modificationNode()->getCachedPipelineNodeOutput(request.time(), true)) {
//            Particles* particles = state.expectMutableObject<Particles>();
//            particles->verifyIntegrity();
//            reuseCachedState(request, particles, state, cachedState);
        }
        return std::move(state);
    }

    // Phase I:
    auto engin1Future = request.modificationNode()->partialResultsCache().getOrCompute(state.data(), [&]() {

        // Get modifier input.
        const Particles* particles = state.expectObject<Particles>();
        particles->verifyIntegrity();
        const Property* posProperty = particles->expectProperty(Particles::PositionProperty);

        // Make sure the PTM modifier has been executed first and its output is available.
        const Property* structureProperty = particles->getProperty(Particles::StructureTypeProperty);
        if(!structureProperty)
            throw Exception(tr("Grain segmentation requires Polyhedral Template Matching (PTM) output. Please insert the PTM modifier into the pipeline first."));
        const Property* orientationProperty = particles->getProperty(Particles::OrientationProperty);
        if(!orientationProperty)
            throw Exception(tr("Grain segmentation requires lattice orientation information. Please activate the 'Lattice orientations' option of the PTM modifier."));
        const Property* correspondenceProperty = particles->expectProperty(QStringLiteral("Correspondences"), Property::Int64);

        const SimulationCell* simCell = state.getObject<SimulationCell>();
        if(simCell && simCell->is2D())
            throw Exception(tr("The grain segmentation modifier does not support 2d simulation cells."));

        // Initialize PTM library.
        ptm_initialize_global();

        // Create engine object. Pass all relevant modifier parameters to the engine as well as the input data.
        auto engine1 = std::make_shared<GrainSegmentationEngine1>(
                posProperty,
                structureProperty,
                orientationProperty,
                correspondenceProperty,
                simCell,
                mergeAlgorithm(),
                handleCoherentInterfaces(),
                outputBonds());

        // Perform the computation in a separate thread.
        return asyncLaunch([engine1 = std::move(engine1)]() mutable {
            engine1->perform();
            return std::move(engine1);
        });
    });

    // Phase II:
    return engin1Future.then(ObjectExecutor(this), [this, state = std::move(state), createdByNode = request.modificationNodeWeak()](std::shared_ptr<const GrainSegmentationEngine1> engine1) {

        auto engine2 = std::make_shared<GrainSegmentationEngine2>(
            std::move(engine1),
            mergingThreshold(),
            orphanAdoption(),
            minGrainAtomCount(),
            colorParticlesByGrain()
        );

        // Perform the computation in a separate thread.
        return asyncLaunch([state = std::move(state), engine2 = std::move(engine2), createdByNode = std::move(createdByNode), bondsVis = OORef<BondsVis>(bondsVis())]() mutable {
            engine2->perform();
            engine2->applyResults(state, createdByNode, bondsVis);
            return std::move(state);
        });
    });
}

}   // End of namespace
