/*******************************************************************************
* Copyright 2021 Intel Corporation
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
*     http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*******************************************************************************/

#pragma once

// IMPORTANT! This file should not be included to any other non-mpi header,
// otherwise it creates dependeny on MPI at the user's application compile time
// TODO: In the future this can be solved via __has_include C++17 feature

#include <mpi.h>
#include <oneapi/dal/array.hpp>
#include "oneapi/dal/detail/communicator.hpp"

namespace spmd = oneapi::dal::preview::spmd;

namespace oneapi::dal::detail {
namespace v1 {

using preview::spmd::communication_error;

inline void mpi_call(int mpi_status) {
    if (mpi_status == MPI_SUCCESS) {
        return;
    }
    switch (mpi_status) {
        case MPI_ERR_BUFFER:
            throw communication_error(dal::detail::error_messages::invalid_buffer());
        case MPI_ERR_COUNT: throw communication_error(dal::detail::error_messages::invalid_count());
        case MPI_ERR_TYPE:
            throw communication_error(dal::detail::error_messages::invalid_data_type());
        case MPI_ERR_OP: throw communication_error(dal::detail::error_messages::invalid_op());
        case MPI_ERR_COMM:
            throw communication_error(dal::detail::error_messages::invalid_mpi_comm());
        case MPI_ERR_ROOT: throw communication_error(dal::detail::error_messages::invalid_root());
        default: throw internal_error(dal::detail::error_messages::unknown_mpi_error());
    }
}

inline MPI_Datatype make_mpi_data_type(const data_type& dtype) {
    switch (dtype) {
        case data_type::int8: return MPI_INT8_T;
        case data_type::int16: return MPI_INT16_T;
        case data_type::int32: return MPI_INT32_T;
        case data_type::int64: return MPI_INT64_T;
        case data_type::uint8: return MPI_UINT8_T;
        case data_type::uint16: return MPI_UINT16_T;
        case data_type::uint32: return MPI_UINT32_T;
        case data_type::uint64: return MPI_UINT64_T;
        case data_type::float32: return MPI_FLOAT;
        case data_type::float64: return MPI_DOUBLE;
        case data_type::bfloat16: return MPI_UINT16_T;
        default: ONEDAL_ASSERT(!"Unknown data type");
    }
    return MPI_DATATYPE_NULL;
}

inline MPI_Op make_mpi_reduce_op(const preview::spmd::reduce_op& op) {
    switch (op) {
        case spmd::reduce_op::max: return MPI_MAX;
        case spmd::reduce_op::min: return MPI_MIN;
        case spmd::reduce_op::sum: return MPI_SUM;
        default: ONEDAL_ASSERT(!"Unknown reduce operation");
    }
    return MPI_OP_NULL;
}

class mpi_request_impl : public preview::spmd::request_iface {
public:
    explicit mpi_request_impl(const MPI_Request& request) : mpi_request_(request) {}

    mpi_request_impl(const mpi_request_impl&) = delete;
    mpi_request_impl& operator=(const mpi_request_impl&) = delete;

    void wait() override {
        mpi_call(MPI_Wait(&mpi_request_, MPI_STATUS_IGNORE));
    }

    bool test() override {
        int flag;
        mpi_call(MPI_Test(&mpi_request_, &flag, MPI_STATUS_IGNORE));
        return bool(flag != 0);
    }

private:
    MPI_Request mpi_request_;
};

/// Implementation of the low-level SPMD communicator interface via MPI
/// TODO: Currently message sizes are limited via `int` type.
///       Large message sizes should be handled on the communicator side in the future.
template <typename MemoryAccessKind>
class mpi_communicator_impl : public via_host_interface_selector<MemoryAccessKind>::type {
public:
    // Explicitly declare all virtual functions with overloads to workaround Clang warning
    // https://stackoverflow.com/questions/18515183/c-overloaded-virtual-function-warning-by-clang
    using base_t = typename via_host_interface_selector<MemoryAccessKind>::type;
    using base_t::bcast;
    using base_t::allgatherv;
    using base_t::allreduce;

    explicit mpi_communicator_impl(std::int64_t default_root = 0)
            : mpi_comm_(MPI_COMM_WORLD),
              default_root_(default_root) {}

#ifdef ONEDAL_DATA_PARALLEL
    //    template<typename T = MemoryAccessKind, spmd::enable_if_device_memory_accessible_t<T>>
    explicit mpi_communicator_impl(sycl::queue& queue, std::int64_t default_root = 0)
            : base_t(queue),
              mpi_comm_(MPI_COMM_WORLD),
              default_root_(default_root) {}
#endif

    std::int64_t get_rank() override {
        if (rank_ < 0) {
            int rank;
            mpi_call(MPI_Comm_rank(mpi_comm_, &rank));
            rank_ = rank;
        }
        return rank_;
    }

    std::int64_t get_rank_count() override {
        if (rank_count_ < 0) {
            int rank_count;
            mpi_call(MPI_Comm_size(mpi_comm_, &rank_count));
            rank_count_ = rank_count;
        }
        return rank_count_;
    }

    std::int64_t get_default_root_rank() override {
        return default_root_;
    }

    void barrier() override {
        mpi_call(MPI_Barrier(mpi_comm_));
    }

    spmd::request_iface* bcast(byte_t* send_buf,
                               std::int64_t count,
                               const data_type& dtype,
                               std::int64_t root) override {
        ONEDAL_ASSERT(root >= 0);

        if (count == 0) {
            return nullptr;
        }

        ONEDAL_ASSERT(send_buf);
        ONEDAL_ASSERT(count > 0);

        // TODO replace with MPI_Ibcast
        mpi_call(MPI_Bcast(send_buf,
                           integral_cast<int>(count),
                           make_mpi_data_type(dtype),
                           integral_cast<int>(root),
                           mpi_comm_));
        return nullptr;
    }

    spmd::request_iface* allgatherv(const byte_t* send_buf,
                                    std::int64_t send_count,
                                    byte_t* recv_buf,
                                    const std::int64_t* recv_counts,
                                    const std::int64_t* displs,
                                    const data_type& dtype) override {
        ONEDAL_ASSERT(root >= 0);

        ONEDAL_ASSERT(send_buf);
        ONEDAL_ASSERT(recv_counts);
        ONEDAL_ASSERT(displs);
        ONEDAL_ASSERT(recv_buf);

        array<int> recv_counts_int;
        array<int> displs_int;

        const std::int64_t rank_count = get_rank_count();
        recv_counts_int.reset(rank_count);
        displs_int.reset(rank_count);

        auto recv_counts_int_ptr = recv_counts_int.get_mutable_data();
        auto displs_int_ptr = displs_int.get_mutable_data();

        [[maybe_unused]] std::int64_t displs_counter = 0;
        for (std::int64_t i = 0; i < rank_count; ++i) {
            ONEDAL_ASSERT(recv_counts[i] > 0);
            ONEDAL_ASSERT(displs[i] >= displs_counter);
            displs_counter += recv_counts[i];

            recv_counts_int_ptr[i] = dal::detail::integral_cast<int>(recv_counts[i]);
            displs_int_ptr[i] = dal::detail::integral_cast<int>(displs[i]);
        }

        MPI_Request mpi_request;
        mpi_call(MPI_Iallgatherv(send_buf,
                                 integral_cast<int>(send_count),
                                 make_mpi_data_type(dtype),
                                 recv_buf,
                                 recv_counts_int.get_data(),
                                 displs_int.get_data(),
                                 make_mpi_data_type(dtype),
                                 mpi_comm_,
                                 &mpi_request));

        return new mpi_request_impl{ mpi_request };
    }

    spmd::request_iface* allreduce(const byte_t* send_buf,
                                   byte_t* recv_buf,
                                   std::int64_t count,
                                   const data_type& dtype,
                                   const spmd::reduce_op& op) override {
        if (count == 0) {
            return nullptr;
        }

        ONEDAL_ASSERT(send_buf);
        ONEDAL_ASSERT(recv_buf);

        // Intel MPI requires buffers to be not aliased
        // However, communicator interface allows aliased buffers
        // TODO: Implement correct aliasing check
        if (send_buf != recv_buf) {
            MPI_Request mpi_request;
            mpi_call(MPI_Iallreduce(send_buf,
                                    recv_buf,
                                    integral_cast<int>(count),
                                    make_mpi_data_type(dtype),
                                    make_mpi_reduce_op(op),
                                    mpi_comm_,
                                    &mpi_request));
            return new mpi_request_impl{ mpi_request };
        }
        else {
            const std::int64_t dtype_size = get_data_type_size(dtype);
            const std::int64_t size = check_mul_overflow(count, dtype_size);
            auto recv_buf_backup = array<byte_t>::empty(size);

            // TODO Replace with MPI_Iallreduce
            mpi_call(MPI_Allreduce(send_buf,
                                   recv_buf_backup.get_mutable_data(),
                                   integral_cast<int>(count),
                                   make_mpi_data_type(dtype),
                                   make_mpi_reduce_op(op),
                                   mpi_comm_));

            memcpy(default_host_policy{}, recv_buf, recv_buf_backup.get_data(), size);

            // We have to copy memory after reduction, this cannot be performed
            // asynchronously in the current implementation, so we return `nullptr`
            // indicating that operation was performed synchronously
            return nullptr;
        }
    }

private:
    MPI_Comm mpi_comm_;
    std::int64_t default_root_;
    std::int64_t rank_ = -1;
    std::int64_t rank_count_ = -1;
};

template <typename MemoryAccessKind>
class mpi_communicator : public spmd::communicator<MemoryAccessKind> {
public:
#ifdef ONEDAL_DATA_PARALLEL
    template <typename T = MemoryAccessKind,
              typename = spmd::enable_if_device_memory_accessible_t<T>>
    explicit mpi_communicator(sycl::queue& queue, std::int64_t default_root = 0)
            : spmd::communicator<MemoryAccessKind>(
                  new mpi_communicator_impl<MemoryAccessKind>(queue, default_root)) {}
#endif
    explicit mpi_communicator(std::int64_t default_root = 0)
            : spmd::communicator<MemoryAccessKind>(
                  new mpi_communicator_impl<MemoryAccessKind>(default_root)) {}
};

} // namespace v1

using v1::mpi_communicator;

} // namespace oneapi::dal::detail
