/* -*- c-file-style: "GNU" -*- */
/*
 * Copyright (C) CNRS, INRIA, Universite Bordeaux 1, Telecom SudParis
 * See COPYING in top-level directory.
 */

#include <assert.h>
#include <stdio.h>
#include "mpi_ev_codes.h"
#include "eztrace_convert.h"
#include "eztrace_convert_mpi.h"
#include "eztrace_hierarchical_array.h"
#include "eztrace_convert_mpi_coll.h"

static struct ezt_list_t finished_collectives[mpi_coll_type_max];

void init_mpi_coll_messages() {
  hierarchical_array_attach(EZTRACE_MPI_STATS_COLL_MSG_ID,
			    sizeof(struct coll_msg_event));

  for(int i=0; i<mpi_coll_type_max; i++) {
    ezt_list_new(&finished_collectives[i]);
  }
}

void print_coll_stats();


void __print_coll_message_header(FILE*stream) {
  fprintf(stream, "#[ID]\tCollective_type\troot_rank\t[rank_0_thread enter_timestamp exit_timestamp] [rank_1_thread enter_timestamp exit_timestamp]...\n");
}

void __print_coll_message(FILE*stream, struct mpi_coll_msg_t *msg) {
  fprintf(stream, "[%p]\t%s\t%d\t", msg, COLL_TYPE_STR(msg->type), ezt_get_global_rank(msg->comm[0], msg->root_process));
  int i;
  for (i = 0; i < msg->comm_size; i++) {
    fprintf(stream, "[%s %lu %lu] ", msg->thread_ids[i],
            msg->times[i][start_coll], msg->times[i][stop_coll]);
  }
  fprintf(stream, "\n");
}

void __print_coll_messages_recurse(FILE*stream, unsigned depth,
                                   p_eztrace_container p_cont);
void __print_coll_messages(FILE*stream) {
  int trace_index;
  struct ezt_list_token_t *token;

  for(trace_index=0; trace_index<NB_TRACES; trace_index++) {
    INIT_MPI_PROCESS_INFO(GET_PROCESS_INFO(trace_index), process_info);

    ezt_list_foreach(&process_info->pending_comm[comm_type_collective], token) {
      struct mpi_pending_coll_comm* p_comm = (struct mpi_pending_coll_comm*) token->data;
      struct mpi_coll_msg_t* msg = p_comm->msg;
      __print_coll_message(stream, msg);
    }
  }
}

void ezt_mpi_dump_coll_messages(FILE* stream) {
  for(int i=0; i<mpi_coll_type_max; i++) {
    struct ezt_list_t*l = &finished_collectives[i];
    struct ezt_list_token_t *token;

    if(l->nb_item > 0) {
      char* path;
      int ret = asprintf(&path, "%s/mpi_collective_message_dump_%s",
			 eztrace_stats_get_output_dir(), COLL_TYPE_STR(i));

      FILE *f = fopen(path, "w");
      if (!f) {
	perror("Error while dumping collective messages");
      }

      __print_coll_message_header(f);
      ezt_list_foreach(l, token) {
	struct mpi_pending_coll_comm* p_comm = (struct mpi_pending_coll_comm*) token->data;
	struct mpi_coll_msg_t* msg = p_comm->msg;
	__print_coll_message(f, msg);
      }

      ret = fclose(f);
      if (ret) {
	perror("Error while dumping collective messages (fclose)");
      }
      printf("\tMPI collective messages dumped in '%s'\n", path);
      free(path);
    }
  }
}

/* create a new collective message */
static struct mpi_coll_msg_t* __new_coll_message(enum coll_type_t type,
						 struct ezt_mpi_comm *comm,
						 int data_size)
{
  struct mpi_coll_msg_t *msg = malloc(sizeof(struct mpi_coll_msg_t));
  int comm_size = comm->comm_size;

  msg->type = type;
  msg->comm_size = comm_size;
  msg->comm = malloc(sizeof(struct ezt_mpi_comm*) * comm_size);
  int i;
  msg->data_size = data_size;
  msg->times = malloc(sizeof(uint64_t*) * comm_size);
  msg->thread_ids = malloc(sizeof(char*) * comm_size);
  msg->link_value = malloc(sizeof(char**) * comm_size);
  msg->link_id = malloc(sizeof(char**) * comm_size);
  msg->requests = malloc(sizeof(struct mpi_request *) * comm_size);
  msg->nb_started = 0;
  msg->root_process = -1;
  for (i = 0; i < comm_size; i++) {
    int j;
    msg->comm[i] = NULL;
    msg->thread_ids[i] = NULL;
    msg->link_value[i] = malloc(sizeof(char*) * comm_size);
    msg->link_id[i] = malloc(sizeof(char*) * comm_size);
    msg->requests[i] = NULL;
    for (j = 0; j < comm_size; j++) {
      CREATE_COLL_LINK_VALUE(msg, i, j);
      CREATE_COLL_MSG_ID(msg, i, j);
    }

    msg->times[i] = malloc(sizeof(uint64_t) * COLL_NB_TIMES);
    for (j = 0; j < COLL_NB_TIMES; j++)
      msg->times[i][j] = TIME_INIT;
  }

  return msg;
}

/* add a new pending collective communicaton */
static struct mpi_pending_coll_comm*
__create_new_pending_collective_comm(int trace_index,
						 struct mpi_coll_msg_t* msg) {
  INIT_MPI_PROCESS_INFO(GET_PROCESS_INFO(trace_index), process_info);

  struct mpi_pending_coll_comm* p_comm = malloc(sizeof(struct mpi_pending_coll_comm));
  p_comm->msg = msg;
  p_comm->token.data = p_comm;
  ezt_list_add(&process_info->pending_comm[comm_type_collective], &p_comm->token);
  return p_comm;
}

/* check if a message matches a collective communication
 * return 1 if the message matches or 0 otherwise
 */
static int __mpi_coll_msg_match(struct mpi_coll_msg_t* msg,
				enum coll_type_t type,
				struct ezt_mpi_comm *comm,
				struct mpi_request *req,
				int time_id) {
  int my_rank = comm->my_rank;

  if ((msg->type == type)
      && (msg->comm_size == comm->comm_size)
      && ((!msg->requests[my_rank]) || (msg->requests[my_rank] == req))
      && (time_id < 0 || !IS_TIME_SET(msg->times[my_rank][time_id]))) {
    /* It looks like this message matches.
     * We need to check wether the communicator for this message is comm or
     * if it matches comm
     */

    int not_this_one = 0;
    int i;
    for (i = 0; i < msg->comm_size; i++) {
      if (msg->comm[i]) {
	/* the ith communicator is not empty */

	if(! __ezt_mpi_communicator_are_equal(msg->comm[i], comm)) {
	  /* the communicator doesn't match. we need to find anothe rmessahe */
	  not_this_one = 1;
	} else {
	  /* the communicator matches */
	  return 1;
	}
      }
    }

    if (!not_this_one) {
      /* all the processes in this collective belong to comm */
      return 1;
    }
  }

  return 0;
}

/* search for a message in the list of pending collective messages */
static __attribute__ ((unused))
struct mpi_pending_coll_comm *
__mpi_find_pending_coll_message(struct mpi_process_info_t*p_info,
				enum coll_type_t type,
				struct ezt_mpi_comm *comm,
				struct mpi_request *req,
				int time_id) {
  struct ezt_list_token_t *token;
  ezt_list_foreach(&p_info->pending_comm[comm_type_collective], token) {
    struct mpi_pending_coll_comm* p_comm = (struct mpi_pending_coll_comm*) token->data;
    struct mpi_coll_msg_t* msg = p_comm->msg;
    if( __mpi_coll_msg_match(msg, type, comm, req, time_id)) {
      return p_comm;
    }
  }
  return NULL;
}


/* search for a message in the list of pending collective messages
 * the returned message has already been matched by the current thread
 */
static struct mpi_pending_coll_comm *
__mpi_find_matched_coll_message(struct mpi_process_info_t*p_info,
				enum coll_type_t type,
				struct ezt_mpi_comm *comm,
				struct mpi_request *req,
				int time_id) {
  struct ezt_list_token_t *token;
  ezt_list_foreach(&p_info->pending_comm[comm_type_collective], token) {
    struct mpi_pending_coll_comm* p_comm = (struct mpi_pending_coll_comm*) token->data;
    struct mpi_coll_msg_t* msg = p_comm->msg;
    if( __mpi_coll_msg_match(msg, type, comm, req, time_id)) {
      if(msg->comm[comm->my_rank] == comm) {
	return p_comm;
      }
    }
  }
  return NULL;
}

/* search for a collective communication message message that matches */
struct mpi_coll_msg_t *
__mpi_find_coll_message(enum coll_type_t type, struct ezt_mpi_comm *comm,
                        struct mpi_request *req, int time_id)
{
  int rank;
  struct ezt_list_token_t *token;

  for(rank=0; rank <comm->comm_size; rank++) {
    /* search for a matching collective message in the processes that belong to
     * the communicator
     */
    INIT_MPI_PROCESS_INFO(GET_PROCESS_INFO(comm->ranks[rank]), p_info);
    ezt_list_foreach(&p_info->pending_comm[comm_type_collective], token) {
      struct mpi_pending_coll_comm* p_comm = (struct mpi_pending_coll_comm*) token->data;
      struct mpi_coll_msg_t* msg = p_comm->msg;
      /* browse the list of pending collective communications in this process */

      if( __mpi_coll_msg_match(msg, type, comm, req, time_id)) {
	return msg;
      }
    }
  }

  /* cannot find a matching collective */
  return NULL;
}

static int __get_local_rank(int global_rank,
		   struct  mpi_coll_msg_t *msg) {
  int i;
  for(i=0; i<msg->comm_size; i++) {
    struct ezt_mpi_comm*comm = msg->comm[i];
    if(is_comm_mine(global_rank, comm))
      return comm->my_rank;
  }
  return -1;
}

struct mpi_coll_msg_t*
__mpi_find_coll_message_by_mpi_req(int rank,
				   struct mpi_request *mpi_req) {
  struct ezt_list_token_t *token;
  INIT_MPI_PROCESS_INFO(GET_PROCESS_INFO(rank), p_info);
  ezt_list_foreach(&p_info->pending_comm[comm_type_collective], token) {
    struct mpi_pending_coll_comm* p_comm = (struct mpi_pending_coll_comm*) token->data;
    struct mpi_coll_msg_t* msg = p_comm->msg;
    /* browse the list of pending collective communications in this process */
    int local_rank = __get_local_rank(rank, msg);
    if(mpi_req ==  msg->requests[local_rank])
      return msg;
  }
  return NULL;
}


struct mpi_coll_msg_t*
__enter_coll(uint64_t time,
	     enum coll_type_t type,
	     struct ezt_mpi_comm *comm,
	     int __attribute__((unused)) my_rank,
	     int data_size,
	     struct mpi_request* req,
	     char* thread_id)
{
  struct mpi_coll_msg_t* msg = __mpi_find_coll_message(type, comm, NULL,
                                                       start_coll);
  if ( !msg) {
    /* We are the first process to arrive in this collective communication.
     * create a new collective message structure
     */
    msg = __new_coll_message(type, comm, data_size);
  }

  /* add this message in the list of the process pending communications */
  __create_new_pending_collective_comm(comm->ranks[my_rank], msg);

  assert(msg->comm[comm->my_rank] == NULL);

  /* todo: record_event(time, type, msg) */

  msg->nb_started++;
  msg->comm[comm->my_rank] = comm;
  assert(msg->comm[comm->my_rank] == comm);

  msg->times[comm->my_rank][start_coll] = time;
  msg->thread_ids[comm->my_rank] = thread_id;
  if (req) {
    msg->requests[comm->my_rank] = req;
    req->coll_msg = msg;
    req->status = mpi_req_status_pending;
  }

  return msg;
}

static void
__store_finished_collective(p_eztrace_container p_cont,
			    uint64_t time,
			    int my_rank,
			    struct mpi_coll_msg_t* p_msg) {
  struct coll_msg_event* msg = hierarchical_array_new_item(p_cont, EZTRACE_MPI_STATS_COLL_MSG_ID);
  msg->time = time;
  msg->my_rank = my_rank;
  msg->msg = p_msg;
}

struct mpi_coll_msg_t*
__leave_coll(uint64_t time,
	     enum coll_type_t type,
	     struct ezt_mpi_comm *comm,
	     int __attribute__((unused)) my_rank,
	     struct mpi_request* req,
	     char* thread_id)
{
  INIT_MPI_PROCESS_INFO(GET_PROCESS_INFO(comm->ranks[comm->my_rank]), p_info);

  struct mpi_pending_coll_comm *pending_comm =
    __mpi_find_matched_coll_message(p_info,
				    type,
				    comm,
				    req,
				    stop_coll);

  assert(pending_comm);
  struct mpi_coll_msg_t* msg = pending_comm->msg;

  assert(msg->comm[comm->my_rank] == comm);

  /* todo: record_event(time, type, msg) */
  msg->times[comm->my_rank][stop_coll] = time;

  assert(msg);
  assert(msg->thread_ids[comm->my_rank] == thread_id);
  assert(msg->requests[comm->my_rank] == req);
  assert(msg->comm[comm->my_rank] == comm);

  if(msg->nb_started != comm->comm_size) {
    msg->times[comm->my_rank][stop_coll] = TIME_INIT;
    /* some process did not reach the start of the collective. We need to wait
     * for them
     */
    return NULL;
  }

  /* remove the collective from the list of pending collectives */
  ezt_list_remove(&pending_comm->token);

  int nb_completed=0;
  for(int i=0; i<comm->comm_size; i++) {
    if(msg->times[i][stop_coll] != TIME_INIT)
      nb_completed++;
  }
  if(nb_completed == comm->comm_size) {

    free(pending_comm);
    struct mpi_pending_coll_comm *finished_comm = malloc(sizeof(struct mpi_pending_coll_comm));
    finished_comm->msg = msg;
    finished_comm->token.data = finished_comm;

    ezt_list_add(&finished_collectives[type], &finished_comm->token);
  } else {
    free(pending_comm);
  }

  __store_finished_collective(GET_PROCESS_CONTAINER(comm->ranks[comm->my_rank]),
			      time,
			      comm->my_rank,
			      msg);

  return msg;
}
