////////////////////////////////////////////////////////////////////////////////////////
//
//  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/particles/objects/Particles.h>
#include <ovito/stdobj/simcell/SimulationCell.h>
#include <ovito/core/app/Application.h>
#include "POSCARExporter.h"

namespace Ovito {

IMPLEMENT_CREATABLE_OVITO_CLASS(POSCARExporter);
DEFINE_PROPERTY_FIELD(POSCARExporter, writeReducedCoordinates);
SET_PROPERTY_FIELD_LABEL(POSCARExporter, writeReducedCoordinates, "Output reduced coordinates");

/******************************************************************************
* Creates a worker performing the actual data export.
*****************************************************************************/
OORef<FileExportJob> POSCARExporter::createExportJob(const QString& filePath, int numberOfFrames)
{
    class Job : public FileExportJob
    {
    public:

        /// Writes the exportable data of a single trajectory frame to the output file.
        virtual SCFuture<void> exportFrameData(boost::anys::unique_any frameData, int frameNumber, const QString& filePath,
                                               TaskProgress& progress) override
        {
            // The exportable frame data.
            OVITO_ASSERT(frameData.has_value());
            const auto state = boost::anys::any_cast<PipelineFlowState>(frameData);

            // Perform the following in a worker thread.
            co_await ExecutorAwaiter(ThreadPoolExecutor());

            // Get particle positions and velocities.
            const Particles* particles = state.expectObject<Particles>();
            BufferReadAccess<Point3> posProperty = particles->expectProperty(Particles::PositionProperty);
            BufferReadAccess<Vector3> velocityProperty = particles->getProperty(Particles::VelocityProperty);
            size_t particleCount = particles->elementCount();

            // Get simulation cell info.
            const SimulationCell* simulationCell = state.getObject<SimulationCell>();
            if(!simulationCell)
                throw Exception(tr("No simulation cell available. Cannot write POSCAR file."));

            // Write POSCAR header including the simulation cell geometry.
            textStream() << "POSCAR file written by " << Application::applicationName() << " " << Application::applicationVersionString() << "\n";
            textStream() << "1\n";
            for(size_t i = 0; i < 3; i++)
                textStream() << simulationCell->matrix()(0, i) << ' ' << simulationCell->matrix()(1, i) << ' ' << simulationCell->matrix()(2, i) << '\n';
            const Vector3& origin = simulationCell->matrix().translation();

            // Count number of particles per particle type.
            QMap<int,int> particleCounts;
            const Property* particleTypeProperty = particles->getProperty(Particles::TypeProperty);
            BufferReadAccess<int32_t> particleTypeArray(particleTypeProperty);
            if(particleTypeProperty) {
                for(int ptype : particleTypeArray)
                    particleCounts[ptype]++;

                // Write line with particle type names.
                for(auto c = particleCounts.begin(); c != particleCounts.end(); ++c) {
                    const ElementType* particleType = particleTypeProperty->elementType(c.key());
                    if(particleType) {
                        QString typeName = particleType->nameOrNumericId();
                        typeName.replace(' ', '_');
                        textStream() << typeName << ' ';
                    }
                    else textStream() << "Type" << c.key() << ' ';
                }
                textStream() << '\n';

                // Write line with particle counts per type.
                for(auto c = particleCounts.begin(); c != particleCounts.end(); ++c) {
                    textStream() << c.value() << ' ';
                }
                textStream() << '\n';
            }
            else {
                // Write line with particle type name.
                textStream() << "A\n";
                // Write line with particle count.
                textStream() << particleCount << '\n';
                particleCounts[0] = particleCount;
            }

            qlonglong totalProgressCount = particleCount;
            if(velocityProperty)
                totalProgressCount += particleCount;
            qlonglong currentProgress = 0;
            progress.setMaximum(totalProgressCount);

            bool writeReducedCoordinates = static_cast<const POSCARExporter*>(this->exporter())->writeReducedCoordinates();

            // Write atomic positions.
            textStream() << (writeReducedCoordinates ? "Direct\n" : "Cartesian\n");
            for(auto c = particleCounts.begin(); c != particleCounts.end(); ++c) {
                int ptype = c.key();
                const Point3* p = posProperty.cbegin();
                for(size_t i = 0; i < particleCount; i++, ++p) {
                    if(particleTypeArray && particleTypeArray[i] != ptype)
                        continue;
                    if(writeReducedCoordinates) {
                        Point3 rp = simulationCell->absoluteToReduced(*p);
                        textStream() << rp.x() << ' ' << rp.y() << ' ' << rp.z() << '\n';
                    }
                    else {
                        textStream() << (p->x() - origin.x()) << ' ' << (p->y() - origin.y()) << ' ' << (p->z() - origin.z()) << '\n';
                    }

                    // Update progress bar and check for user cancellation.
                    progress.setValueIntermittent(currentProgress++);
                }
            }

            // Write atomic velocities.
            if(velocityProperty) {
                textStream() << (writeReducedCoordinates ? "Direct\n" : "Cartesian\n");
                for(auto c = particleCounts.begin(); c != particleCounts.end(); ++c) {
                    int ptype = c.key();
                    const Vector3* v = velocityProperty.cbegin();
                    for(size_t i = 0; i < particleCount; i++, ++v) {
                        if(particleTypeArray && particleTypeArray[i] != ptype)
                            continue;

                        if(writeReducedCoordinates) {
                            Vector3 rv = simulationCell->absoluteToReduced(*v);
                            textStream() << rv.x() << ' ' << rv.y() << ' ' << rv.z() << '\n';
                        }
                        else {
                            textStream() << v->x() << ' ' << v->y() << ' ' << v->z() << '\n';
                        }

                        // Update progress bar and check for user cancellation.
                        progress.setValueIntermittent(currentProgress++);
                    }
                }
            }
        }
    };

    return OORef<Job>::create(this, filePath, true);
}

}   // End of namespace
