/*
This file is part of the EZLib Library
Copyright (C) 2026 Jean-Michel RICHER

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
*/

// ==================================================================
// Author: Jean-Michel Richer
// Email: jean-michel.richer@univ-angers.fr
// Date: Aug 2020
// Last modified: September 2026
// Purpose: interface and class to facilitate use of MPI
// ==================================================================
#pragma once

#include <time.h>
#include <unistd.h>

#include <cassert>
#include <cstdint>
#include <sstream>
#include <stdexcept>
#include <string>
#include <typeinfo>
#include <vector>
using namespace std;
#include <mpi.h>

/**
 * @brief EZ MPI is a wrapper for MPI C++. It simplifies the use
 * of MPI send, receive, gather, scatter functions.
 * For the gather, scatter and reduce functions the processor
 * of rank 0 is considered as the "master" that collects or
 * sends data.
 */

namespace ez {

namespace mpi {

enum { MASTER = 0, SLAVE_1, SLAVE_2, SLAVE_3, SLAVE_4 };

/**
 * @class Process
 * @brief This is the main class that defines a process.
 * This class gathers information about the number of other processes
 * working with this process and the rank of the process also called
 * id here.
 * It enables to perform operations like send and receive but also more
 * complex ones as gather, scatter, broadcast or reduce which are supposed
 * to be initiated by the master process which as the default rank of 0.
 */
class Process {
   protected:
    // number of process working with this one
    int m_nbr_processes;

    // identifier of current process, called rank for MPI
    int m_id;

    // identifier of remote process for to communicate with (send, receive, ...)
    int m_remote;

    // status of last operation
    MPI::Status m_status;

    // message tag if needed (default is 0)
    int m_message_tag;

    // name of process
    string m_name;

    // Linux process identifier (PID)
    int m_pid;

    // verbose mode
    bool m_verbose_flag;

    // verbose mode
    bool m_log_flag;

    // main output stream for current processor
    ostringstream log_stream;

    // temporary output stream
    ostringstream tmp_log;

    // MPI::Finalize already called ?
    bool finalize_already_called;

   private:
    /**
     * Find processor name
     */
    void find_cpu_name() {
        char name[MPI::MAX_PROCESSOR_NAME];
        int length;

        memset(name, 0, MPI::MAX_PROCESSOR_NAME);
        MPI::Get_processor_name(name, length);
        m_name = name;
    }

    /**
     * record output of oss into general output
     */
    // void append();

    /**
     * initialize nbr_process, cpu_rank, processus id
     */
    void init() {
        m_nbr_processes = MPI::COMM_WORLD.Get_size();
        m_id = MPI::COMM_WORLD.Get_rank();
        find_cpu_name();
        m_pid = getpid();
        if (m_verbose_flag) {
            tmp_log << "pid=" << m_pid << ", id=" << m_id << endl;
            flush();
        }
    }

    const std::string current_date_and_time() {
        time_t now = time(0);
        struct tm tstruct;
        char buf[80];
        tstruct = *localtime(&now);
        // strftime(buf, sizeof(buf), "%Y-%m-%d.%X [%s]", &tstruct);
        strftime(buf, sizeof(buf), "%X", &tstruct);

        return buf;
    }

    void flush() {
        string str = current_date_and_time();
        if (m_log_flag) {
            log_stream << str << " cpu " << m_id << "/" << m_nbr_processes
                       << ": " << tmp_log.str();
        }
        if (m_verbose_flag) {
            cerr << str << " cpu " << m_id << "/" << m_nbr_processes << ": "
                 << tmp_log.str();
        }
        tmp_log.str("");
    }

    void print(char v) { tmp_log << v; }
    void print(int v) { tmp_log << v; }
    void print(string v) { tmp_log << v; }
    void print(float v) { tmp_log << v; }
    void print(double v) { tmp_log << v; }

   public:
    /**
     * Default constructor
     * @param argc command line number of arguments
     * @param argv command line arguments
     * @param verbose allow verbose mode
     * @param log allow logging mode
     */
    Process(int argc, char* argv[], bool verbose = true, bool log = true) {
        m_remote = 0;
        m_message_tag = 0;
        m_verbose_flag = verbose;
        m_log_flag = log;
        finalize_already_called = false;
        MPI::Init(argc, argv);
        init();
    }

    /**
     * Free resources of MPI::Init
     */
    void finalize() {
        if (!finalize_already_called) {
            finalize_already_called = true;
            MPI::COMM_WORLD.Barrier();
            MPI::Finalize();
        }
    }

    ~Process() {
        if (!finalize_already_called) {
            finalize_already_called = true;
            MPI::COMM_WORLD.Barrier();
            MPI::Finalize();
        }
    }

    /**
     * set verbose mode
     * @param mode true to enable, false to disable
     */
    void verbose(bool mode) { m_verbose_flag = mode; }

    /**
     * set log mode
     * @param mode true to enable, false to disable
     */
    void log(bool mode) { m_log_flag = mode; }

    /**
     * Get Unix process identifier
     * @return unix process identifier
     */
    int unix_pid() { return m_pid; }

    /**
     * Get processor identifier or rank
     * @return rank of process
     */
    int rank() { return m_id; }

    /**
     * Get processor identifier or rank
     * @return rank of process
     * @note this method is identical to rank()
     */
    int id() { return m_id; }

    /**
     * Get number of processors used
     * @return number of other processes working with this process
     */
    int nbr_processes() { return m_nbr_processes; }

    /**
     * Get processor name
     * @return name of processor on which this process is executed
     */
    string name() { return m_name; }

    /**
     * set remote processor identifier
     * @param rmt_id remote process identifier
     */
    void remote(int rmt_id) { m_remote = rmt_id; }

    /**
     * set message tag
     * @param tag must be an integer between 0 and 32767
     */
    void tag(int tag) { m_message_tag = tag; }

    /**
     * Return true if this processor is the processor of rank 0
     * considered as the master.
     */
    bool is_master() { return (m_id == ez::mpi::MASTER); }

    bool is_slave() { return (m_id != ez::mpi::MASTER); }

    bool is_slave(int n) { return (m_id == n); }

    /**
     * synchronize
     */
    void synchronize() { MPI::COMM_WORLD.Barrier(); }

    /**
     * Determine type of data T and convert it into MPI::Datatype.
     * This function needs to be extended with other types.
     */
    template <class T>
    MPI::Datatype get_type() {
        if (typeid(T) == typeid(char)) {
            return MPI::CHAR;
        } else if (typeid(T) == typeid(int8_t)) {
            return MPI::CHAR;
        } else if (typeid(T) == typeid(uint8_t)) {
            return MPI::CHAR;
        } else if (typeid(T) == typeid(int)) {
            return MPI::INT;
        } else if (typeid(T) == typeid(float)) {
            return MPI::FLOAT;
        } else if (typeid(T) == typeid(double)) {
            return MPI::DOUBLE;
        }
        std::stringstream oss;
        oss << "!!!!! unknown " << typeid(T).name() << endl;
        throw std::runtime_error(oss.str());

        return MPI::INT;
    }

    /**
     * Determine type of data T and convert it into a string
     */
    template <class T>
    std::string get_type_name() {
        if (typeid(T) == typeid(char)) {
            return "char";
        } else if (typeid(T) == typeid(int8_t)) {
            return "char";
        } else if (typeid(T) == typeid(uint8_t)) {
            return "char";
        } else if (typeid(T) == typeid(int)) {
            return "int";
        } else if (typeid(T) == typeid(float)) {
            return "float";
        } else if (typeid(T) == typeid(double)) {
            return "double";
        }
        std::stringstream oss;
        oss << "!!!!! unknown " << typeid(T).name() << endl;
        throw std::runtime_error(oss.str());

        return "char";
    }

    /**
     * Send one instance of data to remote_cpu
     * @param v data to send
     */
    template <class T>
    void send(T& v) {
        MPI::Datatype data_type = get_type<T>();

        if (m_log_flag) {
            tmp_log << "- send value=" << v << " of type "
                    << get_type_name<T>();
            tmp_log << " to process=" << m_remote << endl;
            flush();
        }

        MPI::COMM_WORLD.Send(&v, 1, data_type, m_remote, m_message_tag);
    }

    /**
     * Send an array to remote_cpu
     * @param arr address of the array
     * @param size number of elements to send
     */
    template <class T>
    void send(T* arr, int size, int begin = 0, int end = -1) {
        MPI::Datatype data_type = get_type<T>();

        assert((begin >= 0) and (begin < size));
        if (end < 0) end = size;
        assert((end >= 0) and (end <= size));
        size = end - begin;

        if (m_log_flag) {
            tmp_log << "- send array of type " << get_type_name<T>();
            tmp_log << " from index " << begin << " to " << end;
            tmp_log << " of size " << size << " to process=" << m_remote
                    << endl;
            flush();
        }

        MPI::COMM_WORLD.Send(&arr[begin], size, data_type, m_remote,
                             m_message_tag);
    }

    /**
     * Send data from size data of the vector starting at
     * given index
     * @param v vector
     * @param size number of elements to send
     * @param index where to start, default is first
     *        element of the vector
     *
     */
    template <class T>
    void send(vector<T>& v, int begin = 0, int end = -1) {
        int size = static_cast<int>(v.size());
        assert((begin >= 0) and (begin < size));
        if (end < 0) end = size;
        assert((end >= 0) and (end <= size));
        size = end - begin;

        MPI::Datatype data_type = get_type<T>();

        MPI::COMM_WORLD.Send(&size, 1, MPI::INT, m_remote, m_message_tag);
        T* ptr = v.data();
        MPI::COMM_WORLD.Send(&ptr[begin], size, data_type, m_remote,
                             m_message_tag);

        if (m_log_flag) {
            tmp_log << "- send " << size << " elements of vector";
            tmp_log << " of type " << get_type_name<T>();
            tmp_log << " from index " << begin << " to " << end;
            tmp_log << " of size " << size << " to process=" << m_remote
                    << endl;
            flush();
        }
    }

    /**
     * Receive one instance of data from remote cpu
     * @param v data to receive
     */
    template <class T>
    void recv(T& v) {
        MPI::Datatype data_type = get_type<T>();

        MPI::COMM_WORLD.Recv(
            &v, 1, data_type, m_remote,
            (m_message_tag == 0) ? MPI::ANY_TAG : m_message_tag, m_status);

        if (m_log_flag) {
            tmp_log << "- receive value=" << v << " of type "
                    << get_type_name<T>();
            tmp_log << " from process=" << m_remote << endl;
            flush();
        }
    }

    /**
     * Receive an array of given size
     * @param arr pointer to address of the array
     * @param size number of elements
     */
    template <class T>
    void recv(T* arr, int size) {
        MPI::Datatype data_type = get_type<T>();

        MPI::COMM_WORLD.Recv(
            &arr[0], size, data_type, m_remote,
            (m_message_tag == 0) ? MPI::ANY_TAG : m_message_tag, m_status);

        if (m_log_flag) {
            tmp_log << "- receive array of type " << get_type_name<T>();
            tmp_log << " of size=" << size << " from process=" << m_remote
                    << endl;
            flush();
        }
    }

    template <class T>
    void recv(vector<T>& v) {
        MPI::Datatype data_type = get_type<T>();

        int size = 0;
        MPI::COMM_WORLD.Recv(
            &size, 1, MPI::INT, m_remote,
            (m_message_tag == 0) ? MPI::ANY_TAG : m_message_tag, m_status);

        v.resize(size);

        T* ptr = v.data();
        MPI::COMM_WORLD.Recv(
            ptr, size, data_type, m_remote,
            (m_message_tag == 0) ? MPI::ANY_TAG : m_message_tag, m_status);

        if (m_log_flag) {
            tmp_log << "- receive " << size << " elements of vector";
            tmp_log << " of type " << get_type_name<T>();
            tmp_log << " from process=" << m_remote << endl;
            flush();
        }
    }

    /**
     * Send array and receive value in return, this is an instance
     * of the Sendrecv function.
     * @param arr address of the array to send
     * @param size size of the array to send
     * @param value value to receive
     */
    template <class T, class U>
    void sendrecv(T* array, int size, U& value) {
        MPI::Datatype array_data_type = get_type<T>();
        MPI::Datatype value_data_type = get_type<U>();

        if (m_log_flag) {
            tmp_log << "- sendrecv/send array of type " << get_type_name<T>();
            tmp_log << " of size=" << size << endl;
            flush();
        }

        MPI::COMM_WORLD.Sendrecv(&array[0], size, array_data_type, m_remote, 0,
                                 &value, 1, value_data_type, MPI::ANY_SOURCE,
                                 MPI::ANY_TAG, m_status);

        if (m_log_flag) {
            tmp_log << "sendrecv/receive value=" << value << endl;
            flush();
        }
    }

    /**
     * Perform reduction
     * @param item local data used to perform reduction
     * @param total global data that will contain result
     * @param op operation to perform (MPI::SUM, MPI::MAX, ...)
     */
    template <class T>
    void reduce(T& item, T& total, const MPI::Op& op, int root_id = MASTER) {
        MPI::Datatype data_type = get_type<T>();

        assert((root_id >= 0) and (root_id < m_nbr_processes));

        MPI::COMM_WORLD.Reduce(&item, &total, 1, data_type, op, root_id);

        if (m_log_flag) {
            tmp_log << "- reduction gives value=" << total << endl;
            flush();
        }
    }

    /**
     * Perform gather operation
     * @param lcl_array local array that is send to master process
     * @param glb_array global array that will contain all local arrays
     */
    template <class T>
    void gather(T* lcl_array, int size, T* glb_array, int root_id = MASTER) {
        MPI::Datatype data_type = get_type<T>();

        assert((root_id >= 0) and (root_id < m_nbr_processes));

        MPI::COMM_WORLD.Gather(lcl_array, size, data_type, glb_array, size,
                               data_type, root_id);

        if (m_log_flag) {
            tmp_log << "- gather" << endl;
            flush();
        }
    }

    /**
     * Perform scatter operation
     * @param big_array array of data that will be send by to all processors
     * by parts
     * @param big_size size of the array that will be sent
     * @param small_array local array of data
     */
    template <class T>
    void scatter(T* big_array, int big_size, T* small_array) {
        MPI::Datatype data_type = get_type<T>();

        int small_size = big_size / m_nbr_processes;
        MPI::COMM_WORLD.Scatter(big_array, small_size, data_type, small_array,
                                small_size, data_type, 0);

        if (m_log_flag) {
            tmp_log << "- scatter array of size " << big_size;
            tmp_log << " into " << m_nbr_processes << " chunks of "
                    << small_size;
            tmp_log << " elements" << endl;
            flush();
        }
    }

    /**
     * Perform scatter operation
     * @param big_vec vector of data that will be send by to all processors
     * by parts
     * @param local_vec local vector of data
     */
    template <class T>
    void scatter(vector<T>& big_vec, int big_size, vector<T>& local_vec,
                 int begin = 0, int end = -1) {
        MPI::Datatype data_type = get_type<T>();

        assert((begin >= 0) and (begin < big_size));
        if (end == -1) end = big_size;
        assert((begin < end) and (end <= big_size));
        big_size = (end - begin + 1);

        int local_vec_size = big_size / m_nbr_processes;
        MPI::COMM_WORLD.Scatter(big_vec.data(), local_vec_size, data_type,
                                local_vec.data(), local_vec_size, data_type, 0);

        if (m_log_flag) {
            tmp_log << "- scatter vector of size " << big_size;
            tmp_log << " into " << m_nbr_processes << " chunks of "
                    << local_vec_size;
            tmp_log << " elements" << endl;
            flush();
        }
    }

    // MPI_Bcast(void *buffer, int count, MPI_Datatype datatype, int root,
    // MPI_Comm comm)
    template <class T>
    void broadcast(T* glb_array, int size, int root_id = MASTER) {
        MPI::Datatype data_type = get_type<T>();

        assert((root_id >= 0) and (root_id < m_nbr_processes));

        MPI::COMM_WORLD.Bcast((void*)glb_array, size, data_type, root_id);

        if (m_log_flag) {
            tmp_log << "- broadcast array of size " << size << endl;
            flush();
        }
    }

    template <class T>
    void broadcast(std::vector<T>& vec, int root_id = MASTER) {
        MPI::Datatype data_type = get_type<T>();

        assert((root_id >= 0) and (root_id < m_nbr_processes));

        MPI::COMM_WORLD.Bcast((void*)vec.data(), static_cast<int>(vec.size()),
                              data_type, root_id);

        if (m_log_flag) {
            tmp_log << "- broadcast array of size " << vec.size()
                    << " elements";
            tmp_log << " from process=" << root_id << endl;
            flush();
        }
    }

    typedef std::ostream& (*ManipFn)(std::ostream&);
    typedef std::ios_base& (*FlagsFn)(std::ios_base&);

    template <class T>
    Process& operator<<(vector<T>& v) {
        tmp_log << '[';
        if (v.size() > 0) {
            tmp_log << v[0];
            for (int i = 1; i < static_cast<int>(v.size()); ++i) {
                tmp_log << ' ' << v[i];
            }
        }
        tmp_log << ']';
        return *this;
    }

    template <class T>  // int, double, strings, etc
    Process& operator<<(const T& output) {
        tmp_log << output;
        return *this;
    }

    // endl, flush, setw, setfill, etc.
    Process& operator<<(ManipFn manip) {
        manip(tmp_log);
        if (manip == static_cast<ManipFn>(std::flush) ||
            manip == static_cast<ManipFn>(std::endl)) {
            this->flush();
        }
        return *this;
    }

    // setiosflags, resetiosflags
    Process& operator<<(FlagsFn manip) {
        manip(tmp_log);
        return *this;
    }
    void logs(ostream& out) {
        MPI::COMM_WORLD.Barrier();
        m_verbose_flag = m_log_flag = false;
        if (m_id == 0) {
            out.flush();
            out << std::endl;
            out << "=====================" << std::endl;
            out << "===    LOGGING    ===" << std::endl;
            out << "=====================" << std::endl;
            out << "---------------------" << std::endl;
            out << "CPU " << m_id << " (pid=" << unix_pid() << ")" << std::endl;
            out << "---------------------" << std::endl;
            out << log_stream.str();
            out.flush();
            remote(1);
            int token = -255;
            send(token);
        } else {
            remote(m_id - 1);
            int token;
            recv(token);
            out << "---------------------" << std::endl;
            out << "CPU " << m_id << " (pid=" << unix_pid() << ")" << std::endl;
            out << "---------------------" << std::endl;
            out << log_stream.str();
            out.flush();
            if (m_id < m_nbr_processes - 1) {
                remote(m_id + 1);
                token = -255;
                send(token);
            }
        }
    }

    typedef void (*Code)(Process& p);

    void run(Code code) { code(*this); }
};

}  // end of namespace mpi

}  // end of namespace ez