/*
 * GridTools
 *
 * Copyright (c) 2014-2019, ETH Zurich
 * All rights reserved.
 *
 * Please, refer to the LICENSE file in the root directory.
 * SPDX-License-Identifier: BSD-3-Clause
 */
#pragma once

#include "../../common/array.hpp"
#include "../../common/boollist.hpp"
#include "../GCL.hpp"
#include <cmath>
#include <iostream>
#include <string>

// This file needs to be changed

namespace gridtools {

#ifdef GCL_MPI
    /** \class MPI_3D_process_grid_t
     * Class that provides a representation of a 3D process grid given an MPI CART
     * It requires the MPI CART to be defined before the grid is created
     * \tparam CYCLIC is a template argument matching \ref boollist_concept to specify periodicities
     * \n
     * This is a process grid matching the \ref proc_grid_concept concept
     */
    template <int Ndims>
    struct MPI_3D_process_grid_t {

        /** number of dimensions
         */
        static const int ndims = Ndims;

        typedef std::true_type has_communicator;
        typedef gridtools::boollist<ndims> period_type;

      private:
        MPI_Comm m_communicator; // Communicator that is associated with the MPI CART!
        period_type m_cyclic;
        int m_nprocs;
        gridtools::array<int, ndims> m_dimensions;
        gridtools::array<int, ndims> m_coordinates;

      public:
        MPI_3D_process_grid_t(MPI_3D_process_grid_t const &other) : m_cyclic(other.cyclic()), m_nprocs(other.m_nprocs) {

            MPI_Comm_dup(other.m_communicator, &m_communicator);

            for (uint_t i = 0; i < ndims; ++i) {
                m_dimensions[i] = other.m_dimensions[i];
                m_coordinates[i] = other.m_coordinates[i];
            }
        }

        /** Constructor that takes an MPI CART communicator, already configured, and use it to set up the process grid.
            \param c Object containing information about periodicities as defined in \ref boollist_concept
            \param comm MPI Communicator describing the MPI 3D computing grid
        */
        MPI_3D_process_grid_t(period_type const &c, MPI_Comm const &comm)
            : m_cyclic(c), m_nprocs(0), m_dimensions(), m_coordinates() {

            MPI_Comm_dup(comm, &m_communicator);

            int period[ndims];
            MPI_Cart_get(comm, ndims, &m_dimensions[0], period, &m_coordinates[0]);
            MPI_Comm_size(comm, &m_nprocs);
        }

        /** Constructor that takes an MPI CART communicator, already configured, and use it to set up the process grid.
            \param c Object containing information about periodicities as defined in \ref boollist_concept
            \param comm MPI Communicator describing the MPI 3D computing grid
            \param dims Array of dimensions of the processor grid
        */
        template <typename Array>
        MPI_3D_process_grid_t(period_type const &c, MPI_Comm const &comm, Array const &dims)
            : m_communicator(), m_cyclic(c), m_nprocs(0), m_dimensions(dims), m_coordinates() {
            MPI_Comm_size(comm, &m_nprocs);
            MPI_Dims_create(m_nprocs, dims.size(), &m_dimensions[0]);
            int period[3] = {1, 1, 1};
            MPI_Cart_create(comm, 3, &m_dimensions[0], period, false, &m_communicator);
            MPI_Cart_get(m_communicator, ndims, &m_dimensions[0], period /*does not really care*/, &m_coordinates[0]);
        }

        ~MPI_3D_process_grid_t() { MPI_Comm_free(&m_communicator); }

        /**
           Returns communicator
        */
        MPI_Comm communicator() const { return m_communicator; }

        /** Returns in t_R and t_C the lenght of the dimensions of the process grid AS PRESCRIBED BY THE CONCEPT
            \param[out] t_R Number of elements in first dimension
            \param[out] t_C Number of elements in second dimension
            \param[out] t_S Number of elements in third dimension
        */
        void dims(int &t_R, int &t_C, int &t_S) const {
            GT_STATIC_ASSERT(ndims == 3, "this interface supposes ndims=3");
            t_R = m_dimensions[0];
            t_C = m_dimensions[1];
            t_S = m_dimensions[2];
        }

        /** Returns the dimensions in an array of dimensions (at least of size 3)
            \tparam The array type
            \param array The array where to put the values
        */
        template <class Array>
        void fill_dims(Array &array) const {
            GT_STATIC_ASSERT(ndims == 3, "this interface supposes ndims=3");
            array[0] = m_dimensions[0];
            array[1] = m_dimensions[1];
            array[2] = m_dimensions[2];
        }

        void dims(int &t_R, int &t_C) const {
            GT_STATIC_ASSERT(ndims == 2, "this interface supposes ndims=2");
            t_R = m_dimensions[0];
            t_C = m_dimensions[1];
        }

        /** Returns the number of processors of the processor grid

            \return Number of processors
        */
        uint_t size() const {
            uint_t ret = m_dimensions[0];
            for (uint_t i = 1; i < ndims; ++i)
                ret *= m_dimensions[i];
            return ret;
        }

        /** Returns in t_R and t_C the coordinates ot the caller process in the grid AS PRESCRIBED BY THE CONCEPT
            \param[out] t_R Coordinate in first dimension
            \param[out] t_C Coordinate in second dimension
            \param[out] t_S Coordinate in third dimension
        */
        void coords(int &t_R, int &t_C, int &t_S) const {
            GT_STATIC_ASSERT(ndims == 3, "this interface supposes ndims=3");
            t_R = m_coordinates[0];
            t_C = m_coordinates[1];
            t_S = m_coordinates[2];
        }

        void coords(int &t_R, int &t_C) const {
            GT_STATIC_ASSERT(ndims == 2, "this interface supposes ndims=2");
            t_R = m_coordinates[0];
            t_C = m_coordinates[1];
        }

        /** Returns the process ID of the process with relative coordinates (I,J) with respect to the caller process AS
           PRESCRIBED BY THE CONCEPT
            \tparam I Relative coordinate in the first dimension
            \tparam J Relative coordinate in the second dimension
            \tparam K Relative coordinate in the third dimension
            \return The process ID of the required process
        */
        template <int I, int J, int K>
        int proc() const {
            // int coords[3]={I,J,K};
            return proc(I, J, K);
        }

        int pid() const {
            int rank;
            MPI_Comm_rank(m_communicator, &rank);
            return rank;
        }

        /** Returns the process ID of the process with relative coordinates (I,J) with respect to the caller process AS
           PRESCRIBED BY THE CONCEPT
            \param[in] I Relative coordinate in the first dimension
            \param[in] J Relative coordinate in the second dimension
            \param[in] K Relative coordinate in the third dimension
            \return The process ID of the required process
        */
        int proc(int I, int J, int K) const {
            int _coords[3];

            if (m_cyclic.value(0))
                _coords[0] = (m_coordinates[0] + I + m_dimensions[0]) % m_dimensions[0];
            else {
                _coords[0] = m_coordinates[0] + I;
                if (_coords[0] < 0 || _coords[0] >= m_dimensions[0])
                    return -1;
            }

            if (m_cyclic.value(1))
                _coords[1] = (m_coordinates[1] + J + m_dimensions[1]) % m_dimensions[1];
            else {
                _coords[1] = m_coordinates[1] + J;
                if (_coords[1] < 0 || _coords[1] >= m_dimensions[1])
                    return -1;
            }

            if (m_cyclic.value(2))
                _coords[2] = (m_coordinates[2] + K + m_dimensions[2]) % m_dimensions[2];
            else {
                _coords[2] = m_coordinates[2] + K;
                if (_coords[2] < 0 || _coords[2] >= m_dimensions[2])
                    return -1;
            }

            int pid = 0;
            MPI_Comm_rank(MPI_COMM_WORLD, &pid);
            int res;
            MPI_Cart_rank(m_communicator, _coords, &res);
            return res;
        }

        GT_FUNCTION
        gridtools::array<int, ndims> const &coordinates() const { return m_coordinates; }

        GT_FUNCTION
        gridtools::array<int, ndims> const &dimensions() const { return m_dimensions; }

        /** Returns the process ID of the process with absolute coordinates specified by the input gridtools::array of
           coordinates
            \param[in] crds gridtools::aray of coordinates of the processor of which the ID is needed

            \return The process ID of the required process
        */
        int abs_proc(gridtools::array<int, ndims> const &crds) const {
            return proc(crds[0] - m_coordinates[0], crds[1] - m_coordinates[1], crds[2] - m_coordinates[2]);
        }

        int ntasks() { return m_nprocs; }

        bool periodic(int index) const {
            assert(index < ndims);
            return m_cyclic.value(index);
        }

        array<bool, ndims> periodic() const {
            GT_STATIC_ASSERT(period_type::m_size == ndims, "Dimensions not matching");
            return m_cyclic.value();
        }

        period_type const &cyclic() const { return m_cyclic; }

        int const &coordinates(uint_t const &i) const { return m_coordinates[i]; }
        int const &dimensions(uint_t const &i) const { return m_dimensions[i]; }
    };

#endif

} // namespace gridtools
