#include "moab/ParallelData.hpp"
#include "moab/ParallelComm.hpp"
#include "MBParallelConventions.h"
#include "moab/Interface.hpp"

#include <map>

namespace moab {

    //! return partition sets; if tag_name is input, gets sets with
    //! that tag name, otherwise uses PARALLEL_PARTITION tag
ErrorCode ParallelData::get_partition_sets(Range &part_sets,
                                               const char *tag_name) 
{
  Tag part_tag = 0;
  ErrorCode result;
  
  if (NULL != tag_name) 
    result = mbImpl->tag_get_handle(tag_name, 1, MB_TYPE_INTEGER, part_tag);
  else
    result = mbImpl->tag_get_handle(PARALLEL_PARTITION_TAG_NAME, 1, MB_TYPE_INTEGER, part_tag);
    
  if (MB_SUCCESS != result) return result;
  else if (0 == part_tag) return MB_TAG_NOT_FOUND;
  
  result = mbImpl->get_entities_by_type_and_tag(0, MBENTITYSET, &part_tag, 
                                                NULL, 1, part_sets,
                                                Interface::UNION);
  return result;
}
  

    //! get communication interface sets and the processors with which
    //! this processor communicates; sets are sorted by processor
ErrorCode ParallelData::get_interface_sets(std::vector<EntityHandle> &iface_sets,
                                               std::vector<int> &iface_procs) 
{
#define CONTINUE {result = tmp_result; continue;}
  iface_sets.clear();
  iface_procs.clear();
  
  Tag proc_tag = 0, procs_tag = 0;
  ErrorCode result = MB_SUCCESS;
  int my_rank;
  if (parallelComm) 
    my_rank = parallelComm->proc_config().proc_rank();
  else 
    return MB_FAILURE;

  std::multimap<int, EntityHandle> iface_data;

  for (int i = 0; i < 2; i++) {
    ErrorCode tmp_result;
    
    if (0 == i)
      tmp_result = mbImpl->tag_get_handle(PARALLEL_SHARED_PROC_TAG_NAME, 
                                      1, MB_TYPE_INTEGER, proc_tag);
    else
      tmp_result = mbImpl->tag_get_handle(PARALLEL_SHARED_PROCS_TAG_NAME, 
                                      MAX_SHARING_PROCS, MB_TYPE_INTEGER, proc_tag);
    if (MB_SUCCESS != tmp_result) CONTINUE;

    int tsize;
    tmp_result = mbImpl->tag_get_length(proc_tag, tsize);
    if (0 == tsize || MB_SUCCESS != tmp_result) CONTINUE;
    
    Range proc_sets;
    tmp_result = mbImpl->get_entities_by_type_and_tag(0, MBENTITYSET, 
                                                  &proc_tag, NULL, 1,
                                                  proc_sets, Interface::UNION);
    if (MB_SUCCESS != tmp_result) CONTINUE;
    
    if (proc_sets.empty()) CONTINUE;
      
    std::vector<int> proc_tags(proc_sets.size()*tsize);
    tmp_result = mbImpl->tag_get_data(procs_tag, proc_sets, &proc_tags[0]);
    if (MB_SUCCESS != tmp_result) CONTINUE;
    int k;
    Range::iterator rit;
    
    for (k = 0, rit = proc_sets.begin(); rit != proc_sets.end(); ++rit, k++) {
      for (int j = 0; j < tsize; j++) {
        if (my_rank != proc_tags[2*k+j] && proc_tags[2*k+j] >= 0)
          iface_data.insert(std::pair<int,EntityHandle>(proc_tags[2*k+j], *rit));
      }
    }
  }

    // now get the results in sorted order
  std::multimap<int,EntityHandle>::iterator mit;
  for (mit = iface_data.begin(); mit != iface_data.end(); ++mit)
    iface_procs.push_back((*mit).first),
      iface_sets.push_back((*mit).second);
    
  return result;
}


} // namespace moab

