#include <assert.h>
#include <stdlib.h>

#include "mpitypes.h"

#include "configuration.h"
#include "particles.h"
#include "simulation.h"

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

void mpitype_indexed_int(int count, int *displacements, MPI_Datatype *type);

void mpitype_indexed_double(int count, int *displacemnts, MPI_Datatype *type);

MPI_Aint* part_get_addresses(const part_t *particles, int displacement);

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

void mpitype_conf_init(MPI_Datatype *new_type)
{
   conf_t dummy;
   int i;

   // CONF_T_* constants defined in configuration.h
   int blocklengths[] = {CONF_T_N_INT_MEMBERS, CONF_T_N_DOUBLE_MEMBERS};
   MPI_Datatype types[] = {MPI_INT, MPI_DOUBLE};
   MPI_Aint displacements[2];
   MPI_Aint base;

   MPI_Get_address(&dummy, &base);
   MPI_Get_address(&dummy.CONF_T_FIRST_INT_MEMBER, &displacements[0]);
   MPI_Get_address(&dummy.CONF_T_FIRST_DOUBLE_MEMBER, &displacements[1]);

   for(i = 0; i < 2; i++) displacements[i] -= base;

   MPI_Type_create_struct(2, blocklengths, displacements, types, new_type);
   MPI_Type_commit(new_type);
}

void mpitype_conf_free(MPI_Datatype *type)
{
   MPI_Type_free(type);
}

void mpitype_sim_info_init(MPI_Datatype *new_type)
{
   sim_info_t dummy;
   int i;

   int blocklengths[] = {
      SIM_INFO_T_N_LONG_MEMBERS,
      SIM_INFO_T_N_INT_MEMBERS,
      SIM_INFO_T_N_DOUBLE_MEMBERS
   };
   MPI_Datatype types[] = {MPI_LONG, MPI_INT, MPI_DOUBLE};
   MPI_Aint displacements[3];
   MPI_Aint base;

   MPI_Get_address(&dummy, &base);
   MPI_Get_address(&dummy.SIM_INFO_T_FIRST_LONG_MEMBER, &displacements[0]);
   MPI_Get_address(&dummy.SIM_INFO_T_FIRST_INT_MEMBER, &displacements[1]);
   MPI_Get_address(&dummy.SIM_INFO_T_FIRST_DOUBLE_MEMBER, &displacements[2]);

   for(i = 0; i < 3; i++) displacements[i] -= base;

   MPI_Type_create_struct(3, blocklengths, displacements, types, new_type);
   MPI_Type_commit(new_type);
}

void mpitype_sim_info_free(MPI_Datatype *type)
{
   MPI_Type_free(type);
}

MPI_Aint* part_get_addresses(const part_t *particles, int displacement)
{
   static MPI_Aint addresses[5];
   assert(displacement < particles->count);

   MPI_Get_address(particles->mass + displacement, &addresses[0]);
   MPI_Get_address(particles->velocity.x + displacement, &addresses[1]);
   MPI_Get_address(particles->velocity.y + displacement, &addresses[2]);
   MPI_Get_address(particles->location.x + displacement, &addresses[3]);
   MPI_Get_address(particles->location.y + displacement, &addresses[4]);

   return addresses;
}

void mpitype_part_init_send(part_t *particles, MPI_Datatype *datatypes[], int *send_counts[], int *receiver_ranks[], int *n_receivers)
{
   int part_count = particles->count;
   int *part_owner = particles->owner_rank;

   if(part_count == 0) {
      // no particles are leaving local cell
      *datatypes = NULL;
      *send_counts = NULL;
      *receiver_ranks = NULL;
      *n_receivers = 0;
      return;
   }

   // Only 5 components of part_t structure need to be exchanged:
   // mass, velocity.x, velocity.y, location.x, location.y
   int blocklengths[] = {1, 1, 1, 1, 1};
   MPI_Aint *addresses = part_get_addresses(particles, 0);
   MPI_Datatype struct_component_types[5];

   // Group particles by owner rank; for each group create mpi send type
   int n, i, j, k, recv_rank;

   int recv_capacity = 8;
   MPI_Datatype *types = malloc(recv_capacity * sizeof(*types));
   int *counts = malloc(recv_capacity * sizeof(*counts));
   int *recv_ranks = malloc(recv_capacity * sizeof(*recv_ranks));
   int *part_indices = malloc(part_count * sizeof(*part_indices));
   MPI_Datatype indexed_double_type;
   MPI_Datatype particle_struct_type;

   // Iterate over particles
   for(i = 0, n = 0; i < part_count; i++)
   {
      recv_rank = part_owner[i];
      if(recv_rank < 0) continue; // Particle already processed

      // Find indices of particles with same owner
      for(j = i, k = 0; j < part_count; j++) {
         if(part_owner[j] == recv_rank) {
            part_indices[k] = j;
            k++;
            part_owner[j] = -1; // Mark j-th particle as processed
         }
      }
      // Create (temporary) MPI Datatypes for struct components (5 x indexed double)
      mpitype_indexed_double(k, part_indices, &indexed_double_type);
      struct_component_types[0] = indexed_double_type; // mass
      struct_component_types[1] = indexed_double_type; // velocity.x
      struct_component_types[2] = indexed_double_type; // velocity.y
      struct_component_types[3] = indexed_double_type; // location.x
      struct_component_types[4] = indexed_double_type; // location.y

      // Create (final) MPI Datatype for particle struct
      MPI_Type_create_struct(5, blocklengths, addresses, struct_component_types, &particle_struct_type);
      MPI_Type_commit(&particle_struct_type);

      // Free temporary MPI type
      MPI_Type_free(&indexed_double_type);

      // Store type and rank of corresponding recevier
      if(n >= recv_capacity) {
         recv_capacity += recv_capacity / 2;
         types = realloc(types, recv_capacity * sizeof(*types));
         counts = realloc(counts, recv_capacity * sizeof(*counts));
         recv_ranks = realloc(recv_ranks, recv_capacity * sizeof(*recv_ranks));
      }
      types[n] = particle_struct_type;
      counts[n] = k;
      recv_ranks[n] = recv_rank;
      n++;
   }
   // Set output args
   *datatypes = types;
   *send_counts = counts;
   *receiver_ranks = recv_ranks;
   *n_receivers = n;
}

void mpitype_part_free_send(MPI_Datatype *datatypes[], int *send_counts[], int *receiver_ranks[], int n_receivers)
{
   MPI_Datatype *types;
   int *counts, *receivers;
   int i;

   types = *datatypes; *datatypes = NULL;
   counts = *send_counts; *send_counts = NULL;
   receivers = *receiver_ranks; *receiver_ranks = NULL;

   for(i = 0; i < n_receivers; i++) {
      MPI_Type_free(&types[i]);
   }
   free(types);
   free(counts);
   free(receivers);
}

int mpitype_part_get_count(const MPI_Status *status)
{
   const int n_doubles_per_particle = 5; // mass, velocity.x, velocity.y, location.x, location.y
   int n_doubles;
   MPI_Get_count(status, MPI_DOUBLE, &n_doubles);
   return n_doubles / n_doubles_per_particle;
}

void mpitype_part_init_recv(const part_t* particles, int blocklength, int displacement, MPI_Datatype* type)
{
   int blocklengths[5];
   MPI_Aint *addresses = part_get_addresses(particles, displacement);
   MPI_Datatype struct_component_types[5];
   int i;

   for(i = 0; i < 5; i++) blocklengths[i] = blocklength;

   struct_component_types[0] = MPI_DOUBLE; // mass
   struct_component_types[1] = MPI_DOUBLE; // velocity.x
   struct_component_types[2] = MPI_DOUBLE; // velocity.y
   struct_component_types[3] = MPI_DOUBLE; // location.x
   struct_component_types[4] = MPI_DOUBLE; // location.y

   MPI_Type_create_struct(5, blocklengths, addresses, struct_component_types, type);
   MPI_Type_commit(type);
}

void mpitype_part_free_recv(MPI_Datatype* type)
{
   MPI_Type_free(type);
}

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

void mpitype_indexed_int(int count, int *displacements, MPI_Datatype *type)
{
   MPI_Type_create_indexed_block(count, 1, displacements, MPI_INT, type);
}

void mpitype_indexed_double(int count, int *displacements, MPI_Datatype *type)
{
   MPI_Type_create_indexed_block(count, 1, displacements, MPI_DOUBLE, type);
}


