/** @file phase_transition_mpi.cpp
 *
 *  @author Cory Alexander Balaton (coryab)
 *  @author Janita Ovidie Sandtrøen Willumsen (janitaws)
 *
 *  @version 1.0
 *
 *  @brief Sweep over different temperatures and generate data.
 *
 *  @details This program takes in 4 arguments: the start temperature,
 *  the end temperature, the amount of temperature points to simulate, and
 *  the amount of monte carlo samples to collect, in that order.
 *
 *  @bug No known bugs
 * */
#include "data_type.hpp"
#include "monte_carlo.hpp"
#include "utils.hpp"

#include <getopt.h>
#include <mpi.h>
#include <string>

/** @brief A function that displays how to use the program and quits.*/
void usage(std::string filename)
{
    std::cout
        << "Usage: " << filename
        << " <start temperature> <end temperature> <lattice size> "
           "<points> <cycles> <burn-in-time> <output file>\n"
        << "This should be used with mpiexec or mpirun for maximum "
           "performance\n\n"
        << "\t[ -h | --help ]\n";
    exit(-1);
}

/** @brief The main function.*/
int main(int argc, char **argv)
{
    // Command options
    struct option long_options[] = {{"help", 0, 0, 0}, {NULL, 0, NULL, 0}};

    int option_index = -1;
    int c;

    while (true) {
        c = getopt_long(argc, argv, "h", long_options, &option_index);

        if (c == -1)
            break;

        switch (c) {
        case 0:
            switch (option_index) {
            case 0: // Not a mistake. This just goes to the default.
            default:
                usage(argv[0]);
            }
            break;
        case 'h':
        default:
            usage(argv[0]);
        }
    }
    // Check that the number of arguments is at least 8.
    if (argc < 8) {
        usage(argv[0]);
    }

    // Timing variables
    double t0, t1;
    t0 = MPI_Wtime();

    // Define/initialize variables
    double start = atof(argv[1]), end = atof(argv[2]);
    int points = atoi(argv[3]), cycles = atoi(argv[5]), L = atoi(argv[4]),
        burn_in_time = atoi(argv[6]), N = L * L;
    double dt = (end - start) / points;
    std::ofstream ofile;
    std::string outfile = argv[7];
    data_t data[points];

    // MPI specific variables
    int rank, cluster_size;

    // Initialize MPI
    MPI_Init(&argc, &argv);

    // Get the cluster size and rank
    MPI_Comm_size(MPI_COMM_WORLD, &cluster_size);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    int remainder = points % cluster_size;
    double i_start; // What temperature to start from
    int i_points;   // How many points to simulate

    // Distribute temperature points
    if (rank < remainder) {
        i_points = points / cluster_size + 1;
        i_start = start + dt * i_points * rank;
    }
    else {
        i_points = points / cluster_size;
        i_start = start + dt * (i_points * rank + remainder);
    }

    // Initialize array to contains data for each temperature point
    data_t i_data[i_points];

    // Simulate and save data to array
    for (size_t i = 0; i < i_points; i++) {
        i_data[i] = montecarlo::mcmc_parallel(L, i_start + dt * i, cycles,
                                              burn_in_time);
    }

    // Rank 0 collects all the data and copies it to the "master"
    // data array.
    if (rank == 0) {
        // Copy its own i_data to the data array
        std::copy_n(i_data, i_points, data);

        // Collect i_data from other ranks in order and copy to data.
        for (size_t i = 1; i < cluster_size; i++) {
            if (rank < remainder) {
                MPI_Recv((void *)i_data,
                         sizeof(data_t) * (points / cluster_size + 1), MPI_CHAR,
                         i, MPI_ANY_TAG, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
                std::copy_n(i_data, points / cluster_size + 1,
                            data + (points / cluster_size) * i);
            }
            else {
                MPI_Recv((void *)i_data,
                         sizeof(data_t) * (points / cluster_size), MPI_CHAR, i,
                         MPI_ANY_TAG, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
                std::copy_n(i_data, points / cluster_size,
                            data + (points / cluster_size) * i + remainder);
            }
        }

        // Write everything from data to file
        utils::mkpath(utils::dirname(outfile));
        ofile.open(outfile);

        double temp, CV, X;

        using utils::scientific_format;
        for (size_t i = 0; i < points; i++) {
            temp = start + dt * i;
            CV = (data[i].E2 - data[i].E * data[i].E)
                 / ((double)N * temp * temp);
            X = (data[i].M2 - data[i].M_abs * data[i].M_abs)
                / ((double)N * temp);

            ofile << scientific_format(temp) << ','
                  << scientific_format(data[i].E / N) << ','
                  << scientific_format(data[i].M_abs / N) << ','
                  << scientific_format(CV) << ',' << scientific_format(X)
                  << '\n';
        }
        ofile.close();
    }
    // For all other ranks, send the data to rank 0
    else {
        MPI_Send(i_data, i_points * sizeof(data_t), MPI_CHAR, 0, rank,
                 MPI_COMM_WORLD);
    }

    t1 = MPI_Wtime();

    if (rank == 0) {
        std::cout << "Time: " << t1 - t0 << " seconds\n";
    }

    MPI_Finalize();
}