#include "ep_lib.hpp"
#include <mpi.h>
#include "ep_declaration.hpp"
#include "ep_mpi.hpp"

using namespace std;

namespace ep_lib
{

  void vec_simplify(std::vector<int> *inout_vector)
  {
    std::vector<int> out_vec;
    int found=false;
    for(std::vector<int>::iterator it_in = inout_vector->begin() ; it_in != inout_vector->end(); ++it_in)
    {
      for(std::vector<int>::iterator it = out_vec.begin() ; it != out_vec.end(); ++it)
      {
        if(*it_in == *it)
        {
          found=true;
          break;
        }
        else found=false;
      }
      if(found == false)
      {
        out_vec.push_back(*it_in);
      }
    }
    inout_vector->swap(out_vec);
  }
  
  void vec_simplify(std::vector<int> *in_vector, std::vector<int> *out_vector)
  {
    int found=false;
    for(std::vector<int>::iterator it_in = in_vector->begin() ; it_in != in_vector->end(); ++it_in)
    {
      for(std::vector<int>::iterator it = out_vector->begin() ; it != out_vector->end(); ++it)
      {
        if(*it_in == *it)
        {
          found=true;
          break;
        }
        else found=false;
      }
      if(found == false)
      {
        out_vector->push_back(*it_in);
      }
    }
  }



  int MPI_Comm_split(MPI_Comm comm, int color, int key, MPI_Comm *newcomm)
  {
    int ep_rank, ep_rank_loc, mpi_rank;
    int ep_size, num_ep, mpi_size;

    ep_rank = comm->ep_comm_ptr->size_rank_info[0].first;
    ep_rank_loc = comm->ep_comm_ptr->size_rank_info[1].first;
    mpi_rank = comm->ep_comm_ptr->size_rank_info[2].first;
    ep_size = comm->ep_comm_ptr->size_rank_info[0].second;
    num_ep = comm->ep_comm_ptr->size_rank_info[1].second;
    mpi_size = comm->ep_comm_ptr->size_rank_info[2].second;

    int num_color = 0;

    int color_index;

    vector<int> matched_number;
    vector<int> matched_number_loc;
    
    vector<int> all_color(ep_size);
    vector<int> all_color_loc(num_ep);

    MPI_Allgather(&color, 1, MPI_INT, all_color.data(), 1, MPI_INT, comm);
    MPI_Allgather_local(&color, 1, MPI_INT, all_color_loc.data(), comm);

    list<int> color_list(all_color.begin(), all_color.end());
    list<int> color_list_loc(all_color_loc.begin(), all_color_loc.end());

    vector<int> all_color_simplified;
    vec_simplify(&all_color, &all_color_simplified);
    int number_of_color;
    for(int i=0; i<all_color_simplified.size(); i++)
    {
      if(color == all_color_simplified[i])
      {
        number_of_color = i;
        break;
      }
    }

    matched_number.resize(all_color_simplified.size(), 0);
    matched_number_loc.resize(all_color_simplified.size(), 0);


    while(!color_list.empty())
    {
      int target_color = color_list.front();
      for(list<int>::iterator it = color_list.begin(); it != color_list.end(); ++it)
      {
        if(*it == target_color)
        {
          matched_number[num_color]++;
        }
      }
      for(list<int>::iterator it = color_list_loc.begin(); it != color_list_loc.end(); ++it)
      {
        if(*it == target_color)
        {
          matched_number_loc[num_color]++;
        }
      }
      color_list.remove(target_color);
      color_list_loc.remove(target_color);
      num_color++;
    }
        

    vector<int> all_key(ep_size);
    vector<int> all_key_loc(num_ep);
    
    vector<int> colored_key[num_color];
    vector<int> colored_key_loc[num_color];
    

    MPI_Allgather(&key, 1, MPI_INT, all_key.data(),1, MPI_INT, comm);
    MPI_Allgather_local(&key, 1, MPI_INT, all_key_loc.data(), comm);

    for(int i=0; i<num_ep; i++)
    {
      for(int j = 0; j<num_color; j++)
      {
        if(all_color_loc[i] == all_color_simplified[j])
        {
          colored_key_loc[j].push_back(all_key_loc[i]);
        }
      }
    }
    
    for(int i=0; i<ep_size; i++)
    {
      for(int j = 0; j<num_color; j++)
      {
        if(all_color[i] == all_color_simplified[j])
        {
          colored_key[j].push_back(all_key[i]);
        }
      }
    }
    
    for(int i=0; i<num_color; i++)
    {
      std::sort(colored_key[i].begin(), colored_key[i].end());
      std::sort(colored_key_loc[i].begin(), colored_key_loc[i].end());
    }

    int new_ep_rank;
    
    for(int i=0; i<colored_key[number_of_color].size(); i++)
    {
      if(key == colored_key[number_of_color][i])
      {
        new_ep_rank = i;
        break;
      }
    }
    
    int new_ep_rank_loc;
    
    for(int i=0; i<colored_key_loc[number_of_color].size(); i++)
    {
      if(key == colored_key_loc[number_of_color][i])
      {
        new_ep_rank_loc = i;
        break;
      }
    }
    

    ::MPI_Comm **split_mpi_comm;
    split_mpi_comm = new ::MPI_Comm* [num_color];
    for(int ii=0; ii<num_color; ii++)
      split_mpi_comm[ii] = new ::MPI_Comm;

    for(int j=0; j<num_color; j++)
    {
      if(ep_rank_loc == 0)
      {
        int master_color = 1;
        if(matched_number_loc[j] == 0) master_color = MPI_UNDEFINED;

        ::MPI_Comm_split(to_mpi_comm(comm->mpi_comm), master_color, mpi_rank, split_mpi_comm[j]);
        
        comm->ep_comm_ptr->comm_list[0]->mpi_bridge = split_mpi_comm[j];
      }
      
      MPI_Barrier_local(comm);
      
      int num_new_ep = 0;

      if(new_ep_rank_loc == 0 && color == all_color_simplified[j])
      {
        num_new_ep = matched_number_loc[j];
        MPI_Info info;
        MPI_Comm *ep_comm;

        MPI_Comm_create_endpoints(comm->ep_comm_ptr->comm_list[0]->mpi_bridge, num_new_ep, info, ep_comm);

        comm->ep_comm_ptr->comm_list[0]->mem_bridge = ep_comm;
        
        (*ep_comm)->ep_rank_map->clear();
        
        memcheck("in MPI_Split ep_rank="<< ep_rank <<" : *ep_comm = "<< *ep_comm);
      }
      
      MPI_Barrier_local(comm);
      
      if(color == all_color_simplified[j])
      {
        *newcomm = comm->ep_comm_ptr->comm_list[0]->mem_bridge[new_ep_rank_loc];
        memcheck("in MPI_Split ep_rank="<< ep_rank <<" : *newcomm = "<< *newcomm);

        (*newcomm)->ep_comm_ptr->comm_label = color;
        
        (*newcomm)->ep_comm_ptr->size_rank_info[0].first = new_ep_rank;
        (*newcomm)->ep_comm_ptr->size_rank_info[1].first = new_ep_rank_loc;
        
        int my_triple[3];
        vector<int> my_triple_vector;
        vector<int> my_triple_vector_recv;
        my_triple[0] = new_ep_rank;
        my_triple[1] = new_ep_rank_loc;
        my_triple[2] = (*newcomm)->ep_comm_ptr->size_rank_info[2].first; // new_mpi_rank
        
        int new_ep_size = (*newcomm)->ep_comm_ptr->size_rank_info[0].second;
        int new_num_ep  = (*newcomm)->ep_comm_ptr->size_rank_info[1].second;
        
        int new_mpi_size = (*newcomm)->ep_comm_ptr->size_rank_info[2].second;
        
        if(new_ep_rank_loc == 0) my_triple_vector.resize(3*new_ep_size);
        if(new_ep_rank_loc == 0) my_triple_vector_recv.resize(3*new_ep_size);
        
        MPI_Gather_local(my_triple, 3, MPI_INT, my_triple_vector.data(), 0, *newcomm);
        
        if(new_ep_rank_loc == 0)
        {
          int *recvcounts = new int[new_mpi_size];
          int *displs = new int[new_mpi_size];
          int new_num_epx3 = new_num_ep * 3;
          ::MPI_Allgather(&new_num_epx3, 1, to_mpi_type(MPI_INT), recvcounts, 1, to_mpi_type(MPI_INT), to_mpi_comm((*newcomm)->mpi_comm));
          displs[0]=0;
          for(int i=1; i<new_mpi_size; i++)
            displs[i] = displs[i-1] + recvcounts[i-1];
             
          ::MPI_Allgatherv(my_triple_vector.data(), 3*new_num_ep, to_mpi_type(MPI_INT), my_triple_vector_recv.data(), recvcounts, displs, to_mpi_type(MPI_INT), to_mpi_comm((*newcomm)->mpi_comm));
          
          for(int i=0; i<new_ep_size; i++)
          {
            (*newcomm)->ep_comm_ptr->comm_list[0]->ep_rank_map->insert(std::pair< int, std::pair<int,int> >(my_triple_vector_recv[3*i], my_triple_vector_recv[3*i+1], my_triple_vector_recv[3*i+2]));
          }
          
          (*newcomm)->ep_rank_map = (*newcomm)->ep_comm_ptr->comm_list[0]->ep_rank_map;
          
          delete recvcounts;
          delete displs;
        } 
      }
    }
    
    /*for(int i=0; i<ep_size; i++)
    {
      MPI_Barrier(comm);
      MPI_Barrier(comm);
      if(ep_rank==i) 
      {
        printf("ep_rank_map for endpoint %d = \n", ep_rank);
        for(std::map<int, std::pair<int, int> > :: iterator it = (*newcomm)->ep_rank_map->begin(); it != (*newcomm)->ep_rank_map->end(); it++)
        {
          printf("\t\t\t %d %d %d\n", it->first, it->second.first, it->second.second);
        }
        printf("\n");
      }
      MPI_Barrier(comm);
      MPI_Barrier(comm);
    }*/   
    
    return 0;
  }

}