////////////////////////////////////////////////////////////////////////////////////////
//
//  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.
//
////////////////////////////////////////////////////////////////////////////////////////

#include <ovito/particles/Particles.h>
#include <ovito/mesh/surface/SurfaceMesh.h>
#include <ovito/mesh/surface/SurfaceMeshVis.h>
#include <ovito/stdobj/simcell/SimulationCell.h>
#include <ovito/particles/objects/Particles.h>
#include <ovito/particles/objects/ParticleBondMap.h>
#include <ovito/core/dataset/DataSet.h>
#include <ovito/core/dataset/pipeline/ModificationNode.h>
#include "CoordinationPolyhedraModifier.h"

namespace Ovito {

IMPLEMENT_CREATABLE_OVITO_CLASS(CoordinationPolyhedraModifier);
OVITO_CLASSINFO(CoordinationPolyhedraModifier, "DisplayName", "Coordination polyhedra");
OVITO_CLASSINFO(CoordinationPolyhedraModifier, "Description", "Visualize atomic coordination polyhedra.");
OVITO_CLASSINFO(CoordinationPolyhedraModifier, "ModifierCategory", "Visualization");
DEFINE_REFERENCE_FIELD(CoordinationPolyhedraModifier, surfaceMeshVis);
DEFINE_PROPERTY_FIELD(CoordinationPolyhedraModifier, transferParticleProperties);
SET_PROPERTY_FIELD_LABEL(CoordinationPolyhedraModifier, transferParticleProperties, "Transfer particle properties to mesh");

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

    // Create the vis element for rendering the polyhedra generated by the modifier.
    if(!flags.testFlag(ObjectInitializationFlag::DontInitializeObject)) {
        setSurfaceMeshVis(OORef<SurfaceMeshVis>::create(flags));
        surfaceMeshVis()->setShowCap(false);
        surfaceMeshVis()->setSmoothShading(false);
        surfaceMeshVis()->setSurfaceTransparency(FloatType(0.25));
        surfaceMeshVis()->setObjectTitle(tr("Polyhedra"));
        if(this_task::isInteractive())
            surfaceMeshVis()->setHighlightEdges(true);
    }
}

/******************************************************************************
* Replaces any references the modifier has to the given visual element with a new compatible object.
******************************************************************************/
void CoordinationPolyhedraModifier::replaceVisualElement(DataVis* visElement, const std::function<OORef<DataVis>(const QString&)>& getReplacement)
{
    if(surfaceMeshVis() == visElement)
        setSurfaceMeshVis(static_object_cast<SurfaceMeshVis>(getReplacement(tr("Polyhedra"))));
}

/******************************************************************************
* Asks the modifier whether it can be applied to the given input data.
******************************************************************************/
bool CoordinationPolyhedraModifier::OOMetaClass::isApplicableTo(const DataCollection& input) const
{
    if(const Particles* particles = input.getObject<Particles>()) {
        return particles->bonds() != nullptr;
    }
    return false;
}

/******************************************************************************
* Modifies the input data.
******************************************************************************/
Future<PipelineFlowState> CoordinationPolyhedraModifier::evaluateModifier(const ModifierEvaluationRequest& request, PipelineFlowState&& state)
{
    // Get modifier input.
    const Particles* particles = state.expectObject<Particles>();
    particles->verifyIntegrity();
    const Property* positions = particles->expectProperty(Particles::PositionProperty);
    const Property* selection = particles->getProperty(Particles::SelectionProperty);

    particles->expectBonds()->verifyIntegrity();
    const Property* bondTopology = particles->expectBondsTopology();
    const Property* bondPeriodicImages = particles->bonds()->getProperty(Bonds::PeriodicImageProperty);
    const SimulationCell* simCell = state.getObject<SimulationCell>();

    if(!selection)
        throw Exception(tr("Please first select some particles, for which coordination polyhedra should be generated."));

    // Collect the set of particle properties that should be transferred over to the surface mesh vertices and mesh regions.
    std::vector<ConstPropertyPtr> particleProperties;
    if(transferParticleProperties()) {
        for(const Property* property : particles->properties()) {
            // Certain properties should never be transferred to the mesh vertices.
            if(property->typeId() == Particles::SelectionProperty) continue;
            if(property->typeId() == Particles::PositionProperty) continue;
            if(property->typeId() == Particles::ColorProperty) continue;
            if(property->typeId() == Particles::VectorColorProperty) continue;
            if(property->typeId() == Particles::PeriodicImageProperty) continue;
            if(property->typeId() == Particles::TransparencyProperty) continue;
            particleProperties.push_back(property);
        }
    }

    // Create the output surface mesh.
    SurfaceMesh* mesh = state.createObjectWithVis<SurfaceMesh>(QStringLiteral("coord-polyhedra"), request.modificationNode(), surfaceMeshVis(), tr("Coordination polyhedra"));
    mesh->setDomain(simCell);

    // Perform the main work in a separate thread.
    return asyncLaunch([
            state = std::move(state),
            positions,
            selection,
            simCell,
            bondTopology,
            bondPeriodicImages,
            mesh = std::move(mesh),
            particleProperties = std::move(particleProperties)]() mutable
    {
        TaskProgress progress(this_task::ui());
        progress.setText(tr("Generating coordination polyhedra"));

        SurfaceMeshBuilder meshBuilder(mesh);

        // Create the "Region" face property.
        meshBuilder.createFaceProperty(DataBuffer::Uninitialized, SurfaceMeshFaces::RegionProperty);
        PropertyContainer::Grower regionGrower(meshBuilder.mutableRegions());

        // Determine number of selected particles.
        size_t npoly = selection->nonzeroCount();
        progress.setMaximum(npoly);

        ParticleBondMap bondMap(bondTopology, bondPeriodicImages);

        BufferReadAccess<Point3> positionsArray(positions);

        // Working variables.
        std::vector<Point3> neighborPositions;
        std::vector<size_t> neighborIndices;
        SurfaceMesh::size_type oldVertexCount = 0;

        // After construction of the mesh, this array will contain for each
        // mesh vertex the index of the particle it was created from.
        std::vector<size_t> vertexToParticleMap;
        // After construction of the mesh, this array will contain for each
        // mesh region the index of the particle it was created for.
        std::vector<size_t> regionToParticleMap;
        regionToParticleMap.reserve(npoly);

        // Iterate over all input particles.
        BufferReadAccess<SelectionIntType> selectionAcc(selection);
        for(size_t i = 0; i < positionsArray.size(); i++) {
            // Construct coordination polyhedron only for selected particles.
            if(selectionAcc[i] == 0) continue;

            // Collect the bonds that are part of the coordination polyhedron.
            const Point3& p1 = positionsArray[i];
            for(BondWithIndex bond : bondMap.bondsOfParticle(i)) {
                if(bond.index2 < positions->size()) {
                    Vector3 delta = positionsArray[bond.index2] - p1;
                    if(simCell) {
                        if(bond.pbcShift.x()) delta += simCell->matrix().column(0) * (FloatType)bond.pbcShift.x();
                        if(bond.pbcShift.y()) delta += simCell->matrix().column(1) * (FloatType)bond.pbcShift.y();
                        if(bond.pbcShift.z()) delta += simCell->matrix().column(2) * (FloatType)bond.pbcShift.z();
                    }
                    neighborPositions.push_back(p1 + delta);
                    neighborIndices.push_back(bond.index2);
                }
            }

            // Include the central particle in the point list too.
            neighborPositions.push_back(p1);
            neighborIndices.push_back(i);
            regionToParticleMap.push_back(i);

            // Construct the polyhedron (i.e. convex hull) from the point list.
            if(particleProperties.empty()) {
                // Note: We are moving the point list into the convex hull method to avoid an extra copy.
                meshBuilder.constructConvexHull(std::move(neighborPositions), regionGrower.grow(1));
            }
            else {
                // Note: We keep our own copy of the point list so that we can determine the insertion order afterwards.
                meshBuilder.constructConvexHull(neighborPositions, regionGrower.grow(1));

                // Find each input point among the newly added vertices of the mesh.
                // This will help us later to transfer the particle properties to the corresponding mesh vertices.
                BufferReadAccess<Point3> vertexPositions = meshBuilder.expectVertexProperty(SurfaceMeshVertices::PositionProperty);
#ifndef OVITO_USE_SYCL
                for(const Point3& vpos : std::move(vertexPositions).subrange(oldVertexCount)) {
                    auto idx = neighborIndices.cbegin();
                    for(const Point3& p : neighborPositions) {
                        OVITO_ASSERT(idx != neighborIndices.cend());
                        if(vpos == p) {
                            vertexToParticleMap.push_back(*idx);
                            break;
                        }
                        ++idx;
                    }
                }
#else
                OVITO_ASSERT(false);  // This code path is not supported in the SYCL version of OVITO.
#endif
                OVITO_ASSERT(vertexToParticleMap.size() == meshBuilder.vertexCount());
                oldVertexCount = meshBuilder.vertexCount();
            }

            // Clear point list for next loop iteration.
            neighborPositions.clear();
            neighborIndices.clear();

            progress.incrementValue();
        }
        regionGrower.commit();
        OVITO_ASSERT(regionToParticleMap.size() == meshBuilder.regionCount());

        // Transfer particle properties to the mesh vertices and mesh regions if requested.
        if(!particleProperties.empty()) {
            OVITO_ASSERT(vertexToParticleMap.size() == meshBuilder.vertexCount());
            for(const ConstPropertyPtr& particleProperty : particleProperties) {

                // Create the corresponding output mesh vertex property.
                PropertyPtr vertexProperty;
                if(particleProperty->typeId() < Property::FirstSpecificProperty && SurfaceMeshVertices::OOClass().isValidStandardPropertyId(particleProperty->typeId())) {
                    // Input property is also a standard property for mesh vertices.
                    vertexProperty = meshBuilder.createVertexProperty(DataBuffer::Uninitialized, static_cast<SurfaceMeshVertices::Type>(particleProperty->typeId()));
                    OVITO_ASSERT(vertexProperty->dataType() == particleProperty->dataType());
                    OVITO_ASSERT(vertexProperty->stride() == particleProperty->stride());
                }
                else if(SurfaceMeshVertices::OOClass().standardPropertyTypeId(particleProperty->name()) != 0) {
                    // Input property name is that of a standard property for mesh vertices.
                    // Must rename the property to avoid conflict, because user properties may not have a standard property name.
                    QString newPropertyName = particleProperty->name() + tr("_particles");
                    vertexProperty = meshBuilder.createVertexProperty(DataBuffer::Uninitialized, newPropertyName, particleProperty->dataType(), particleProperty->componentCount(), particleProperty->componentNames());
                }
                else {
                    // Input property is a user property for mesh vertices.
                    vertexProperty = meshBuilder.createVertexProperty(DataBuffer::Uninitialized, particleProperty->name(), particleProperty->dataType(), particleProperty->componentCount(), particleProperty->componentNames());
                }
                // Copy particle property values to mesh vertices using precomputed index mapping.
                particleProperty->mappedCopyTo(*vertexProperty, vertexToParticleMap);
                // Also adapt element types of the property.
                vertexProperty->setElementTypes(particleProperty->elementTypes());

                // Create the corresponding output mesh region property.
                PropertyPtr regionProperty;
                if(particleProperty->typeId() < Property::FirstSpecificProperty && SurfaceMeshRegions::OOClass().isValidStandardPropertyId(particleProperty->typeId())) {
                    // Input property is also a standard property for mesh regions.
                    regionProperty = meshBuilder.createRegionProperty(DataBuffer::Uninitialized, static_cast<SurfaceMeshRegions::Type>(particleProperty->typeId()));
                    OVITO_ASSERT(regionProperty->dataType() == particleProperty->dataType());
                    OVITO_ASSERT(regionProperty->stride() == particleProperty->stride());
                }
                else if(SurfaceMeshRegions::OOClass().standardPropertyTypeId(particleProperty->name()) != 0) {
                    // Input property name is that of a standard property for mesh regions.
                    // Must rename the property to avoid conflict, because user properties may not have a standard property name.
                    QString newPropertyName = particleProperty->name() + tr("_particles");
                    regionProperty = meshBuilder.createRegionProperty(DataBuffer::Uninitialized, newPropertyName, particleProperty->dataType(), particleProperty->componentCount(), particleProperty->componentNames());
                }
                else {
                    // Input property is a user property for mesh regions.
                    regionProperty = meshBuilder.createRegionProperty(DataBuffer::Uninitialized, particleProperty->name(), particleProperty->dataType(), particleProperty->componentCount(), particleProperty->componentNames());
                }
                // Copy particle property values to mesh regions using precomputed index mapping.
                particleProperty->mappedCopyTo(*regionProperty, regionToParticleMap);
                // Also adapt element types of the property.
                regionProperty->setElementTypes(particleProperty->elementTypes());
            }
        }

        // Create the "Particle index" region property, which contains the index of the particle that is at the center of each coordination polyhedron.
        PropertyPtr particleIndexProperty = meshBuilder.createRegionProperty(DataBuffer::Uninitialized, QStringLiteral("Particle Index"), Property::Int64);
        std::copy(regionToParticleMap.cbegin(), regionToParticleMap.cend(), BufferWriteAccess<int64_t, access_mode::discard_write>(particleIndexProperty).begin());

        return std::move(state);
    });
}

}   // End of namespace
