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

#define NB_THREADS 4



struct map {
  int width;
  int height;

  double **cur_values;
  double **next_values;
  pthread_barrier_t barrier;
};

void swap_values(struct map* m) {
  double ** tmp = m->next_values;
  m->next_values = m->cur_values;
  m->cur_values =  tmp;  
}

void compute_next_value(struct map* m, int i, int j) {
  m->next_values[i][j] = (m->cur_values[i-1][j] + m->cur_values[i+1][j] + m->cur_values[i][j-1] + m->cur_values[i][j+1] + m->cur_values[i][j]) / 5;
}

void compute_next_values(struct map* m) {
  for(int i=1; i<m->height-1; i++) {
    for(int j=1; j<m->width-1; j++) {
      compute_next_value(m, i, j);
    }
  }

  swap_values(m);
}

void print_color(double v) {
  printf("\x1b[0m"); // reset colors

  if(v >= 10) printf("\x1b[41m"); // red background
  else if(v > 1) printf("\x1b[43m"); // yellow background
  else if(v < -10) printf("\x1b[44m"); // blue background
  else if(v < -1) printf("\x1b[46m"); // cyan background

  printf("%3.2d ", (int) v);
  printf("\x1b[0m"); // reset colors
}

void clear_screen() {
  printf("\e[1;1H\e[2J");
}

void print_map(struct map* m) {
  clear_screen();
  for(int i=0; i<m->height; i++) {
    for(int j=0; j<m->width; j++) {
      print_color(m->cur_values[i][j]);
    }
    printf("\n");
  }
}

void heat_random_point(struct map* m) {
  int i = 0;
  int j = 0;
  while(i==0 || i >= m->height-1) {
   i =  rand() % m->height;
  }
  while(j==0 || j >= m->width-1) {
   j =  rand() % m->width;
  }

  m->cur_values[i][j] = 99;

  if(rand()%2) // randomly add cold points
    m->cur_values[i][j] = -m->cur_values[i][j];
}

void init_map(struct map* m, int height, int width, int nb_heat_points) {
  m->width = width;
  m->height = height;
  m->cur_values = malloc(sizeof(double*) * height);
  m->next_values = malloc(sizeof(double*) * height);

  for(int i=0; i<height; i++) {
    m->cur_values[i] = malloc(sizeof(double) * width);
    m->next_values[i] = malloc(sizeof(double) * width);
    for(int j=0; j<height; j++) {
      m->cur_values[i][j] = 0;
      m->next_values[i][j] = 0;
    }
  }
 
  for(int i = 0; i<nb_heat_points; i++)
    heat_random_point(m);
}

struct thread_param {
  struct map* map;
  int i_start;
  int i_stop;
  int niter;
};

void* worker(void* arg) {
  struct thread_param *param = arg;
  static _Atomic int next_rank;
  int my_rank = next_rank++;

  printf("Le thread %d traitera les lignes %d à %d\n", my_rank, param->i_start, param->i_stop);

  for(int i=0; i<param->niter; i++) {
    for(int i = param->i_start; i <= param->i_stop; i++) {
      for(int j=1; j<param->map->width; j++) {
	compute_next_value(param->map, i, j);
      }
    }
    pthread_barrier_wait(&param->map->barrier);
  }    
  
  return NULL;
}

int main(int argc, char**argv) {
  int world_height = 4096;
  int world_width  = 4096;
  int nb_heat_points = 10;
  int niter = 50;

  srand(time(NULL));
  struct map m;
  init_map(&m, world_height, world_width, nb_heat_points);

  pthread_barrier_init(&m.barrier, NULL, NB_THREADS+1);

  struct thread_param params[NB_THREADS];
  pthread_t tids[NB_THREADS];
  int i_start = 1;
  int i_stop = 0;


  struct timespec t1, t2;
  clock_gettime(CLOCK_MONOTONIC, &t1);
  
  for(int i=0; i<NB_THREADS; i++) {
    i_start = i_stop+1;
    i_stop = i_start + world_height/NB_THREADS;
    if(i_stop > world_height-2)
      i_stop = world_height-2;
    params[i].map = &m;
    params[i].i_start = i_start;
    params[i].i_stop = i_stop;
    params[i].niter = niter;

    pthread_create(&tids[i], NULL, worker, &params[i]);
  }

  for(int i=0; i<niter; i++) {
    pthread_barrier_wait(&m.barrier);
    swap_values(&m);
  }
  
  for(int i=0; i<NB_THREADS; i++) {
    pthread_join(tids[i], NULL);
  }


  clock_gettime(CLOCK_MONOTONIC, &t2);
  double duration = (t2.tv_sec - t1.tv_sec) + (t2.tv_nsec - t1.tv_nsec)/1e9;
  printf("%d * %d (%d iterations) computed in %lf s\n", world_height, world_width, niter, duration);

  return EXIT_SUCCESS;
}