#ifndef CAFFE2_MPI_MPI_OPS_H_
#define CAFFE2_MPI_MPI_OPS_H_

#include <mpi.h>

#include "caffe2/core/operator.h"
#include "caffe2/mpi/mpi_common.h"

namespace caffe2 {

// TODO(jiayq): if needed, write up the use of color and key with MPI split.
// Currently, the operator simply creates a communicator that has the
// same topology as the Caffe2 global communicator.
template <class Context>
class MPICreateCommonWorldOp final : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;
  MPICreateCommonWorldOp(const OperatorDef& operator_def, Workspace* ws)
      : Operator<Context>(operator_def, ws) {}

  bool RunOnDevice() override {
    OperatorBase::Outputs()[0]->Reset(new MPICommonWorldWrapper());
    return true;
  }
};

template <class Context>
class MPIBroadcastOp final : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;
  MPIBroadcastOp(const OperatorDef& operator_def, Workspace* ws)
      : Operator<Context>(operator_def, ws),
        root_(OperatorBase::template GetSingleArgument<int>("root", 0)) {}
  ~MPIBroadcastOp() {}

  bool RunOnDevice() override {
    MPI_Comm comm = OperatorBase::Input<MPICommonWorldWrapper>(0).comm();
    CAFFE_ENFORCE(
        OperatorBase::OutputIsTensorType(0, Context::GetDeviceType()),
        "Output is of wrong type.");
    auto* output = Output(0);
    // Make sure that output is already allocated.
    CAFFE_ENFORCE(
        output->numel() > 0,
        "Broadcast op uses in-place operation so the output "
        "should be already allocated.");
    MPI_CHECK(MPI_Bcast(
        output->raw_mutable_data(),
        output->nbytes(),
        MPIDataTypeWrapper<char>::type(),
        root_,
        comm));
    return true;
  }

 protected:
  int root_;
};

// MPIReduceOp does Reduce using MPI. Currently, only SUM is supported.
template <typename T, class Context>
class MPIReduceOp final : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;
  MPIReduceOp(const OperatorDef& operator_def, Workspace* ws)
      : Operator<Context>(operator_def, ws),
        root_(OperatorBase::template GetSingleArgument<int>("root", 0)) {}
  ~MPIReduceOp() {}

  bool RunOnDevice() override {
    MPI_Comm comm = OperatorBase::Input<MPICommonWorldWrapper>(0).comm();
    auto& input = Input(1);
    auto* output = Output(0, input.sizes(), at::dtype<T>());
    MPI_CHECK(MPI_Reduce(
        const_cast<T*>(input.template data<T>()),
        output->template mutable_data<T>(),
        input.numel(),
        MPIDataTypeWrapper<T>::type(),
        MPI_SUM,
        root_,
        comm));
    return true;
  }

 protected:
  int root_;
};

// MPIAllgatherOp does MPIAllgather using MPI.
template <typename T, class Context>
class MPIAllgatherOp final : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;
  USE_SIMPLE_CTOR_DTOR(MPIAllgatherOp);

  bool RunOnDevice() override {
    MPI_Comm comm = OperatorBase::Input<MPICommonWorldWrapper>(0).comm();
    auto& input = Input(1);
    auto* output = Output(0);
    vector<int64_t> output_dims = input.sizes().vec();
    output_dims[0] *= OperatorBase::Input<MPICommonWorldWrapper>(0).size();
    output->Resize(output_dims);
    MPI_CHECK(MPI_Allgather(
        const_cast<T*>(input.template data<T>()),
        input.numel(),
        MPIDataTypeWrapper<T>::type(),
        output->template mutable_data<T>(),
        input.numel(),
        MPIDataTypeWrapper<T>::type(),
        comm));
    return true;
  }
};

// MPIAllreduceOp does MPIAllreduce using MPI. Currently, only SUM is supported.
template <typename T, class Context>
class MPIAllreduceOp final : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;
  USE_SIMPLE_CTOR_DTOR(MPIAllreduceOp);

  bool RunOnDevice() override {
    MPI_Comm comm = OperatorBase::Input<MPICommonWorldWrapper>(0).comm();
    auto& input = Input(1);
    auto* output = Output(0, input.sizes(), at::dtype<T>());
    void* source;
    if (output->template mutable_data<T>() == input.template data<T>()) {
      // We are doing in-place call. Special case handling.
      source = MPI_IN_PLACE;
    } else {
      // Normal allreduce takes the source from the input.
      source = const_cast<T*>(input.template data<T>());
    }
    MPI_CHECK(MPI_Allreduce(
        source,
        output->template mutable_data<T>(),
        input.numel(),
        MPIDataTypeWrapper<T>::type(),
        MPI_SUM,
        comm));
    return true;
  }
};

template <class Context>
class MPISendTensorOp final : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;
  MPISendTensorOp(const OperatorDef& def, Workspace* ws)
      : Operator<Context>(def, ws),
        OP_SINGLE_ARG(int, "dst", dst_, MPI_ANY_SOURCE),
        OP_SINGLE_ARG(int, "tag", tag_, MPI_ANY_TAG),
        OP_SINGLE_ARG(bool, "raw_buffer", raw_buffer_, false) {
    CAFFE_ENFORCE(raw_buffer_, "non-raw-buffer transfer not supported yet.");
    CAFFE_ENFORCE(
        dst_ != MPI_ANY_SOURCE || def.input_size() == 4,
        "You should explicitly specify the to rank either via "
        "argument or via input blobs.");
    CAFFE_ENFORCE(
        tag_ != MPI_ANY_TAG || def.input_size() == 4,
        "You should explicitly specify the tag either via "
        "argument or via input blobs.");
  }

  bool RunOnDevice() override {
    MPI_Comm comm = OperatorBase::Input<MPICommonWorldWrapper>(COMM).comm();
    auto& input = Input(INPUT);
    if (InputSize() == 4) {
      dst_ = OperatorBase::Input<Tensor>(DST, CPU).template data<int>()[0];
      tag_ = OperatorBase::Input<Tensor>(TAG, CPU).template data<int>()[0];
    }
    if (raw_buffer_) {
      // We need to do a const cast to cope with the fact that, before OpenMPI
      // 1.7, MPI_Send expects a non-const pointer although it uses it in a
      // const way.
      MPI_CHECK(MPI_Send(
          const_cast<void*>(input.raw_data()),
          input.nbytes(),
          MPI_CHAR,
          dst_,
          tag_,
          comm));
    } else {
      CAFFE_NOT_IMPLEMENTED;
    }
    return true;
  }

 protected:
  int dst_;
  int tag_;
  bool raw_buffer_;

  INPUT_TAGS(COMM, INPUT, DST, TAG);
};

template <class Context>
class MPIReceiveTensorOp final : public Operator<Context> {
 public:
  USE_OPERATOR_CONTEXT_FUNCTIONS;
  MPIReceiveTensorOp(const OperatorDef& def, Workspace* ws)
      : Operator<Context>(def, ws),
        OP_SINGLE_ARG(int, "src", src_, MPI_ANY_SOURCE),
        OP_SINGLE_ARG(int, "tag", tag_, MPI_ANY_TAG),
        OP_SINGLE_ARG(bool, "raw_buffer", raw_buffer_, false) {
    CAFFE_ENFORCE(raw_buffer_, "non-raw-buffer transfer not supported yet.");
  }

  bool RunOnDevice() override {
    MPI_Comm comm = OperatorBase::Input<MPICommonWorldWrapper>(COMM).comm();
    if (InputSize() == 4) {
      src_ = OperatorBase::Input<Tensor>(SRC_IN, CPU).template data<int>()[0];
      tag_ = OperatorBase::Input<Tensor>(TAG_IN, CPU).template data<int>()[0];
    }
    MPI_Status status;
    if (raw_buffer_) {
      auto* output = Output(OUTPUT);
      MPI_CHECK(MPI_Recv(
          output->raw_mutable_data(),
          output->nbytes(),
          MPI_CHAR,
          src_,
          tag_,
          comm,
          &status));
    } else {
      CAFFE_NOT_IMPLEMENTED;
    }
    auto* src_out = OperatorBase::Output<Tensor>(SRC_OUT, CPU);
    src_out->Resize();
    src_out->template mutable_data<int>()[0] = status.MPI_SOURCE;
    auto* tag_out = OperatorBase::Output<Tensor>(TAG_OUT, CPU);
    tag_out->Resize();
    tag_out->template mutable_data<int>()[0] = status.MPI_TAG;
    return true;
  }

 protected:
  int src_;
  int tag_;
  bool raw_buffer_;
  INPUT_TAGS(COMM, INPUT, SRC_IN, TAG_IN);
  OUTPUT_TAGS(OUTPUT, SRC_OUT, TAG_OUT);
};

} // namespace caffe2

#endif // CAFFE2_MPI_MPI_OPS_H_