////////////////////////////////////////////////////////////////////////////////////////
//
//  Copyright 2023 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_OVITO_CLASS(GrainSegmentationModifier);
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);

/******************************************************************************
* Constructs the modifier object.
******************************************************************************/
GrainSegmentationModifier::GrainSegmentationModifier(ObjectInitializationFlags flags) : AsynchronousModifier(flags),
        _mergeAlgorithm(GraphClusteringAutomatic),
        _minGrainAtomCount(100),
        _mergingThreshold(0.0),
        _orphanAdoption(true),
        _outputBonds(false),
        _colorParticlesByGrain(true),
        _handleCoherentInterfaces(true)
{
    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::PreliminaryStateAvailable);
    }
    AsynchronousModifier::propertyChanged(field);
}

/******************************************************************************
* Creates a computation engine that will compute the modifier's results.
******************************************************************************/
Future<AsynchronousModifier::EnginePtr> GrainSegmentationModifier::createEngine(const ModifierEvaluationRequest& request, const PipelineFlowState& input)
{
    // Get modifier input.
    const Particles* particles = input.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("Correspondences", Property::Int64);

    const SimulationCell* simCell = input.expectObject<SimulationCell>();
    if(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.
    return std::make_shared<GrainSegmentationEngine1>(
            request,
            particles,
            posProperty,
            structureProperty,
            orientationProperty,
            correspondenceProperty,
            simCell,
            mergeAlgorithm(),
            handleCoherentInterfaces(),
            outputBonds());
}

/******************************************************************************
* Injects the computed results of the engine into the data pipeline.
******************************************************************************/
void GrainSegmentationEngine1::applyResults(const ModifierEvaluationRequest& request, PipelineFlowState& state)
{
    GrainSegmentationModifier* modifier = static_object_cast<GrainSegmentationModifier>(request.modifier());
    OVITO_ASSERT(modifier);

    Particles* particles = state.expectMutableObject<Particles>();
    particles->verifyIntegrity();

    if(_inputFingerprint.hasChanged(particles))
        throw Exception(GrainSegmentationModifier::tr("Cached modifier results are obsolete, because the number or the storage order of input particles has changed."));

    // Output the edges of the neighbor graph.
    if(_outputBondsToPipeline && modifier->outputBonds()) {

        std::vector<Bond> bonds;
        std::vector<FloatType> disorientations;
        BufferReadAccess<Point3> positionsArray(particles->expectProperty(Particles::PositionProperty));

        for (auto edge: neighborBonds()) {
            if (isCrystallineBond(edge)) {
                Bond bond;
                bond.index1 = edge.a;
                bond.index2 = edge.b;
                disorientations.push_back(edge.disorientation);

                // Determine PBC bond shift using minimum image convention.
                Vector3 delta = positionsArray[bond.index1] - positionsArray[bond.index2];
                for(size_t dim = 0; dim < 3; dim++) {
                    if(cell() && cell()->pbcFlags()[dim])
                        bond.pbcShift[dim] = (int)std::floor(cell()->inverseMatrix().prodrow(delta, dim) + FloatType(0.5));
                    else
                        bond.pbcShift[dim] = 0;
                }

                bonds.push_back(bond);
            }
        }

        // Output disorientation angles as a bond property.
        PropertyPtr neighborDisorientationAngles = Bonds::OOClass().createUserProperty(DataBuffer::Uninitialized, bonds.size(), DataBuffer::FloatDefault, 1, QStringLiteral("Disorientation"));
        BufferWriteAccess<FloatType, access_mode::discard_write> disorientationAnglesAccess(neighborDisorientationAngles);
        for(size_t i = 0; i < disorientations.size(); i++) {
            disorientationAnglesAccess[i] = disorientations[i];
        }
        disorientationAnglesAccess.reset();

        particles->addBonds(bonds, modifier->bondsVis(), { std::move(neighborDisorientationAngles) });
    }

    // Output a data plot with the dendrogram points.
    if(mergeSize() && mergeDistance())
        state.createObject<DataTable>(QStringLiteral("grains-merge"), request.modificationNode(), DataTable::Scatter, GrainSegmentationModifier::tr("Merge size vs. distance"), mergeSize(), mergeDistance());

    // Output a data plot with the log-log dendrogram points.
    if(logMergeSize() && logMergeDistance())
        state.createObject<DataTable>(QStringLiteral("grains-log"), request.modificationNode(), DataTable::Scatter, GrainSegmentationModifier::tr("Log distance vs. log merge size"), logMergeDistance(), logMergeSize());

    if(modifier->mergeAlgorithm() == GrainSegmentationModifier::GraphClusteringAutomatic)
        state.addAttribute(QStringLiteral("GrainSegmentation.auto_merge_threshold"), QVariant::fromValue(suggestedMergingThreshold()), request.modificationNode());
}

/******************************************************************************
* Injects the computed results of the engine into the data pipeline.
******************************************************************************/
void GrainSegmentationEngine2::applyResults(const ModifierEvaluationRequest& request, PipelineFlowState& state)
{
    // Output the results from the 1st algorithm stage.
    _engine1->applyResults(request, state);

    GrainSegmentationModifier* modifier = static_object_cast<GrainSegmentationModifier>(request.modifier());
    OVITO_ASSERT(modifier);

    Particles* particles = state.expectMutableObject<Particles>();

    // Output per-particle properties.
    if(atomClusters()) {
        particles->createProperty(atomClusters());

        if(modifier->colorParticlesByGrain()) {

            // Assign colors to particles according to the grains they belong to.
            BufferReadAccess<ColorG> grainColorsArray(_grainColors);
            BufferWriteAccess<ColorG, access_mode::discard_write> particleColorsArray = particles->createProperty(Particles::ColorProperty);
            boost::transform(BufferReadAccess<int64_t>(atomClusters()), particleColorsArray.begin(), [&](int64_t cluster) {
                if(cluster != 0)
                    return grainColorsArray[cluster - 1];
                else
                    return ColorG(0.8f, 0.8f, 0.8f); // Special color for non-crystalline particles not part of any grain.
            });
        }
    }

    // Output a data table with the list of grains.
    // The X-column consists of the grain IDs, the Y-column contains the grain sizes.
    DataTable* grainTable = state.createObject<DataTable>(QStringLiteral("grains"), request.modificationNode(), DataTable::Scatter, GrainSegmentationModifier::tr("Grain list"), _grainSizes, _grainIds);
    // Add extra columns to the table containing other per-grain data.
    grainTable->createProperty(_grainColors);
    grainTable->createProperty(_grainStructureTypes);
    grainTable->createProperty(_grainOrientations);

    size_t numGrains = 0;
    if(atomClusters()->size() != 0)
        numGrains = *boost::max_element(BufferReadAccess<int64_t>(atomClusters()));

    state.addAttribute(QStringLiteral("GrainSegmentation.grain_count"), QVariant::fromValue(numGrains), request.modificationNode());

    state.setStatus(PipelineStatus(
        GrainSegmentationModifier::tr("Found %1 grains").arg(numGrains),
        GrainSegmentationModifier::tr("%1 grains").arg(numGrains)));
}

}   // End of namespace
