#include "moab/Core.hpp"
#include "moab/Range.hpp"
#include "TestUtil.hpp"
#include "MBTagConventions.hpp"
#include "moab/ParallelComm.hpp"
#include "moab/EntityType.hpp"
#include <stdlib.h>
#include <algorithm>
#include <mpi.h>

using namespace moab;

void test_create_tag();
void test_global_id(Interface& mb,
                    ParallelComm *pcomm,
                    Range &orig_ents);

int main( int argc, char* argv[] )
{
  int failures = 0;
  MPI_Init(&argc, &argv);
  
  failures += RUN_TEST( test_create_tag );
  
  if (failures) 
    std::cerr << "<<<< " << failures << " TESTS FAILED >>>>" << std::endl;
    
  return failures;
}

void setup_mesh( Interface& mesh, ParallelComm *pcomm );

void test_create_tag()
{ 
  Core mb;
   
  int rank,size ;
  MPI_Comm_rank( MPI_COMM_WORLD, &rank );
  MPI_Comm_size(MPI_COMM_WORLD, &size);

  ParallelComm *pcomm = new ParallelComm( &mb, MPI_COMM_WORLD   );

  setup_mesh(mb, pcomm);

  Range range_elem; 
  ErrorCode rval;
  rval = mb.get_entities_by_type( 0, MBQUAD, range_elem );
  CHECK_ERR(rval);

  test_global_id(mb, pcomm, range_elem);
}

void test_global_id(Interface& mb,
                    ParallelComm *pcomm,
                    Range &orig_ents)
{
  ErrorCode result = MB_SUCCESS;
  Range all_dim[4];
  Range::const_iterator orig_it = orig_ents.begin();
  for(; orig_it != orig_ents.end(); orig_it++)
  {
    for (int i = 0; i < 4; i++)
        result = mb.get_adjacencies(&(*orig_it), 1, i, true, all_dim[i],
                 Interface::UNION);
  }

  //since I don't have the interfaceSet set, I'll just call to get owned-only
  //global ids for vertices
  result = pcomm->assign_global_ids(all_dim, 0, 1, true, true); 

  //check for the result.
  int v_gid[4];
  Tag gid_tag;
  mb.tag_get_handle(GLOBAL_ID_TAG_NAME, 1, MB_TYPE_INTEGER, gid_tag);

  for (int i = 0; i < 4; i++)
  {
    EntityHandle h_vt = all_dim[0][i];
    mb.tag_get_data(gid_tag, &h_vt, 1,  &v_gid[i]);
  }

  int rank ;
  MPI_Comm_rank( MPI_COMM_WORLD, &rank );
  std::cout << "rank = " << rank << ",  vertice's global ids = " << v_gid[0] << "\t" << v_gid[1] << "\t" <<
  v_gid[2] << "\t" << v_gid[3] << "\n";

  //output is
//rank = 0,  vertice's global ids = 1   2       3       4
//rank = 2,  vertice's global ids = 9	10	0	9   <--only first two
//rank = 3,  vertice's global ids = 13	14	0	13  <--are valid.
//rank = 1,  vertice's global ids = 5	6	0	5   
}
void setup_mesh( Interface& mb, ParallelComm *pcomm )
{
  //              Mesh Ids                                                                                                                  
  //   2.0   5 ------- 6 -------- 7 -------- 8 --------- 9                                         
  //          |        |          |          |            |                      
  //          |    1   |    2     |     3    |      4     |                          
  //          |        |          |          |            |                                   
  //   1.0   0 ------- 1 -------- 2 -------- 3 --------- 4  
  //        0.0       1.0        2.0        3.0          4.0                                                                                      
  //                                                                                                                                        
  //      Node Ids at corners                                                                                                               
  //      Element Ids in centers                                                                                                            
  //                                                                                                                                        
  //             Mesh Owners                                                                                                                
  //                                                                                                                                        
  //   2.0   0 ------- 0 -------- 1 -------- 2 ---------- 3                      
  //          |        |          |          |            |                      
  //          |    0   |    1     |     2    |      3     |                      
  //          |        |          |          |            |                      
  //   1.0   0 ------- 0 -------- 1 -------- 2 -----------3

  //                                                                                                                                        
  //       0.0       1.0        2.0                                                                                                         
  //                                                                                                                                        
  //      Node Owners at corners                                                                                                            
  //      Element Owners in centers    
  Range vertex_handles;
  const double vertex_coords[] = { 0, 0, 0,  1, 0, 0,  2, 0, 0,
                                   3, 0, 0,  4, 0, 0,  0, 1, 0,
                                   1, 1, 0,  2, 1, 0,  3, 1, 0,  4,1,0 };
                                   
  const unsigned num_vtx = sizeof(vertex_coords)/(3*sizeof(double));
  
  ErrorCode rval = mb.create_vertices( vertex_coords, num_vtx, vertex_handles );
  CHECK_ERR(rval);
  CHECK_EQUAL( num_vtx, (unsigned)vertex_handles.size() );
  
  EntityHandle elements;
  EntityHandle conn[4][4] = { 
    {  vertex_handles[0],  vertex_handles[1],  vertex_handles[6],  vertex_handles[5] },   
    {vertex_handles[1] , vertex_handles[2], vertex_handles[7], vertex_handles[6] },
    {vertex_handles[2],  vertex_handles[3],  vertex_handles[8],  vertex_handles[7] },   
    {vertex_handles[3], vertex_handles[4], vertex_handles[9], vertex_handles[8] }}; 
  int rank;
  MPI_Comm_rank( MPI_COMM_WORLD, &rank );
  rval = mb.create_element( MBQUAD, conn[rank], 4, elements );
  CHECK_ERR(rval);

  Range edges;
  mb.get_adjacencies(&elements, 1, 1, true, edges, Interface::UNION);
  std::cout << "edges.size() = " << edges.size() << "\n";
  std::vector<EntityHandle> h_edges;
  h_edges.resize(edges.size());
  Range::iterator it = edges.begin();
  int j = 0;
  for(; it != edges.end(); it++, j++)
    h_edges[j] = *it;

  Tag proc_tag;
  mb.tag_get_handle(PARALLEL_TO_PROCESSOR_NAME, 1, MB_TYPE_INTEGER, proc_tag, MB_TAG_EXCL|MB_TAG_SPARSE);

  Tag shared_tag = pcomm->pstatus_tag();

  Tag shared_proc = pcomm->sharedp_tag();

  Tag shared_handle = pcomm->sharedh_tag();

  //Range::iterator it = vertex_handles.begin();
  int val = 3, pstat = 0, proc = 1;
  switch(rank)
  {
    case 0:
      mb.tag_set_data(proc_tag, &elements, 1, &val);
      mb.tag_set_data(proc_tag, &conn[0][0], 4, &val);
      mb.tag_set_data(shared_tag, &conn[0][0], 4, &pstat);
      mb.tag_set_data(shared_tag, &h_edges[0], 4, &pstat);
      mb.tag_set_data(shared_proc, &conn[0][1], 2, &proc);
      mb.tag_set_data(shared_handle, &conn[0][1], 1, &conn[1][0]);
      mb.tag_set_data(shared_handle, &conn[0][2], 1, &conn[1][3]);
      mb.tag_set_data(shared_proc, &h_edges[1], 1, &proc);
      
      break;
    case 1:
      mb.tag_set_data(proc_tag, &elements, 1, &val);
      mb.tag_set_data(proc_tag, &conn[1][0], 4, &val); 

      pstat = PSTATUS_NOT_OWNED;
      mb.tag_set_data(shared_tag, &h_edges[3], 1, &pstat);
      mb.tag_set_data(shared_tag, &conn[1][0], 1, &pstat);
      mb.tag_set_data(shared_tag, &conn[1][3], 1, &pstat);
      
      pstat = 0;
      mb.tag_set_data(shared_tag, &conn[1][1], 1, &pstat);
      mb.tag_set_data(shared_tag, &conn[1][2], 1, &pstat);
      mb.tag_set_data(shared_tag, &h_edges[0], 3, &pstat);

      proc = 2;
      mb.tag_set_data(shared_proc, &conn[1][1], 1, &proc);
      mb.tag_set_data(shared_proc, &conn[1][2], 1, &proc);
      mb.tag_set_data(shared_proc, &h_edges[1], 1, &proc);
      proc = 0;
      mb.tag_set_data(shared_proc, &h_edges[3], 1, &proc);
      mb.tag_set_data(shared_proc, &conn[1][0], 1, &proc);
      mb.tag_set_data(shared_proc, &conn[1][3], 1, &proc);
 
      mb.tag_set_data(shared_handle,&conn[1][0], 1, &conn[0][1]);
      mb.tag_set_data(shared_handle,&conn[1][1], 1, &conn[2][0]);
      mb.tag_set_data(shared_handle,&conn[1][2], 1, &conn[2][3]);
      mb.tag_set_data(shared_handle,&conn[1][3], 1, &conn[0][2]); 
      break;
    case 2:
      val = 1;
      mb.tag_set_data(proc_tag, &elements, 1, &val);
      mb.tag_set_data(proc_tag, &conn[2][0], 1, &val);
      mb.tag_set_data(proc_tag, &conn[2][3], 1, &val);
      pstat = PSTATUS_NOT_OWNED;
      mb.tag_set_data(shared_tag, &conn[2][0], 1, &pstat);
      mb.tag_set_data(shared_tag, &conn[2][3], 1, &pstat);
      mb.tag_set_data(proc_tag, &h_edges[3], 1, &pstat);
      val = 0;
      mb.tag_set_data(proc_tag, &conn[2][1], 1, &val);
      mb.tag_set_data(proc_tag, &conn[2][2], 1, &val);
      pstat = 0;
      mb.tag_set_data(shared_tag, &conn[2][1], 1, &pstat);
      mb.tag_set_data(shared_tag, &conn[2][2], 1, &pstat);
      mb.tag_set_data(proc_tag, &h_edges[0], 3, &pstat);

      proc = 1;
      mb.tag_set_data(shared_proc, &conn[2][0], 1, &proc);
      mb.tag_set_data(shared_proc, &conn[2][3], 1, &proc);
      mb.tag_set_data(shared_proc, &h_edges[3], 1, &proc);
      proc = 3;
      mb.tag_set_data(shared_proc, &conn[2][1], 1, &proc);
      mb.tag_set_data(shared_proc, &conn[2][2], 1, &proc);
      mb.tag_set_data(shared_proc, &h_edges[1], 1, &proc);

      mb.tag_set_data(shared_handle, &conn[2][0], 1, &conn[1][1]);
      mb.tag_set_data(shared_handle, &conn[2][1], 1, &conn[3][0]);
      mb.tag_set_data(shared_handle, &conn[2][2], 1, &conn[3][3]);
      mb.tag_set_data(shared_handle, &conn[2][3], 1, &conn[1][2]);      
      break;
    case 3:
      val = 0;
      mb.tag_set_data(proc_tag, &elements, 1, &val);
      mb.tag_set_data(proc_tag, &conn[3][0], 4, &val);
      pstat = PSTATUS_NOT_OWNED;
      mb.tag_set_data(shared_tag, &conn[3][0], 1, &pstat);
      mb.tag_set_data(shared_tag, &conn[3][3], 1, &pstat);
      mb.tag_set_data(proc_tag, &h_edges[3], 1, &pstat);
      pstat = 0;
      mb.tag_set_data(shared_tag, &conn[3][1], 2, &pstat);
      mb.tag_set_data(shared_tag, &h_edges[0], 3, &pstat);

      proc = 2;
      mb.tag_set_data(shared_proc, &conn[3][0], 1, &proc);
      mb.tag_set_data(shared_proc, &conn[3][3], 1, &proc);
      mb.tag_set_data(shared_proc, &h_edges[3], 1,  &proc);

      mb.tag_set_data(shared_handle, &conn[3][0], 1, &conn[2][1]);
      mb.tag_set_data(shared_handle, &conn[3][3], 1, &conn[2][2]);
      break;
  }
}