#include <stdlib.h>
#include <stdio.h>

#include "mesh.h"
#include "mpitypes.h"
#include "mpicomm.h"

// ---------------------------------------------- Helper function declarations

void mesh_find_neighbors(mesh_t *mesh, const conf_t *configuration, const box_t *domains);

void mesh_print_neighbors(const mesh_t *mesh, const conf_t *configuraiton, const box_t *domains, FILE *f);

void mesh_setup_index_mappings(mesh_t *mesh);

void halo_info_init(halo_info *halo, const box_t *domain);

void halo_info_setup_index_mapping(halo_info *halo, const mesh_t *mesh);

void halo_info_free(halo_info *halo);

int communicator_reordered(MPI_Comm new_comm);

// ===========================================================================

void mesh_init(mesh_t *mesh, const conf_t *configuration, const box_t *decomposition)
{
   if(debug_enabled(configuration)) {
      printf("Examining neighborhood topology for MPI Graph communicator creation...\n\n");
   }

   MPI_Comm_dup(MPI_COMM_WORLD, &mesh->communicator);
   mesh->n_neighbors = 0;
   mesh->neighbors = NULL;
   mesh_find_neighbors(mesh, configuration, decomposition);
   mesh_print_neighbors(mesh, configuration, decomposition, stdout);

   if(trace_enabled(configuration)) {
     mesh_print_neighbors(mesh, configuration, decomposition, stdout);
   }

   if(configuration->transfer_mode == SPARSE_COLLECTIVE) {
      int ranks_reordered;

      // Create MPI Graph communicator
      if(debug_enabled(configuration)) {
         printf("Creating MPI Graph communicator...\n\n");
      }
      MPI_Comm_free(&mesh->communicator);
      mpi_create_graph_communicator(mesh, configuration, &mesh->communicator);

      ranks_reordered = communicator_reordered(mesh->communicator);
      if(info_enabled(configuration)) {
         printf("INFO: MPI reordered ranks: %s\n\n", ranks_reordered ? "YES" : "NO");
      }

      if(ranks_reordered) {
         if(debug_enabled(configuration)) {
            printf("Re-examining neighborhood topology after rank-reordering...\n\n");
         }
         mesh_find_neighbors(mesh, configuration, decomposition);
         mesh_print_neighbors(mesh, configuration, decomposition, stdout);
      }
   }

   // Setup index mappings (plus send and receive types for int and double)
   if(debug_enabled(configuration)) {
      printf("Setting up index mappings and MPI types...\n\n");
   }
   mesh_setup_index_mappings(mesh);
}

void mesh_find_neighbors(mesh_t *mesh, const conf_t *configuration, const box_t *domains)
{
   int mpi_rank, nprocs;
   int i, n;
   box_t tmp;
   box_t local_domain, extended_local_domain, extended_remote_domain;
   box_t halo_incoming, halo_outgoing;

   MPI_Comm_rank(mesh->communicator, &mpi_rank);
   MPI_Comm_size(mesh->communicator, &nprocs);

   // Determine local domain (without and with halo cells)
   local_domain = domains[mpi_rank];
   tmp = local_domain;
   box_grow(&tmp, 1);
   box_intersect(&tmp, &configuration->global_domain, &extended_local_domain);

   // Reserve space for neighbors
   int capacity = 32;
   neighbor_info *neighbors = mesh->neighbors;
   neighbors = realloc(neighbors, capacity * sizeof(*neighbors));

   for(i = 0, n = 0; i < nprocs; i++) {
      if(i == mpi_rank) continue;

      // Determine cells for incoming halo data
      box_intersect(&extended_local_domain, &domains[i], &halo_incoming);
      if(box_is_empty(&halo_incoming)) continue; // Not a neighbor

      // Reserve more space for neighbors if necessary
      if(n >= capacity) {
         capacity += capacity / 2;
         neighbors = realloc(neighbors, capacity * sizeof(*neighbors));
      }

      // Determine cells for outgoing halo data
      extended_remote_domain = domains[i];
      box_grow(&extended_remote_domain, 1);
      box_intersect(&local_domain, &extended_remote_domain, &halo_outgoing);

      // Initialize neighbor info
      neighbors[n].mpi_rank = i;
      halo_info_init(&neighbors[n].halo_incoming, &halo_incoming);
      halo_info_init(&neighbors[n].halo_outgoing, &halo_outgoing);

      n++;
   }

   // Update mesh data
   mesh->local_domain = local_domain;
   mesh->extended_local_domain = extended_local_domain;

   mesh->n_neighbors = n;
   mesh->neighbors = neighbors;
}

void mesh_print_neighbors(const mesh_t *mesh, const conf_t *configuration, const box_t *domains, FILE *f)
{
   if(debug_enabled(configuration)) {
      int mesh_comm_rank, comm_world_rank, neighbor_rank, i, n;
      char buf[128];

      MPI_Comm_rank(mesh->communicator, &mesh_comm_rank);
      MPI_Comm_rank(MPI_COMM_WORLD, &comm_world_rank);

      n = mesh->n_neighbors;

      if(mesh_comm_rank != comm_world_rank) {
         fprintf(f, "Found %d neighbors for rank %d (%d in MPI_COMM_WORLD):\n", n, mesh_comm_rank, comm_world_rank);
      } else {
         fprintf(f, "Found %d neighbors for rank %d:\n", n, mesh_comm_rank);
      }

      if(trace_enabled(configuration)) {
         fprintf(f, " * %4d: %s\n---\n", mesh_comm_rank,
                box_to_string(&mesh->local_domain, buf, sizeof(buf)));
         for(i = 0; i < n; i++) {
            neighbor_rank = mesh->neighbors[i].mpi_rank;
            fprintf(f, " * %4d: %s\n", neighbor_rank,
                   box_to_string(&domains[neighbor_rank], buf, sizeof(buf)));
         }
         fprintf(f, "\n");
      } else {
         for(i = 0; i < n; i++) {
            neighbor_rank = mesh->neighbors[i].mpi_rank;
            fprintf(f, " %d", neighbor_rank);
         }
         fprintf(f, "\n\n");
      }
   }
}

int mesh_idx(const mesh_t *mesh, int x, int y, int z)
{
   const int *coords = mesh->extended_local_domain.coords;
   const int *e = mesh->extended_local_domain.extents;

   x -= coords[X];
   y -= coords[Y];
   z -= coords[Z];

   return x + e[X] * (y + e[Y] * z);
}

void mesh_setup_index_mappings(mesh_t *mesh)
{
   int mpi_rank, nprocs, i;
   neighbor_info *neighbors = mesh->neighbors;

   MPI_Comm_rank(mesh->communicator, &mpi_rank);
   MPI_Comm_size(mesh->communicator, &nprocs);

   for(i = 0; i < mesh->n_neighbors; i++)
   {
      halo_info_setup_index_mapping(&neighbors[i].halo_incoming, mesh);
      halo_info_setup_index_mapping(&neighbors[i].halo_outgoing, mesh);
   }
}

void mesh_free(mesh_t *mesh)
{
   int i;

   for(i = 0; i < mesh->n_neighbors; i++) {
      halo_info_free(&mesh->neighbors[i].halo_incoming);
      halo_info_free(&mesh->neighbors[i].halo_outgoing);
   }

   mesh->n_neighbors = 0;
   mesh->neighbors = realloc(mesh->neighbors, 0);
   MPI_Comm_free(&mesh->communicator);
}

void halo_info_init(halo_info *halo, const box_t *domain)
{
   halo->domain = *domain;
   halo->cell_indices = NULL;
   halo->transfer_type_int = MPI_DATATYPE_NULL;
   halo->transfer_type_double = MPI_DATATYPE_NULL;
}

void halo_info_setup_index_mapping(halo_info *halo, const mesh_t *mesh)
{
   int x, y, z;
   const int *c = halo->domain.coords;
   const int *e = halo->domain.extents;
   int count = halo->domain.volume;
   int *indices;

   indices = malloc(count * sizeof(*indices));

   int i = 0;
   for(z = c[Z]; z < c[Z] + e[Z]; z++) {
      for(y = c[Y]; y < c[Y] + e[Y]; y++) {
         for(x = c[X]; x < c[X] + e[X]; x++) {
            indices[i] = mesh_idx(mesh, x, y, z);
            i++;
         }
      }
   }

   halo->cell_indices = indices;
   mpitype_indexed_int(count, indices, &halo->transfer_type_int);
   mpitype_indexed_double(count, indices, &halo->transfer_type_double);
}

void halo_info_free(halo_info *halo)
{
   free(halo->cell_indices);
   halo->cell_indices = NULL;
   if(halo->transfer_type_int != MPI_DATATYPE_NULL) {
      MPI_Type_free(&halo->transfer_type_int);
   }
   if(halo->transfer_type_double != MPI_DATATYPE_NULL) {
      MPI_Type_free(&halo->transfer_type_double);
   }
}

// ---------------------------------------------------------- Helper functions

int communicator_reordered(MPI_Comm new_comm)
{
   int comm_world_rank, new_rank;
   int local_ranks_different;
   int ranks_reordered = 0;

   MPI_Comm_rank(MPI_COMM_WORLD, &comm_world_rank);
   MPI_Comm_rank(new_comm, &new_rank);

   local_ranks_different = comm_world_rank != new_rank;
   MPI_Allreduce(
         &local_ranks_different, &ranks_reordered, 1, MPI_INT,
         MPI_LOR, MPI_COMM_WORLD
   );
   return ranks_reordered;
}