#include <math.h>
#include <mpi.h>
#include <stdlib.h>

#include "simulation.h"
#include "random.h"
#include "mpicomm.h"

void sim_info_local(const sim_t *s, sim_info_t *local_info);
void allocate_memory(sim_t *s);
void deallocate_memory(sim_t *s);
void init_local_particles(sim_t *s);
void classify_particles(sim_t *s);
void relocate_particles(sim_t *s);

void sim_init(sim_t *s, const conf_t* configuration, const MPI_Comm communicator)
{
   s->configuration = configuration;
   s->communicator = communicator;

   MPI_Comm_rank(communicator, &s->mpi_rank);

   // allocate storage and initialize local particles with random values
   allocate_memory(s);
   init_local_particles(s);
}

void sim_free(sim_t *s)
{
   deallocate_memory(s);
}

void sim_info(const sim_t *s, sim_info_t *info)
{
   sim_info_t local_info;

   sim_info_local(s, &local_info);
   mpi_reduce_sim_info(&local_info, info);
}

void sim_info_local(const sim_t *s, sim_info_t *local_info)
{
   const part_t *particles = &s->local_particles;
   int count = particles->count;

   double *vx = particles->velocity.x;
   double *vy = particles->velocity.y;
   double *vsqr = malloc(count * sizeof(*vsqr));
   double vsqr_max = 0.;
   int i;

   for(i = 0; i < count; i++) {
      vsqr[i] = vx[i] * vx[i] + vy[i] * vy[i];
   }
   for(i = 0; i < count; i++) {
      if(vsqr_max < vsqr[i]) vsqr_max = vsqr[i];
   }
   free(vsqr);

   local_info->n_particles_total = s->local_particles.count;
   local_info->min_particles_per_cell = count;
   local_info->max_particles_per_cell = count;
   local_info->max_velocity = sqrt(vsqr_max);
}

void sim_calc_forces(sim_t *s)
{
   part_t *particles = &(s->local_particles);
   int count = particles->count;

   // just use random values for demonstration purpose
   double max_force = s->configuration->max_force;
   vec_t* force = &(particles->force);
   vec_set_random(force, max_force, count);
}

void sim_move_particles(sim_t *s)
{
   part_t *particles = &s->local_particles;
   int count = particles->count;

   const vec_t *force = &particles->force;
   const double *mass = particles->mass;
   double delta_t = s->configuration->delta_t;

   vec_t tmp_vector;
   vec_init(&tmp_vector, count);

   vec_t *acceleration = &tmp_vector;
   vec_t *velocity = &particles->velocity;
   vec_t *location = &particles->location;

   // acceleration = force / mass
   vec_arr_divide(acceleration, force, mass, count);

   // velocity += acceleration * delta_t
   vec_add_scaled(velocity, acceleration, delta_t, count);

   // location += velocity * delta_t, taking care of periodicity
   double modulus_x = s->configuration->ncells_x;
   double modulus_y = s->configuration->ncells_y;
   vec_add_scaled_mod(location, velocity, delta_t, modulus_x, modulus_y, count);

   vec_free(&tmp_vector);
}

void sim_communicate_particles(sim_t *s)
{
   int mode = s->configuration->transmission_mode;

   classify_particles(s);
   relocate_particles(s);

   switch(mode) {
      case DYNAMIC:
          mpi_communicate_particles_dynamic(s); break;
      case COLLECTIVE:
          mpi_communicate_particles_collective(s); break;
      case RMA:
          mpi_communicate_particles_rma(s); break;
   }
}

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

void allocate_memory(sim_t *s)
{
   // set up storage for particles in local cell and temporary working vector
   int count = s->configuration->particles_per_cell_initial;
   int capacity = count + count / 4;

   part_init(&s->local_particles, capacity);

   // set up buffer for particles leaving the local cell
   // (capacity will be increased on demand)
   part_init(&s->leaving_particles, 0);
}

void deallocate_memory(sim_t *s)
{
   part_free(&s->local_particles);

   part_free(&s->leaving_particles);
}

void init_local_particles(sim_t *s)
{
   int count = s->configuration->particles_per_cell_initial;
   part_t *particles = &s->local_particles;
   part_resize(particles, count);

   // particle mass
   double min_mass = s->configuration->min_mass;
   double max_mass = s->configuration->max_mass;
   array_set_random(particles->mass, min_mass, max_mass, count);

   // particle location
   int nx = s->configuration->ncells_x;
   int cell_x = s->mpi_rank % nx;
   int cell_y = s->mpi_rank / nx;
   double min_x = cell_x, max_x = min_x + 1.;
   double min_y = cell_y, max_y = min_y + 1.;
   vec_set_random_box(&particles->location, min_x, max_x, min_y, max_y, count);

   // particle velocity
   double vmax = s->configuration->max_velocity_initial;
   vec_set_random(&particles->velocity, vmax, count);

   // forces
   vec_set_zero(&particles->force, count);
}

void classify_particles(sim_t *s)
{
   part_t *particles = &s->local_particles;
   double *x = particles->location.x;
   double *y = particles->location.y;
   int *owner = particles->owner_rank;
   int count = particles->count;
   int i;

   if(s->configuration->use_cartesian_topology) {
      int coords[2];
      for(i = 0; i < count; i++) {
         coords[0] = (int)x[i];
         coords[1] = (int)y[i];
         MPI_Cart_rank(s->communicator, coords, &owner[i]);
      }
   } else {
      int nx = s->configuration->ncells_x;

      for(i = 0; i < count; i++) {
         owner[i] = (int)x[i] + nx * (int)y[i];
      }
   }
}

void relocate_particles(sim_t *s)
{
   int count = s->local_particles.count;
   int my_rank = s->mpi_rank;

   part_t *particles = &s->local_particles;
   int *owner = particles->owner_rank;

   int i, j, k;
   int block_start, block_len;

   i = 0; j = 0; k = 0;
   while(k < count)
   {
      // scan for particles leaving local cell
      block_start = k;
      for(; k < count && owner[k] != my_rank; k++);
      block_len = k - block_start;

      if(block_len > 0) {
         part_copy(&s->leaving_particles, j, &s->local_particles, block_start, block_len);
         j += block_len;
      }

      // scan for particles staying in local cell
      block_start = k;
      for(; k < count && owner[k] == my_rank; k++);
      block_len = k - block_start;
      if(block_len > 0) {
         if(i != block_start) {
            part_copy(&s->local_particles, i, &s->local_particles, block_start, block_len);
         }
         i += block_len;
      }
   }
   s->local_particles.count = i;
   s->leaving_particles.count = j;
}