io.h

#pragma once

#include "particle.h"
#include "units.h"
#include "group.h"
#include <fstream>
#include <type_traits>
#include <concepts>
#include <spdlog/spdlog.h>
#include <numeric>

namespace Faunus {

class Space;
class Group;

[[maybe_unused]] std::string addGrowingSuffix(
    const std::string&); 

#ifndef __cplusplus
#define __cplusplus
#endif
#ifndef CPLUSPLUS
#define CPLUSPLUS
#endif

namespace XdrFile {
#include <xdrfile_trr.h>
#include <xdrfile_xtc.h>

struct XdrFileDeleter
{
    void
    operator()(XdrFile::XDRFILE* xdrfile); 
};

using XDRFILE_unique =
    std::unique_ptr<XDRFILE, XdrFileDeleter>; 
} // namespace XdrFile

namespace IO {

std::unique_ptr<std::ostream> openCompressedOutputStream(const std::string& filename,
                                                         bool throw_on_error = false);

template <typename TKey, typename TValue>
void writeKeyValuePairs(std::ostream& stream, const std::map<TKey, TValue>& data,
                        const std::string& sep = " ", const std::string& end = "\n")
{
    if (stream) {
        for (const auto& [key, value] : data) {
            stream << key << sep << value << end;
        }
    }
}

template <typename TKey, typename TValue>
void writeKeyValuePairs(const std::string& filename, const std::map<TKey, TValue>& data)
{
    if (!data.empty()) {
        std::ofstream file(filename);
        writeKeyValuePairs(file, data);
    }
}

} // namespace IO

class StructureFileReader
{
  private:
    virtual void loadHeader(std::istream& stream) = 0;       
    virtual void loadFooter(std::istream& stream);           
    virtual Particle loadParticle(std::istream& stream) = 0; 
    void checkLoadedParticles() const; 

  protected:
    static size_t getNumberOfAtoms(const std::string& line); 
    void getNextLine(std::istream& stream, std::string& destination,
                     const std::string& comment_identifiers); 
    size_t expected_number_of_particles = 0;

    void handleChargeMismatch(Particle& particle,
                              int atom_index) const; 
    static void handleRadiusMismatch(const Particle& particle, double radius,
                                     int atom_index); 

  public:
    ParticleVector particles;
    Point box_length = Point::Zero();
    std::vector<std::string> comments;
    bool prefer_charges_from_file =
        true; 

    ParticleVector& load(std::istream& stream); 
    ParticleVector& load(std::string_view filename);
    virtual ~StructureFileReader() = default;

    bool box_dimension_support = false;
    bool particle_charge_support = false;
    bool particle_radius_support = false;
};

class CoarseGrainedFastaFileReader : public StructureFileReader
{
  private:
    void loadHeader(std::istream& stream) override;
    Particle loadParticle(std::istream& stream) override;
    double bond_length = 0.0;
    Point new_particle_position = Point::Zero();

  public:
    explicit CoarseGrainedFastaFileReader(double bond_length,
                                          const Point& initial_particle_position = Point(0, 0, 0));
    void setBondLength(double bond_length);
    std::string loadSequence(std::istream& stream);
    static char getFastaLetter(std::istream& stream);
};

class AminoAcidModelReader : public StructureFileReader
{
  private:
    void loadHeader(std::istream& stream) override;
    Particle loadParticle(std::istream& stream) override;

  public:
    AminoAcidModelReader();
};

class PQRReader : public StructureFileReader
{
  private:
    void loadHeader(std::istream& stream) override;
    Particle loadParticle(std::istream& stream) override;

  public:
    PQRReader();
};

class XYZReader : public StructureFileReader
{
  private:
    void loadHeader(std::istream& stream) override;
    Particle loadParticle(std::istream& stream) override;
};

class SpheroCylinderXYZReader : public StructureFileReader
{
  private:
    void loadHeader(std::istream& stream) override;
    Particle loadParticle(std::istream& stream) override;
};

class GromacsReader : public StructureFileReader
{
  private:
    void loadBoxInformation(
        std::istream& stream); 
    void loadHeader(std::istream& stream) override;
    Particle loadParticle(std::istream& stream) override;

  public:
    GromacsReader();
};

class StructureFileWriter
{
  private:
    virtual void saveHeader(std::ostream& stream,
                            int number_of_particles) const = 0; 
    virtual void
    saveFooter(std::ostream& stream) const; 
    virtual void saveParticle(std::ostream& stream,
                              const Particle& particle) = 0;  
    void saveGroup(std::ostream& stream, const Group& group); 

    template <class ParticleIter>
    void saveParticles(std::ostream& stream, ParticleIter begin, ParticleIter end)
    {
        group_index = 0;
        particle_index = 0;
        std::for_each(begin, end, [&](const auto& particle) {
            saveParticle(stream, particle);
            particle_index++;
        });
    }

  protected:
    static const std::string generated_by_faunus_comment;
    bool particle_is_active = true;
    std::string group_name;
    std::size_t particle_index = 0;
    std::size_t group_index = 0;
    Point box_dimensions = Point::Zero();

  public:
    template <std::forward_iterator ParticleIter>
    void save(std::ostream& stream, ParticleIter begin, ParticleIter end, const Point& box_length)
    {
        if (auto number_of_particles = std::distance(begin, end); number_of_particles > 0) {
            box_dimensions = box_length;
            saveHeader(stream, number_of_particles);
            saveParticles(stream, begin, end);
            saveFooter(stream);
        }
    }

    void save(std::ostream& stream, const RequireGroups auto& groups, const Point& box_length)
    {
        group_index = 0;
        particle_index = 0;
        box_dimensions = box_length;

        auto group_sizes = groups | std::views::transform(&Group::capacity);
        const auto number_of_particles =
            std::accumulate(group_sizes.begin(), group_sizes.end(), 0u);
        saveHeader(stream, number_of_particles);

        std::ranges::for_each(groups, [&](const auto& group) { saveGroup(stream, group); });
        saveFooter(stream);
    }

    template <class... Args> void save(const std::string& filename, const Args&... args)
    {
        if (std::ofstream stream(filename); stream) {
            faunus_logger->debug("writing to {}", filename);
            save(stream, args...);
        }
        else {
            throw std::runtime_error("writeKeyValuePairs error: "s + filename);
        }
    }

    virtual ~StructureFileWriter() = default;
};

class AminoAcidModelWriter : public StructureFileWriter
{
  private:
    void saveHeader(std::ostream& stream, int number_of_particles) const override;
    void saveParticle(std::ostream& stream, const Particle& particle) override;
};

class XYZWriter : public StructureFileWriter
{
  private:
    void saveHeader(std::ostream& stream, int number_of_particles) const override;
    void saveParticle(std::ostream& stream, const Particle& particle) override;
};

class SpheroCylinderXYZWriter : public StructureFileWriter
{
  private:
    void saveHeader(std::ostream& stream, int number_of_particles) const override;
    void saveParticle(std::ostream& stream, const Particle& particle) override;
};

class PQRWriter : public StructureFileWriter
{
  private:
    void saveHeader(std::ostream& stream, int number_of_particles) const override;
    void saveFooter(std::ostream& stream) const override;
    void saveParticle(std::ostream& stream, const Particle& particle) override;

  public:
    enum class Style
    {
        PQR_LEGACY,
        PDB,
        PQR
    }; 
    Style style = Style::PQR_LEGACY;
    explicit PQRWriter(Style style = Style::PQR_LEGACY);
};

class GromacsWriter : public StructureFileWriter
{
  private:
    void saveHeader(std::ostream& stream, int number_of_particles) const override;
    void saveFooter(std::ostream& stream) const override;
    void saveParticle(std::ostream& stream, const Particle& particle) override;
};

namespace PQRTrajectoryReader {
bool readAtomRecord(const std::string& record, Particle& particle,
                    double& radius); 
void loadTrajectory(const std::string& filename,
                    std::vector<ParticleVector>& destination); 
} // namespace PQRTrajectoryReader

struct TrajectoryFrame;

struct XTCTrajectoryFrame
{
    XdrFile::matrix xtc_box; 
    std::unique_ptr<XdrFile::rvec[]> xtc_coordinates; 
    int xtc_step = 0;                                 
    float xtc_time = 0.0;                             
    int number_of_atoms = 0;  
    float precision = 1000.0; 

    XTCTrajectoryFrame(int number_of_atoms);
    XTCTrajectoryFrame(const TrajectoryFrame& frame);

    template <RequirePointIterator begin_iterator, class end_iterator>
    XTCTrajectoryFrame(int step, float time, const Point& box, begin_iterator coordinates_begin,
                       end_iterator coordinates_end)
    {
        initNumberOfAtoms(std::distance(coordinates_begin, coordinates_end));
        importFrame(step, time, box, coordinates_begin, coordinates_end);
    }

    XTCTrajectoryFrame& operator=(const TrajectoryFrame& frame);
    void importFrame(const TrajectoryFrame& frame);

    template <RequirePointIterator begin_iterator, class end_iterator>
    void importFrame(const int step, const float time, const Point& box,
                     begin_iterator coordinates_begin, end_iterator coordinates_end)
    {
        importTimestamp(step, time);
        importBox(box);
        importCoordinates(coordinates_begin, coordinates_end, 0.5 * box);
    }

    void exportFrame(TrajectoryFrame& frame) const;

    template <RequirePointIterator begin_iterator, RequirePointIterator end_iterator>
    void exportFrame(int& step, float& time, Point& box, begin_iterator coordinates_begin,
                     end_iterator coordinates_end) const
    {
        exportTimestamp(step, time);
        exportBox(box);
        exportCoordinates(coordinates_begin, coordinates_end, 0.5 * box);
    }

  protected:
    using XTCFloat = float;
    using XTCMatrix =
        Eigen::Matrix<XTCFloat, 3, 3, Eigen::RowMajor>; //<! eigen equivalent of the box tensor
    using XTCVector =
        Eigen::Matrix<XTCFloat, 3, 1>; //<! eigen equivalent of the single set of coordinates
    void importTimestamp(int step, float time);
    void importBox(const Point& box);
    void importCoordinates(const PointVector& coordinates, const Point& offset);

    template <RequirePointIterator begin_iterator, typename end_iterator>
    void importCoordinates(begin_iterator begin, end_iterator end, const Point& offset) const
    {
        int i = 0;
        std::ranges::for_each(begin, end, [&](const auto& position) {
            if (i >= number_of_atoms) {
                throw std::runtime_error("too many particles for XTC frame");
            }
            const XTCVector xtc_pos = ((position + offset) / 1.0_nm).template cast<XTCFloat>();
            std::copy(xtc_pos.data(), xtc_pos.data() + DIM, xtc_coordinates.get()[i++]);
        });
        if (i != number_of_atoms) {
            throw std::runtime_error("too few particles for XTC frame");
        }
    }

    void exportTimestamp(int& step, float& time) const;
    void exportBox(Point& box) const;
    void exportCoordinates(PointVector& coordinates, const Point& offset) const;

    template <RequirePointIterator begin_iterator, typename end_iterator>
    void exportCoordinates(begin_iterator begin, end_iterator end, const Point& offset) const
    {
        // comparison of i is probably faster than prior call of std::distance
        int i = 0;
        for (auto coordinates_it = begin; coordinates_it != end; ++coordinates_it) {
            if (i >= number_of_atoms) {
                throw std::runtime_error("number of particles in the loaded XTC frame too low");
            }
            XTCVector xtc_atom_coordinates(xtc_coordinates.get()[i++]);
            *coordinates_it = Point(xtc_atom_coordinates.cast<double>() * 1.0_nm) - offset;
        }
        if (i < number_of_atoms) {
            throw std::runtime_error("number of particles in the loaded XTC frame too high");
        }
    }

  private:
    void initNumberOfAtoms(int);
};

struct TrajectoryFrame
{
    Point box;               
    PointVector coordinates; 
    int step = 0;            
    float timestamp = 0.0;   

    TrajectoryFrame() = default;
    TrajectoryFrame(const Point& box, const PointVector& coordinates, int step, float timestamp);
    explicit TrajectoryFrame(const XTCTrajectoryFrame& xtc_frame);
    void operator=(const XTCTrajectoryFrame& xtc_frame);
};

class XTCReader
{
    int return_code = XdrFile::exdrOK; 
    XdrFile::XDRFILE_unique xdrfile;   
    std::unique_ptr<XTCTrajectoryFrame> xtc_frame;

  public:
    std::string filename; 

    explicit XTCReader(const std::string& filename);

    int getNumberOfCoordinates();
    bool read(TrajectoryFrame& frame);

    template <RequirePointIterator begin_iterator, typename end_iterator>
    bool read(int& step, float& time, Point& box, begin_iterator coordinates_begin,
              end_iterator coordinates_end)
    {
        if (readFrame()) {
            xtc_frame->exportFrame(step, time, box, coordinates_begin, coordinates_end);
            return true;
        }
        return false;
    }

  protected:
    bool readFrame();
};

class XTCWriter
{
    int return_code = XdrFile::exdrOK;             
    XdrFile::XDRFILE_unique xdrfile;               
    std::unique_ptr<XTCTrajectoryFrame> xtc_frame; 
    int step_counter = 0;                          
    float time_delta = 1.0_ps; 

  public:
    const std::string filename; 

    explicit XTCWriter(const std::string& filename);
    void write(const TrajectoryFrame& frame);
    void writeNext(const TrajectoryFrame& frame);

    template <RequirePointIterator begin_iterator, typename end_iterator>
    void writeNext(const Point& box, begin_iterator coordinates_begin, end_iterator coordinates_end)
    {
        if (!xtc_frame) {
            auto number_of_atoms = std::ranges::distance(coordinates_begin, coordinates_end);
            faunus_logger->trace("preparing xtc output for {} particles", number_of_atoms);
            xtc_frame = std::make_unique<XTCTrajectoryFrame>(number_of_atoms);
        }
        xtc_frame->importFrame(step_counter, static_cast<float>(step_counter) * time_delta, box,
                               coordinates_begin, coordinates_end);
        writeFrame();
        ++step_counter;
    }

  protected:
    void writeFrame();
    void writeFrameAt(int step, float time);
};

std::vector<AtomData::index_type>
fastaToAtomIds(std::string_view fasta_sequence); 

ParticleVector fastaToParticles(std::string_view fasta_sequence, double bond_length = 7.0,
                                const Point& origin = {0, 0, 0});

ParticleVector loadStructure(std::string_view filename, bool prefer_charges_from_file = true);

std::unique_ptr<StructureFileWriter> createStructureFileWriter(const std::string& suffix);

} // namespace Faunus