//============================================================================
//  Copyright (c) Kitware, Inc.
//  All rights reserved.
//  See LICENSE.txt for details.
//
//  This software is distributed WITHOUT ANY WARRANTY; without even
//  the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
//  PURPOSE.  See the above copyright notice for more information.
//============================================================================

#ifndef vtk_m_filter_flow_internal_Messenger_h
#define vtk_m_filter_flow_internal_Messenger_h

#include <vtkm/Particle.h> //DRP. Added for instrumentation defines...
#include <vtkm/filter/flow/Tracer.h>


#include <vtkm/Types.h>
#include <vtkm/filter/flow/vtkm_filter_flow_export.h>
#include <vtkm/thirdparty/diy/diy.h>

#include <list>
#include <map>
#include <set>
#include <vector>

#ifdef VTKM_ENABLE_MPI
#include <mpi.h>
#endif

namespace vtkm
{
namespace filter
{
namespace flow
{
namespace internal
{

class VTKM_FILTER_FLOW_EXPORT Messenger
{
public:
enum CommunicationType {SYNC, ASYNC, ASYNC_PROBE};
struct AsyncCommParamType
{
  int MessageSize = 1;
  int NumberOfReceivers = 64;
  int NumberOfParticles = 128;
  int NumberOfBlockIds = 2;
};

public:
  VTKM_CONT Messenger(vtkmdiy::mpi::communicator& comm);
  VTKM_CONT virtual ~Messenger()
  {
#ifdef VTKM_ENABLE_MPI
    this->CleanupRequests();
#endif
  }

  int GetRank() const { return this->Rank; }
  int GetNumRanks() const { return this->NumRanks; }

#ifdef VTKM_ENABLE_MPI
  VTKM_CONT void RegisterTag(int tag, std::size_t numRecvs, std::size_t size);

  bool UsingSyncCommunication() const { return this->CommType == SYNC; }
  bool UsingAsyncCommunication() const { return this->CommType == ASYNC; }
  bool UsingAsyncProbeCommunication() const { return this->CommType == ASYNC_PROBE; }

protected:
  static std::size_t CalcMessageBufferSize(std::size_t msgSz);

  void InitializeBuffers();
  void CheckPendingSendRequests();
  void CleanupRequests(int tag = TAG_ANY);
  void SendData(int dst, int tag, vtkmdiy::MemoryBuffer& buff)
  {
    if (this->UsingSyncCommunication())
      this->SendDataSync(dst, tag, buff);
    else if (this->UsingAsyncCommunication() || this->UsingAsyncProbeCommunication())
      this->SendDataAsync(dst, tag, buff);
    else
      throw vtkm::cont::ErrorFilterExecution("Unsupported communication method");
  }

  bool RecvData(const std::set<int>& tags,
                std::vector<std::pair<int, vtkmdiy::MemoryBuffer>>& buffers,
                std::vector<std::pair<int,int>>& bufferRecvTimeWindows, //DRP
                bool blockAndWait = false)
  {
    if (this->UsingSyncCommunication())
      return this->RecvDataSync(tags, buffers, blockAndWait);
    else if (this->UsingAsyncCommunication())
      return this->RecvDataAsync(tags, buffers, bufferRecvTimeWindows, blockAndWait);
    else if (this->UsingAsyncProbeCommunication())
      return this->RecvDataAsyncProbe(tags, buffers, blockAndWait);

    throw vtkm::cont::ErrorFilterExecution("Unsupported communication method");
    return false;
  }

  Messenger::CommunicationType CommType = SYNC;
  Messenger::AsyncCommParamType AsyncCommParams;

private:
  void SendDataAsync(int dst, int tag, const vtkmdiy::MemoryBuffer& buff);
  void SendDataSync(int dst, int tag, vtkmdiy::MemoryBuffer& buff);
  bool RecvDataAsync(const std::set<int>& tags,
                     std::vector<std::pair<int, vtkmdiy::MemoryBuffer>>& buffers,
                     std::vector<std::pair<int,int>>& bufferRecvTimeWindows, //DRP
                     bool blockAndWait);
  bool RecvDataAsyncProbe(const std::set<int>& tags,
                          std::vector<std::pair<int, vtkmdiy::MemoryBuffer>>& buffers,
                          bool blockAndWait);
  bool RecvDataSync(const std::set<int>& tags,
                    std::vector<std::pair<int, vtkmdiy::MemoryBuffer>>& buffers,
                    bool blockAndWait);
  void PostRecv(int tag);
  void PostRecv(int tag, std::size_t sz, int src = -1);


  //Message headers.
  typedef struct
  {
    int rank, tag;
    std::size_t id, numPackets, packet, packetSz, dataSz;
  } Header;

  void PrepareForSend(int tag, const vtkmdiy::MemoryBuffer& buff, std::vector<char*>& buffList);
  vtkm::Id GetMsgID() { return this->MsgID++; }
  static bool PacketCompare(const char* a, const char* b);
  void ProcessReceivedBuffers(std::vector<char*>& incomingBuffers,
                              std::vector<std::pair<int, vtkmdiy::MemoryBuffer>>& buffers,
                              std::vector<std::pair<int,int>>& bufferRecvTimeWindows);

  // Send/Recv buffer management structures.
  using RequestTagPair = std::pair<MPI_Request, int>;
  using RankIdPair = std::pair<int, int>;

  //Member data
  // <tag, {dst, buffer}>
  std::map<int, std::vector<std::pair<int, vtkmdiy::MemoryBuffer>>> SyncSendBuffers;
  std::map<int, std::pair<std::size_t, std::size_t>> MessageTagInfo;
  MPI_Comm MPIComm;
  std::size_t MsgID;
  int NumRanks;
  int Rank;
  std::map<RequestTagPair, char*> RecvBuffers;
  std::map<RankIdPair, std::list<char*>> RecvPackets;
#ifdef VTKm_INSTRUMENT_PARTICLE_ADVECTION
  std::map<RankIdPair, std::pair<int,int>> RecvPacketTimes;
#endif
  std::map<RequestTagPair, char*> SendBuffers;
  static constexpr int TAG_ANY = -1;

  vtkm::Id NumberOfReceivers=0;

  void CheckRequests(const std::map<RequestTagPair, char*>& buffer,
                     const std::set<int>& tags,
                     bool BlockAndWait,
                     std::vector<RequestTagPair>& reqTags);

  int TestSomeRequests(std::vector<MPI_Request>& requests,
                       std::vector<int>& indices,
                       int& num) const;
  int TestAllRequests(std::vector<MPI_Request>& requests,
                      std::vector<int>& indices,
                      int& num) const;


#else
protected:
  static constexpr int NumRanks = 1;
  static constexpr int Rank = 0;
#endif
};


template <typename T>
std::ostream& operator<<(std::ostream& os, const std::vector<T>& v)
{
  os << "[";
  for (std::size_t i = 0; i < v.size(); ++i)
  {
    os << v[i];
    if (i != v.size() - 1)
      os << ", ";
  }
  os << "]";
  return os;
}

}
}
}
} // vtkm::filter::flow::internal

#endif // vtk_m_filter_flow_internal_Messenger_h
