#ifndef CONNECTIVITY_DISPATCHER_HPP
#define CONNECTIVITY_DISPATCHER_HPP

#include <Mesh.hpp>
#include <ItemValue.hpp>
#include <ItemValueUtils.hpp>

#include <unordered_map>
#include <ConnectivityDescriptor.hpp>

template <int Dimension>
class ConnectivityDispatcher
{
 public:
  using ConnectivityType = Connectivity<Dimension>;

 private:
  const ConnectivityType& m_connectivity;
  ConnectivityDescriptor m_new_descriptor;
  std::shared_ptr<ConnectivityType> m_dispatched_connectivity;

  template <ItemType item_type>
  struct DispatchedItemInfo
  {
    using ItemId = ItemIdT<item_type>;
    ItemValue<const int, item_type> m_new_owner;
    Array<const unsigned int> m_list_to_send_size_by_proc;
    std::vector<Array<const ItemId>> m_list_to_send_by_proc;
    Array<const unsigned int> m_list_to_recv_size_by_proc;
    std::unordered_map<int, int> m_number_to_id_map;
    std::vector<Array<const ItemId>> m_recv_id_correspondance_by_proc;
  };

  DispatchedItemInfo<ItemType::cell> m_dispatched_cell_info;
  DispatchedItemInfo<ItemType::face> m_dispatched_face_info;
  DispatchedItemInfo<ItemType::edge> m_dispatched_edge_info;
  DispatchedItemInfo<ItemType::node> m_dispatched_node_info;

  template <ItemType item_type>
  PASTIS_INLINE
  DispatchedItemInfo<item_type>& _dispatchedInfo()
  {
    if constexpr (item_type == ItemType::cell) {
      return m_dispatched_cell_info;
    } else if constexpr (item_type == ItemType::face) {
      return m_dispatched_face_info;
    } else if constexpr (item_type == ItemType::edge) {
      return m_dispatched_edge_info;
    } else {
      return m_dispatched_node_info;
    }
  }

  template <ItemType item_type>
  PASTIS_INLINE
  const DispatchedItemInfo<item_type>& _dispatchedInfo() const
  {
    if constexpr (item_type == ItemType::cell) {
      return m_dispatched_cell_info;
    } else if constexpr (item_type == ItemType::face) {
      return m_dispatched_face_info;
    } else if constexpr (item_type == ItemType::edge) {
      return m_dispatched_edge_info;
    } else {
      return m_dispatched_node_info;
    }
  }

  template <typename ItemToItem>
  struct DispatchedItemOfItemInfo
  {
    std::vector<Array<const int>> m_number_of_sub_item_per_item_to_recv_by_proc;
    std::vector<Array<const int>> m_sub_item_numbers_to_recv_by_proc;
  };

  DispatchedItemOfItemInfo<NodeOfCell> m_dispatched_node_of_cell_info;
  DispatchedItemOfItemInfo<EdgeOfCell> m_dispatched_edge_of_cell_info;
  DispatchedItemOfItemInfo<FaceOfCell> m_dispatched_face_of_cell_info;

  DispatchedItemOfItemInfo<NodeOfEdge> m_dispatched_node_of_edge_info;
  DispatchedItemOfItemInfo<FaceOfEdge> m_dispatched_face_of_edge_info;
  DispatchedItemOfItemInfo<CellOfEdge> m_dispatched_cell_of_edge_info;

  DispatchedItemOfItemInfo<NodeOfFace> m_dispatched_node_of_face_info;
  DispatchedItemOfItemInfo<EdgeOfFace> m_dispatched_edge_of_face_info;
  DispatchedItemOfItemInfo<CellOfFace> m_dispatched_cell_of_face_info;

  DispatchedItemOfItemInfo<EdgeOfNode> m_dispatched_edge_of_node_info;
  DispatchedItemOfItemInfo<FaceOfNode> m_dispatched_face_of_node_info;
  DispatchedItemOfItemInfo<CellOfNode> m_dispatched_cell_of_node_info;

  template <typename ItemOfItem>
  PASTIS_INLINE
  DispatchedItemOfItemInfo<ItemOfItem>& _dispatchedInfo()
  {
    if constexpr (std::is_same_v<NodeOfCell, ItemOfItem>) {
      return m_dispatched_node_of_cell_info;
    } else if constexpr (std::is_same_v<EdgeOfCell, ItemOfItem>) {
      return m_dispatched_edge_of_cell_info;
    } else if constexpr (std::is_same_v<FaceOfCell, ItemOfItem>) {
      return m_dispatched_face_of_cell_info;
    } else if constexpr (std::is_same_v<NodeOfEdge, ItemOfItem>) {
      return m_dispatched_node_of_edge_info;
    } else if constexpr (std::is_same_v<FaceOfEdge, ItemOfItem>) {
      return m_dispatched_face_of_edge_info;
    } else if constexpr (std::is_same_v<CellOfEdge, ItemOfItem>) {
      return m_dispatched_cell_of_edge_info;
    } else if constexpr (std::is_same_v<NodeOfFace, ItemOfItem>) {
      return m_dispatched_node_of_face_info;
    } else if constexpr (std::is_same_v<EdgeOfFace, ItemOfItem>) {
      return m_dispatched_edge_of_face_info;
    } else if constexpr (std::is_same_v<CellOfFace, ItemOfItem>) {
      return m_dispatched_cell_of_face_info;
    } else if constexpr (std::is_same_v<EdgeOfNode, ItemOfItem>) {
      return m_dispatched_edge_of_node_info;
    } else if constexpr (std::is_same_v<FaceOfNode, ItemOfItem>) {
      return m_dispatched_face_of_node_info;
    } else if constexpr (std::is_same_v<CellOfNode, ItemOfItem>) {
      return m_dispatched_cell_of_node_info;
    } else {
      static_assert(is_false_v<ItemOfItem>, "Unexpected ItemOfItem type");
    }
  }

  template <typename ItemOfItem>
  PASTIS_INLINE
  const DispatchedItemOfItemInfo<ItemOfItem>& _dispatchedInfo() const
  {
    if constexpr (std::is_same_v<NodeOfCell, ItemOfItem>) {
      return m_dispatched_node_of_cell_info;
    } else if constexpr (std::is_same_v<EdgeOfCell, ItemOfItem>) {
      return m_dispatched_edge_of_cell_info;
    } else if constexpr (std::is_same_v<FaceOfCell, ItemOfItem>) {
      return m_dispatched_face_of_cell_info;
    } else if constexpr (std::is_same_v<NodeOfEdge, ItemOfItem>) {
      return m_dispatched_node_of_edge_info;
    } else if constexpr (std::is_same_v<FaceOfEdge, ItemOfItem>) {
      return m_dispatched_face_of_edge_info;
    } else if constexpr (std::is_same_v<CellOfEdge, ItemOfItem>) {
      return m_dispatched_cell_of_edge_info;
    } else if constexpr (std::is_same_v<NodeOfFace, ItemOfItem>) {
      return m_dispatched_node_of_face_info;
    } else if constexpr (std::is_same_v<EdgeOfFace, ItemOfItem>) {
      return m_dispatched_edge_of_face_info;
    } else if constexpr (std::is_same_v<CellOfFace, ItemOfItem>) {
      return m_dispatched_cell_of_face_info;
    } else if constexpr (std::is_same_v<EdgeOfNode, ItemOfItem>) {
      return m_dispatched_edge_of_node_info;
    } else if constexpr (std::is_same_v<FaceOfNode, ItemOfItem>) {
      return m_dispatched_face_of_node_info;
    } else if constexpr (std::is_same_v<CellOfNode, ItemOfItem>) {
      return m_dispatched_cell_of_node_info;
    } else {
      static_assert(is_false_v<ItemOfItem>, "Unexpected ItemOfItem type");
    }
  }

  template <ItemType item_type>
  void _buildNewOwner();

  template <ItemType item_type>
  void _buildItemListToSend();

  void _buildCellNumberIdMap();

  template <typename ItemOfItemT>
  void _buildSubItemNumberToIdMap();

  template <ItemType item_type>
  void _buildItemToExchangeLists();

  template <ItemType item_type>
  void _buildNumberOfItemToExchange();

  template <typename ItemOfItemT>
  void _buildItemToSubItemDescriptor();

  void _dispatchEdges();
  void _dispatchFaces();

  template<typename DataType, ItemType item_type, typename ConnectivityPtr>
  void _gatherFrom(const ItemValue<DataType, item_type, ConnectivityPtr>& data_to_gather,
                   std::vector<std::remove_const_t<DataType>>& gathered_vector);

  template<typename DataType, typename ItemOfItem, typename ConnectivityPtr>
  void _gatherFrom(const SubItemValuePerItem<DataType, ItemOfItem, ConnectivityPtr>& data_to_gather,
                   std::vector<Array<std::remove_const_t<DataType>>>& gathered_vector);

  template <typename SubItemOfItemT>
  void _buildNumberOfSubItemPerItemToRecvByProc();

  template <typename SubItemOfItemT>
  void _buildSubItemNumbersToRecvByProc();

  template <ItemType item_type>
  void _buildRecvItemIdCorrespondanceByProc();

  template <ItemType item_type>
  void _buildItemReferenceList();

 public:
  std::shared_ptr<const ConnectivityType>
  dispatchedConnectivity() const
  {
    return m_dispatched_connectivity;
  }

  template<typename DataType, ItemType item_type, typename ConnectivityPtr>
  std::vector<Array<const DataType>>
  exchange(ItemValue<DataType, item_type, ConnectivityPtr> item_value) const
  {
    using ItemId = ItemIdT<item_type>;
    using MutableDataType = std::remove_const_t<DataType>;
    std::vector<Array<const DataType>> item_value_to_send_by_proc(parallel::size());

    const auto& item_list_to_send_by_proc = this->_dispatchedInfo<item_type>().m_list_to_send_by_proc;

    for (size_t i=0; i<parallel::size(); ++i) {
      const Array<const ItemId>& item_list = item_list_to_send_by_proc[i];
      Array<MutableDataType> item_value_list(item_list.size());
      parallel_for (item_list.size(), PASTIS_LAMBDA(const ItemId& item_id) {
          item_value_list[item_id] = item_value[item_list[item_id]];
        });
      item_value_to_send_by_proc[i] = item_value_list;
    }

    std::vector<Array<MutableDataType>> recv_item_value_by_proc(parallel::size());
    {
      const auto& list_to_recv_size_by_proc = this->_dispatchedInfo<item_type>().m_list_to_recv_size_by_proc;
      for (size_t i=0; i<parallel::size(); ++i) {
        recv_item_value_by_proc[i] = Array<MutableDataType>(list_to_recv_size_by_proc[i]);
      }
    }
    parallel::exchange(item_value_to_send_by_proc, recv_item_value_by_proc);

    std::vector<Array<const DataType>> const_recv_item_value_by_proc(parallel::size());
    for (size_t i=0; i<parallel::size(); ++i) {
      const_recv_item_value_by_proc[i] = recv_item_value_by_proc[i];
    }

    return const_recv_item_value_by_proc;
  }

  template<typename DataType, ItemType item_type, typename ConnectivityPtr>
  ItemValue<std::remove_const_t<DataType>, item_type, ConnectivityPtr>
  dispatch(ItemValue<DataType, item_type, ConnectivityPtr> item_value) const
  {
    using ItemId = ItemIdT<item_type>;

    Assert(m_dispatched_connectivity.use_count()> 0,
           "cannot dispatch quantity before connectivity");

    const auto& item_list_to_send_by_proc = this->_dispatchedInfo<item_type>().m_list_to_send_by_proc;

    using MutableDataType = std::remove_const_t<DataType>;
    std::vector<Array<DataType>> item_value_to_send_by_proc(parallel::size());
    for (size_t i=0; i<parallel::size(); ++i) {
      const Array<const ItemId>& item_list = item_list_to_send_by_proc[i];
      Array<MutableDataType> item_value_list(item_list.size());
      parallel_for (item_list.size(), PASTIS_LAMBDA(const ItemId& item_id) {
          item_value_list[item_id] = item_value[item_list[item_id]];
        });
      item_value_to_send_by_proc[i] = item_value_list;
    }

    const auto& item_list_to_recv_size_by_proc = this->_dispatchedInfo<item_type>().m_list_to_recv_size_by_proc;
    std::vector<Array<MutableDataType>> recv_item_value_by_proc(parallel::size());
    for (size_t i=0; i<parallel::size(); ++i) {
      recv_item_value_by_proc[i] = Array<MutableDataType>(item_list_to_recv_size_by_proc[i]);
    }

    parallel::exchange(item_value_to_send_by_proc, recv_item_value_by_proc);

    const auto& recv_item_id_correspondance_by_proc =
        this->_dispatchedInfo<item_type>().m_recv_id_correspondance_by_proc;
    ItemValue<MutableDataType, item_type> new_item_value(*m_dispatched_connectivity);
    for (size_t i_rank=0; i_rank<parallel::size(); ++i_rank) {
      const auto& recv_item_id_correspondance = recv_item_id_correspondance_by_proc[i_rank];
      const auto& recv_item_value = recv_item_value_by_proc[i_rank];
      parallel_for(recv_item_value.size(), PASTIS_LAMBDA(size_t r) {
          const ItemId& item_id = recv_item_id_correspondance[r];
          new_item_value[item_id] = recv_item_value[r];
        });
    }
    return new_item_value;
  }

  ConnectivityDispatcher(const ConnectivityType& mesh);
  ConnectivityDispatcher(const ConnectivityDispatcher&) = delete;
  ~ConnectivityDispatcher() = default;
};


#endif // CONNECTIVITY_DISPATCHER_HPP