io.hpp

#pragma once

#include <iostream>
#include <filesystem>
#include <fstream>
#include <stdint.h>
#include <limits.h>
#include <unordered_map>

// Only need to include MPI headers with the MPI supporting implementations (> CSTONE)
#if defined WASH_CSTONE || defined WASH_WONE  
#include <mpi.h>
#endif

#include "wash.hpp"
#include "vector.hpp"
#include "particle_data.hpp"

#if SIZE_MAX == UCHAR_MAX
   #define MPI_SIZE_T MPI_UNSIGNED_CHAR
#elif SIZE_MAX == USHRT_MAX
   #define MPI_SIZE_T MPI_UNSIGNED_SHORT
#elif SIZE_MAX == UINT_MAX
   #define MPI_SIZE_T MPI_UNSIGNED
#elif SIZE_MAX == ULONG_MAX
   #define MPI_SIZE_T MPI_UNSIGNED_LONG
#elif SIZE_MAX == ULLONG_MAX
   #define MPI_SIZE_T MPI_UNSIGNED_LONG_LONG
#else
   #error "Unknown size_t size"
#endif

namespace fs = std::filesystem;

namespace wash {

namespace io {

    const std::string get_simulation_name();

    const std::string get_output_name();

    struct SimulationData {
        // Number of particles (rows) in the data 
        size_t particle_count;
        // `particle_count * sum(dim[labels[i]])` doubles representing particle data
        // each row is a particle, columns are the property/force data
        std::vector<double> data;
        // labels for each particle property/force. Should contain global IDs and all predeined properties
        std::vector<std::string> labels;
        // dimension for each label
        std::vector<unsigned short> dim;

        friend std::ostream& operator<<(std::ostream& os, const SimulationData& data) {
            os << "SimulationData[ count=" << data.particle_count << ", data=" << data.data.size() << ", ";
            os << " labels=";
            for (int i = 0; i < data.labels.size(); i++) {
                os << "(" << data.labels.at(i) << "," << data.dim.at(i) << "),";
            }
            os << "]";

            return os;
        }
    };

    // Return a SimulationData struct holding a copy of all data at that point. 
    SimulationData copy_simulation_data();

    class IOManager {
    public:
        using WriterFuncT = std::function<int(const io::IOManager&, const SimulationData&, const size_t)>;
        static const std::unordered_map<std::string, std::string> label_map;
    private:
        std::string path; 
        WriterFuncT writer;
        size_t output_nth;

        size_t rank;
        size_t size;
        bool gather;

        bool do_timings;

        SimulationData data;
    public:
    
        IOManager(const std::string format, WriterFuncT writer, const size_t nth, const size_t rank = 0, const size_t size = 1, const bool timings = true) : writer(writer), output_nth(nth), rank(rank), size(size), do_timings(timings) {
            path = "./out/" + get_simulation_name() + std::string("/");
            std::cout << "IO Manager: Output Path: " << path << "; Type: " << format << "; Rank " << rank << "; of " << size << ";" << std::endl;

            if (!fs::exists(path)) {
                try {
                    fs::create_directories(path);
                } catch (const std::exception& e) {
                    std::cerr << "Error creating directory: " << e.what() << std::endl;
                    throw std::runtime_error("Error initialising IO manager: Creating Output directory");
                }
            }

            if (do_timings) {
                new_timings();
            }
        }

        IOManager(const std::string format, WriterFuncT writer) : IOManager(format, writer, 1) {}

        IOManager(const std::string format, WriterFuncT writer, const size_t rank, const size_t size) : IOManager(format, writer, 1, rank, size) {} 

        IOManager();

        const std::string get_path() const {
            return this->path;
        }

        size_t get_rank() const {
            return this->rank;
        }

        size_t get_size() const {
            return this->size;
        }

        void set_gather(bool value = true) {
            this->gather = value;
        }

        const std::string& expand_label(const std::string& label) const {
            auto find = label_map.find(label);
            if (find != label_map.end()) {
                return find->second;
            }

            return label;
        }

        const SimulationData get_simulation_data() {
            data = copy_simulation_data();
#if defined WASH_WSER || defined WASH_WISB || defined WASH_WEST
            // Just return the copied data in non-MPI capable implementations
            return data;
#elif defined WASH_CSTONE || defined WASH_WONE 
            // If we're not running in more than one rank, or not using gather then just return
            if (size == 1 || !gather) {
                return data;
            // Else we need to gather from all ranks to write the data out from the 1st rank
            } else {
                std::vector<int> recv_counts(size, 1);
                std::vector<int> displs (size);
                for (size_t i = 0 ; i < size; i++) { // Displace each element `i` from start
                    displs[i] = i;
                }

                size_t particle_counts[size];
                MPI_Gatherv(&data.particle_count, 1, MPI_SIZE_T, particle_counts, recv_counts.data(), displs.data(), MPI_SIZE_T, 0, MPI_COMM_WORLD);

                // Have to downcast here for MPI - only accepts an int. 
                int int_particle_counts[size];
                unsigned sim_particle_count = 0;
                for (size_t idx = 0; idx < size; idx++) {
                    int_particle_counts[idx] = particle_counts[idx]; 
                    sim_particle_count += particle_counts[idx];
                }

                int total_width = 0;
                for (size_t idx = 0; idx < data.labels.size(); idx++) {
                    total_width += data.dim[idx];
                } 

                // Row = particle_data, cols = force/property in labels order
                std::vector<double> sim_data(sim_particle_count * total_width);
                
                std::vector<int> send_sizes(size); 
                for (size_t i = 0; i < size; i++) {
                    send_sizes[i] = total_width * int_particle_counts[i];
                }
                
                std::vector<int> displs_2(size); 
                displs_2[0] = 0;
                for (size_t i = 1; i < size; i++) {
                    displs_2[i] = displs_2[i-1] + send_sizes[i];
                }

                MPI_Gatherv(data.data.data(), data.data.size(), MPI_DOUBLE, sim_data.data(), send_sizes.data(), displs_2.data(), MPI_DOUBLE, 0, MPI_COMM_WORLD );
                
                if (rank == 0) {
                    return SimulationData { .particle_count = sim_particle_count, .data = sim_data, .labels = data.labels, .dim = data.dim };
                } else {
                    return data;
                }
            }
#endif
        }

        void write_iteration(const size_t iteration) {
            if (writer && iteration % this->output_nth == 0 && this->path != "") {
                auto sim_data = get_simulation_data();
                if (gather && rank != 0) { // Don't write if we're using the gather and not rnk 0 
                    return;
                }

                writer(*this, sim_data, iteration);
            }
        }

        void write_timings(const std::string& event_name, const int tag, const int64_t time_taken) const {
            if (do_timings) {
                std::string fpath = (new std::string(this->path))->append(get_output_name() + ".timings.csv");

                std::ios_base::openmode mode = std::ofstream::app;
                std::ofstream outputFile(fpath, mode);

                outputFile << event_name << "," << tag << "," << time_taken << "," << rank << "," << size << std::endl;

                outputFile.close();
            }
        }

        void new_timings() {
            try {
#if defined WASH_CSTONE || defined WASH_WONE
                if (rank != 0) { // Only perform this operation on the primary rank.
                    return;
                }                
#endif 
                std::string fpath = (new std::string(this->path))->append(get_output_name() + ".timings.csv");
                
                if (fs::exists(fpath)) {
                    fs::path old_fpath = this->path + get_output_name() + ".old.timings.csv";
                    fs::rename(fpath, old_fpath);
                }
            } catch (const std::exception& ex) {
                std::cerr << "Error trying to create timings output " << ex.what() << std::endl;
                throw ex;
            }
        }
    };
    
    IOManager::WriterFuncT return_writer(const std::string format);

}

    io::IOManager create_io(const std::string format, const size_t output_nth, const bool use_gather = false, const size_t rank = 0, const size_t size = 1, const bool timings = true);

    std::vector<double> copy_variables();

    // std::vector<std::string> get_force_scalars_names();

    // std::vector<std::string> get_force_vectors_names();

    std::vector<std::string> get_variables_names();
}