////////////////////////////////////////////////////////////////////////////////////////
//
//  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/modifier/analysis/StructureIdentificationModifier.h>
#include <ovito/particles/objects/Bonds.h>
#include <ovito/particles/objects/Particles.h>

namespace Ovito {

/**
 * \brief A modifier that performs the common neighbor analysis (CNA) to identify
 *        local coordination structures.
 */
class OVITO_PARTICLES_EXPORT CommonNeighborAnalysisModifier : public StructureIdentificationModifier
{
    OVITO_CLASS(CommonNeighborAnalysisModifier)

public:

    enum CNAMode {
        FixedCutoffMode,    ///< Performs the conventional CNA using a global cutoff radius.
        AdaptiveCutoffMode, ///< Performs the adaptive CNA, which picks an optimal cutoff for each atom.
        IntervalCutoffMode, ///< Performs the interval CNA, which performs an interval analysis to find the optimal cutoff for each atom.
        BondMode,           ///< Performs the CNA based on the existing network of bonds.
    };
    Q_ENUM(CNAMode);

    /// The maximum number of neighbor atoms taken into account for the common neighbor analysis.
    static constexpr int MAX_NEIGHBORS = 14;

    /// The structure types recognized by the common neighbor analysis.
    enum StructureType {
        OTHER = 0,              //< Unidentified structure
        FCC,                    //< Face-centered cubic
        HCP,                    //< Hexagonal close-packed
        BCC,                    //< Body-centered cubic
        ICO,                    //< Icosahedral structure

        NUM_STRUCTURE_TYPES     //< This just counts the number of defined structure types.
    };
    Q_ENUM(StructureType);

    /// Pair of neighbor atoms that form a bond (bit-wise storage).
    using CNAPairBond = unsigned int;

    /**
     * A bit-flag array indicating which pairs of neighbors are bonded
     * and which are not.
     */
    struct NeighborBondArray
    {
        /// Two-dimensional bit array that stores the bonds between neighbors.
        unsigned int neighborArray[32];

        /// Resets all bits.
        NeighborBondArray() {
            memset(neighborArray, 0, sizeof(neighborArray));
        }

        /// Returns whether two nearest neighbors have a bond between them.
        inline bool neighborBond(int neighborIndex1, int neighborIndex2) const {
            OVITO_ASSERT(neighborIndex1 < 32);
            OVITO_ASSERT(neighborIndex2 < 32);
            return (neighborArray[neighborIndex1] & (1<<neighborIndex2));
        }

        /// Sets whether two nearest neighbors have a bond between them.
        inline void setNeighborBond(int neighborIndex1, int neighborIndex2, bool bonded) {
            OVITO_ASSERT(neighborIndex1 < 32);
            OVITO_ASSERT(neighborIndex2 < 32);
            if(bonded) {
                neighborArray[neighborIndex1] |= (1<<neighborIndex2);
                neighborArray[neighborIndex2] |= (1<<neighborIndex1);
            }
            else {
                neighborArray[neighborIndex1] &= ~(1<<neighborIndex2);
                neighborArray[neighborIndex2] &= ~(1<<neighborIndex1);
            }
        }
    };

public:

    /// Constructor.
    void initializeObject(ObjectInitializationFlags flags);

    /// Find all atoms that are nearest neighbors of the given pair of atoms.
    static int findCommonNeighbors(const NeighborBondArray& neighborArray, int neighborIndex, unsigned int& commonNeighbors);

    /// Finds all bonds between common nearest neighbors.
    static int findNeighborBonds(const NeighborBondArray& neighborArray, unsigned int commonNeighbors, int numNeighbors, CNAPairBond* neighborBonds);

    /// Find all chains of bonds between common neighbors and determine the length
    /// of the longest continuous chain.
    static int calcMaxChainLength(CNAPairBond* neighborBonds, int numBonds);

protected:

    /// Creates the algorithm that will perform the structure identification.
    virtual std::shared_ptr<Algorithm> createAlgorithm(const ModifierEvaluationRequest& request, const PipelineFlowState& input) override;

private:

    /// Base class for CNA compute engines.
    class CNAAlgorithm : public StructureIdentificationModifier::Algorithm
    {
    public:

        /// Inherit constructor.
        using Algorithm::Algorithm;

        /// Computes the structure identification statistics.
        virtual std::vector<int64_t> computeStructureStatistics(const Property* structures, PipelineFlowState& state, const OOWeakRef<const PipelineNode>& createdByNode, const std::any& modifierParameters) const override;

    protected:

        /// Determines the coordination structure of a single particle using the common neighbor analysis method.
        StructureType determineStructureAdaptive(NearestNeighborFinder& neighList, size_t particleIndex);

        /// Determines the coordination structure of a single particle using the common neighbor analysis method.
        StructureType determineStructureInterval(NearestNeighborFinder& neighList, size_t particleIndex);

        /// Determines the coordination structure of a single particle using the common neighbor analysis method.
        template<class NeighborFinderType>
        static StructureType determineStructureFixed(size_t particleIndex, const NeighborFinderType& neighFinder, const std::array<bool, NUM_STRUCTURE_TYPES>& typesToIdentify);

        /// Determines the coordination signature for structures with 12 neighbors.
        static StructureType analyzeSmallSignature(NeighborBondArray& neighborArray);

        /// Determines the coordination signature for structures with 14 neighbors.
        static StructureType analyzeLargeSignature(NeighborBondArray& neighborArray);
    };

    /// Analysis engine that performs the conventional common neighbor analysis.
    class FixedCNAAlgorithm : public CNAAlgorithm
    {
    public:

        /// Constructor.
        FixedCNAAlgorithm(const StructureIdentificationModifier& modifier, const PipelineFlowState& input, FloatType cutoff) :
            CNAAlgorithm(modifier, input),
            _cutoff(cutoff) {}

        /// Performs the atomic structure classification.
        virtual void identifyStructures() override;

    private:

        /// The CNA cutoff radius.
        const FloatType _cutoff;
    };

    /// Analysis engine that performs the adaptive common neighbor analysis.
    class AdaptiveCNAAlgorithm : public CNAAlgorithm
    {
    public:

        /// Inherit constructor.
        using CNAAlgorithm::CNAAlgorithm;

        /// Performs the atomic structure classification.
        virtual void identifyStructures() override;
    };

    /// Analysis engine that performs the interval common neighbor analysis.
    class IntervalCNAAlgorithm : public CNAAlgorithm
    {
    public:

        /// Inherit constructor.
        using CNAAlgorithm::CNAAlgorithm;

        /// Performs the atomic structure classification.
        virtual void identifyStructures() override;
    };

    /// Analysis engine that performs the common neighbor analysis based on existing bonds.
    class BondCNAAlgorithm : public CNAAlgorithm
    {
    public:

        /// Constructor.
        BondCNAAlgorithm(const StructureIdentificationModifier& modifier, const PipelineFlowState& input) : CNAAlgorithm(modifier, input)
        {
            const Bonds* bonds = particles()->expectBonds();
            bonds->verifyIntegrity();
            _bondTopology = bonds->expectProperty(Bonds::TopologyProperty);
            _bondPeriodicImages = bonds->getProperty(Bonds::PeriodicImageProperty);
            _cnaIndices = Bonds::OOClass().createUserProperty(DataBuffer::Uninitialized, bonds->elementCount(), Property::Int32, 3, tr("CNA Indices"));
        }

        /// Performs the atomic structure classification.
        virtual void identifyStructures() override;

        /// Computes the structure identification statistics.
        virtual std::vector<int64_t> computeStructureStatistics(const Property* structures, PipelineFlowState& state, const OOWeakRef<const PipelineNode>& createdByNode, const std::any& modifierParameters) const override;

        /// Returns the output bonds property that stores the computed CNA indices.
        const PropertyPtr& cnaIndices() const { return _cnaIndices; }

        /// Returns the input bonds between particles.
        const ConstPropertyPtr& bondTopology() const { return _bondTopology; }

        /// Returns the periodic image shift vector for input bonds.
        const ConstPropertyPtr& bondPeriodicImages() const { return _bondPeriodicImages; }

    private:

        ConstPropertyPtr _bondTopology;
        ConstPropertyPtr _bondPeriodicImages;
        PropertyPtr _cnaIndices;
    };

    /// The cutoff radius used for the conventional CNA.
    DECLARE_MODIFIABLE_PROPERTY_FIELD_FLAGS(FloatType{3.2}, cutoff, setCutoff, PROPERTY_FIELD_MEMORIZE);

    /// Controls how the CNA is performed.
    DECLARE_MODIFIABLE_PROPERTY_FIELD_FLAGS(CNAMode{AdaptiveCutoffMode}, mode, setMode, PROPERTY_FIELD_MEMORIZE);
};

}   // End of namespace
