////////////////////////////////////////////////////////////////////////////////////////
//
//  Copyright 2026 OVITO GmbH, Germany
//
//  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.
//
////////////////////////////////////////////////////////////////////////////////////////

#pragma once


#include <ovito/particles/Particles.h>
#include <ovito/particles/objects/Bonds.h>
#include <ovito/stdobj/simcell/SimulationCell.h>
#include <ovito/stdobj/properties/Property.h>
#include <ovito/core/dataset/pipeline/Modifier.h>

#include <utility>

namespace Ovito {

/**
 * \brief Extends the current particle selection by adding particles to the selection
 *        that are neighbors of an already selected particle.
 */
class OVITO_PARTICLES_EXPORT ExpandSelectionModifier : public Modifier
{
    /// Give this modifier class its own metaclass.
    class OOMetaClass : public Modifier::OOMetaClass
    {
    public:

        /// Inherit constructor from base metaclass.
        using Modifier::OOMetaClass::OOMetaClass;

        /// Asks the metaclass whether the modifier can be applied to the given input data.
        virtual bool isApplicableTo(const DataCollection& input) const override;
    };

    OVITO_CLASS_META(ExpandSelectionModifier, OOMetaClass)

public:
    enum ExpansionMode
    {
        BondedNeighbors,   ///< Expands the selection to particles that are bonded to an already selected particle.
        CutoffRange,       ///< Expands the selection to particles that within a cutoff range of an already selected particle.
        NearestNeighbors,  ///< Expands the selection to the N nearest particles of already selected particles.
        Molecule           ///< Expands the selection to all particles in the molecule
    };
    Q_ENUM(ExpansionMode);

    /// Compile-time constant for the maximum number of nearest neighbors that can be taken into account.
    enum { MAX_NEAREST_NEIGHBORS = 30 };

public:

    /// Is called by the pipeline system before a new modifier evaluation begins.
    virtual void preevaluateModifier(const ModifierEvaluationRequest& request, PipelineEvaluationResult::EvaluationTypes& evaluationTypes, TimeInterval& validityInterval) const override;

    /// Modifies the input data.
    virtual Future<PipelineFlowState> evaluateModifier(const ModifierEvaluationRequest& request, PipelineFlowState&& state) override;

    /// Indicates that a preliminary viewport update will be performed immediately after this modifier
	/// has computed new results.
    virtual bool shouldRefreshViewportsAfterEvaluation() override { return true; }

private:

    /// The modifier's compute engine.
    class ExpandSelectionEngine
    {
    public:

        /// Constructor.
        ExpandSelectionEngine(OOWeakRef<const PipelineNode> createdByNode,
                              ConstPropertyPtr positions,
                              const SimulationCell* simCell,
                              ConstPropertyPtr inputSelection,
                              int numIterations)
            : _createdByNode(std::move(createdByNode)),
              _numIterations(numIterations),
              _positions(std::move(positions)),
              _simCell(simCell),
              _inputSelection(inputSelection),
              _outputSelection(inputSelection.makeCopy())
        {
        }

        /// Destructor.
        virtual ~ExpandSelectionEngine() = default;

        /// Computes the modifier's results.
        virtual void perform();

        /// Performs one iteration of the expansion.
        virtual void expandSelection(TaskProgress& progress) = 0;

        /// Injects the computed results into the data pipeline.
        void applyResults(PipelineFlowState& state);

        const PropertyPtr& outputSelection() { return _outputSelection; }

        void setOutputSelection(PropertyPtr ptr) { _outputSelection = std::move(ptr); }

        size_t numSelectedParticlesInput() const { return _numSelectedParticlesInput; }

        size_t numSelectedParticlesOutput() const { return _numSelectedParticlesOutput; }

        void setNumSelectedParticlesInput(size_t count) { _numSelectedParticlesInput = count; }

        void setNumSelectedParticlesOutput(size_t count) { _numSelectedParticlesOutput = count; }

        const SimulationCellData& simCell() const { return _simCell; }

        const ConstPropertyPtr& positions() const { return _positions; }

        const ConstPropertyPtr& inputSelection() const { return _inputSelection; }

        const OOWeakRef<const PipelineNode>& createdByNode() const { return _createdByNode; }

    protected:

        const int _numIterations;
        const SimulationCellData _simCell;
        ConstPropertyPtr _positions;
        ConstPropertyPtr _inputSelection;
        PropertyPtr _outputSelection;
        size_t _numSelectedParticlesInput;
        size_t _numSelectedParticlesOutput;
        OOWeakRef<const PipelineNode> _createdByNode;
    };

    /// Computes the expanded selection by using the nearest neighbor criterion.
    class ExpandSelectionNearestEngine : public ExpandSelectionEngine
    {
    public:

        /// Constructor.
        ExpandSelectionNearestEngine(OOWeakRef<const PipelineNode> createdByNode, ConstPropertyPtr positions, const SimulationCell* simCell, ConstPropertyPtr inputSelection, int numIterations, int numNearestNeighbors) :
            ExpandSelectionEngine(std::move(createdByNode), std::move(positions), simCell, std::move(inputSelection), numIterations),
            _numNearestNeighbors(numNearestNeighbors) {}

        /// Expands the selection by one step.
        virtual void expandSelection(TaskProgress& progress) override;

    private:

        const int _numNearestNeighbors;
    };

    /// Computes the expanded selection when using a cutoff range criterion.
    class ExpandSelectionCutoffEngine : public ExpandSelectionEngine
    {
    public:

        /// Constructor.
        ExpandSelectionCutoffEngine(OOWeakRef<const PipelineNode> createdByNode, ConstPropertyPtr positions, const SimulationCell* simCell, ConstPropertyPtr inputSelection, int numIterations, FloatType cutoff) :
            ExpandSelectionEngine(std::move(createdByNode), std::move(positions), simCell, std::move(inputSelection), numIterations),
            _cutoffRange(cutoff) {}

        /// Expands the selection by one step.
        virtual void expandSelection(TaskProgress& progress) override;

    private:

        const FloatType _cutoffRange;
    };

    /// Computes the expanded selection when using bonds.
    class ExpandSelectionBondedEngine : public ExpandSelectionEngine
    {
    public:

        /// Constructor.
        ExpandSelectionBondedEngine(OOWeakRef<const PipelineNode> createdByNode, ConstPropertyPtr positions, const SimulationCell* simCell, ConstPropertyPtr inputSelection, int numIterations, ConstPropertyPtr bondTopology) :
            ExpandSelectionEngine(std::move(createdByNode), std::move(positions), simCell, std::move(inputSelection), numIterations),
            _bondTopology(std::move(bondTopology)) {}

        /// Expands the selection by one step.
        virtual void expandSelection(TaskProgress& progress) override;

    private:

        ConstPropertyPtr _bondTopology;
    };

    /// Computes the expanded selection by using the molecule id.
    class ExpandSelectionMoleculeEngine : public ExpandSelectionEngine
    {
    public:
        /// Constructor.
        ExpandSelectionMoleculeEngine(OOWeakRef<const PipelineNode> createdByNode,
                                      ConstPropertyPtr moleculeIdProperty,
                                      ConstPropertyPtr inputSelection)
            : ExpandSelectionEngine(std::move(createdByNode), nullptr, nullptr, std::move(inputSelection), /*num iterations*/ 1),
              _moleculeIdProperty(std::move(moleculeIdProperty))
        {
        }

        /// Expands the selection by one step.
        virtual void expandSelection(TaskProgress& progress) override;

    private:
        ConstPropertyPtr _moleculeIdProperty;
    };

private:

    /// The expansion mode.
    DECLARE_MODIFIABLE_PROPERTY_FIELD_FLAGS(ExpansionMode{CutoffRange}, mode, setMode, PROPERTY_FIELD_MEMORIZE);

    /// The selection cutoff range.
    DECLARE_MODIFIABLE_PROPERTY_FIELD_FLAGS(FloatType{3.2}, cutoffRange, setCutoffRange, PROPERTY_FIELD_MEMORIZE);

    /// The number of nearest neighbors to select.
    DECLARE_MODIFIABLE_PROPERTY_FIELD_FLAGS(int{1}, numNearestNeighbors, setNumNearestNeighbors, PROPERTY_FIELD_MEMORIZE);

    /// The number of expansion steps to perform.
    DECLARE_MODIFIABLE_PROPERTY_FIELD(int{1}, numberOfIterations, setNumberOfIterations);
};

}   // End of namespace
