// ========================================================================== //
// Class inclusion
#include "Options/options.hpp"

// Local libraries
#include "Utilities/string_operations.hpp"

// Parallellisation
#if USE_MPI == 1
    #include <mpi.h>
#endif

// Standard libraries
#include <filesystem>
#include <fstream>
#include <iostream>
#include <map>
#include <sstream>
#include <string>
#include <vector>

#ifndef MPI_VERSION
    #define MPI_VERSION 0
#endif

// ========================================================================== //
// IO functions

// Function designed to read in file with options
void Options::read_file(std::string opt_file) {

    // Save the filename to the class
    filename = opt_file;

#if MPI_VERSION > 0
    int mpi_rank, mpi_size;
    MPI_Comm_rank(MPI_COMM_WORLD, &mpi_rank);
    MPI_Comm_size(MPI_COMM_WORLD, &mpi_size);
    if (mpi_rank != 0) {
        mpi_share();
        return;
    }
#endif

    // Initialise the streams and variables needed.
    std::ifstream dataFile(opt_file);
    std::string   line;

    // Loop through all lines in the file
    while (getline(dataFile, line)) {

        std::string              cell;
        std::stringstream        lineStream(line);
        std::vector<std::string> input;

        // Read the line and count how many entries are there
        int n = 0;
        while (getline(lineStream, cell, ' ')) {
            input.push_back(cell);
            n++;
        }

        // This is a newline
        if (n == 0) continue;

        // Ignore know comments:
        if (std::string::npos != input[0].find("#") ||
            std::string::npos != input[0].find("//") ||
            std::string::npos != input[0].find("!") ||
            std::string::npos != input[0].find("%"))
            continue;

        // Prevent multiple definitions of same option
        if (options.count(input[0]) > 0)
            Utilities::error(__FILE__, __LINE__,
                             "Option definitioned already: " + input[0]);

        // Add input value to list of values
        std::vector<std::string> val;
        for (size_t i = 1; i < input.size(); i++) val.push_back(input[i]);

        this->set(input[0], val);
    }
    dataFile.close();

#if MPI_VERSION > 0
    mpi_share();
#endif
}

// Function designed to write options into file
void Options::write_file() { write_file(filename); }
void Options::write_file(std::string filename) {

#if MPI_VERSION > 0
    int mpi_rank;
    MPI_Comm_rank(MPI_COMM_WORLD, &mpi_rank);
    if (mpi_rank != 0) return;
#endif

    // Inialise the streams and variables needed.
    std::ofstream dataFile(filename, std::ofstream::trunc);

    // Loop through all lines in the file
    for (auto opt = options.begin(); opt != options.end(); opt++) {

        dataFile << (*opt).first;
        for (auto v : values[(*opt).second]) dataFile << " " << v;
        dataFile << "\n";
    }
    dataFile.close();
}

void send_string(std::string &str, int from, int mpi_rank, int mpi_size) {

    int str_length = str.size() + 1;
    MPI_Bcast(&str_length, 1, MPI_INT, from, MPI_COMM_WORLD);

    char buf[str_length];

    // Send the string over
    if (mpi_rank == from) {
        for (int rank = 0; rank < mpi_size; rank++) {
            if (rank == from) continue;
            MPI_Send(str.c_str(), str_length, MPI_CHAR, rank, rank,
                     MPI_COMM_WORLD);
        }
    }

    // Recieve the string
    else {
        MPI_Status stat;
        MPI_Recv(&buf, str_length, MPI_CHAR, 0, mpi_rank, MPI_COMM_WORLD,
                 &stat);
        str = std::string(buf);
    }
}

// ========================================================================== //
// MPI Communication

void Options::mpi_share() {

#if MPI_VERSION == 0
    return;
#endif

    // Setup variables etc
    int mpi_rank, mpi_size;
    MPI_Comm_rank(MPI_COMM_WORLD, &mpi_rank);
    MPI_Comm_size(MPI_COMM_WORLD, &mpi_size);

    MPI_Barrier(MPI_COMM_WORLD);

    // ====================================================================== //
    // Send over and construct the map

    int nr_labels = options.size();

    // Broadcast the number of known options
    MPI_Bcast(&nr_labels, 1, MPI_INT, 0, MPI_COMM_WORLD);

    // Setup the required vectors
    std::vector<std::string> labels(nr_labels);
    std::vector<int>         indices(nr_labels);
    std::vector<int>         labels_sizes(nr_labels);

    // Construct the vectors on rank 0
    if (mpi_rank == 0) {
        int idx = 0;
        for (auto opt : options) {
            labels[idx] = opt.first, indices[idx] = opt.second;
            labels_sizes[idx] = opt.first.size();
            idx++;
        }
    }

    // Broadcast the index list and the string lengths
    MPI_Bcast(indices.data(), nr_labels, MPI_INT, 0, MPI_COMM_WORLD);
    MPI_Bcast(labels_sizes.data(), nr_labels, MPI_INT, 0, MPI_COMM_WORLD);

    // Send over the strings
    for (int i = 0; i < nr_labels; i++) {
        send_string(labels[i], 0, mpi_rank, mpi_size);
        options[labels[i]] = indices[i];
    }

    // ====================================================================== //
    // Send the value lists

    std::vector<std::vector<std::string>> value_list(nr_labels);
    std::vector<int>                      value_lengths(nr_labels);

    // Construct value list and length
    if (mpi_rank == 0) {
        for (int i = 0; i < nr_labels; i++) {
            int idx            = indices[i];
            value_list[idx]    = values[idx];
            value_lengths[idx] = values[idx].size();
        }
    }

    MPI_Bcast(value_lengths.data(), nr_labels, MPI_INT, 0, MPI_COMM_WORLD);

    // Send over the strings
    for (int i = 0; i < nr_labels; i++) {
        size_t idx = indices[i];
        if (mpi_rank != 0)
            value_list[idx] = std::vector<std::string>(value_lengths[idx]);

        for (int j = 0; j < value_lengths[idx]; j++) {
            send_string(value_list[idx][j], 0, mpi_rank, mpi_size);
        }
    }
    values = value_list;
}

// End of File
// ========================================================================== //