/* (c) 2005 Tomas Plachetka */

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/time.h>

#include "tpl.h"

/*****************************************************************************
* Local preprocessor defines
******************************************************************************/

#define DEBUG(x)

/*****************************************************************************
* Local typedefs
******************************************************************************/

enum 
{
  MSG_INIT = 0,
  MSG_PING,
  MSG_PONG,
  MSG_QUIT
};

enum
{
  RANK_PING = 0,
  RANK_PONG
};

/*****************************************************************************
* Local variables
******************************************************************************/

static int nr_msgs;
static int msg_size;

static int rank_ping = RANK_PING;
static int rank_pong = RANK_PONG;
static  int my_rank;


/*****************************************************************************
* Static functions
******************************************************************************/

static int match_all(int sender, int tag, void *message);
static void ping();
static void pong();
static void *ping_listener(void *arg);
static void *pong_listener(void *arg);

static int match_all(int sender, int tag, void *message)
{
  return(1);
}

static void *ping_listener(void *arg)
{
  int counter;
  int sender;
  int tag;
  void *message;
  char *databuf;
  
  if (msg_size > 0)
  {
    if ((databuf = (char *) malloc(msg_size * sizeof(char))) == NULL)
    {
      printf("Could not allocate memory\n");
      exit(1);
    }
  }
  
  for (counter = 0; counter < nr_msgs; counter++)
  {
    DEBUG(printf("Receiving %d\n", my_rank);fflush(stdout);)
    tpl_recv(match_all, &sender, &tag, &message);
    tpl_upkchar(message, databuf, msg_size);
    tpl_destroy(message);
    DEBUG(printf("Received %d\n", my_rank);fflush(stdout);)
  }
  
  if (msg_size > 0)
  {
    free(databuf);
  }
}

static void ping(void)
{
  int counter;
  void *message;
  char *databuf;
  pthread_t thr_ping;
  struct timeval tv1, tv2;
  struct timezone tz1, tz2;
  float elapsed_time, mbytes;

  if (msg_size > 0)
  {
    if ((databuf = (char *) malloc(msg_size * sizeof(char))) == NULL)
    {
      printf("Could not allocate memory\n");
      exit(1);
    }
  }
  
  /* Start timer. */
  gettimeofday(&tv1, &tz1);
    
  /* Start receiving thread. */
  pthread_create(&thr_ping, NULL, ping_listener, NULL);
      
  for (counter = 0; counter < nr_msgs; counter++)
  {
    tpl_create(&message);
    DEBUG(printf("Sending %d\n", my_rank);fflush(stdout);)
    tpl_pkchar(message, databuf, msg_size);
    tpl_send(&rank_pong, 1, MSG_PING, message);
    DEBUG(printf("Sent %d\n", my_rank);fflush(stdout);)
  }

  if (msg_size > 0)
  {
    free(databuf);
  }

  DEBUG(printf("Joining %d\n", my_rank);fflush(stdout);)  
  pthread_join(thr_ping, NULL);
  DEBUG(printf("Joined %d\n", my_rank);fflush(stdout);)
  
  /* Stop timer. */
  gettimeofday(&tv2, &tz2);

  elapsed_time = tv2.tv_sec - tv1.tv_sec +
   (tv2.tv_usec - tv1.tv_usec) / 1e6;
        
  mbytes = 2.0 * (float) nr_msgs * (float) msg_size / 1048576.0;
  printf("%d x %d Bytes x 2 = %f MBytes, %f s, %f MB/s, %f msg/s, %f s/msg\n", 
   nr_msgs, msg_size, mbytes, elapsed_time, mbytes / elapsed_time, 
   nr_msgs / elapsed_time, elapsed_time / nr_msgs);
}

static void *pong_listener(void *arg)
{
  int counter;
  int sender;
  int tag;
  char *databuf;
  void *message;

  if (msg_size > 0)
  {
    if ((databuf = (char *) malloc(msg_size * sizeof(char))) == NULL)
    {
      printf("Could not allocate memory\n");
      exit(1);
    }
  }

  for (counter = 0; counter < nr_msgs; counter++)
  {
    DEBUG(printf("Receiving %d\n", my_rank);fflush(stdout);)
    tpl_recv(match_all, &sender, &tag, &message);
    tpl_upkchar(message, databuf, msg_size);
    tpl_destroy(message);
    DEBUG(printf("Received %d\n", my_rank);fflush(stdout);)
  }

  if (msg_size > 0)
  {
    free(databuf);
  }
}

static void pong(void)
{
  int counter;
  char *databuf;
  void *sendbuf;
  int tag;
  int sender;
  pthread_t thr_pong;
  
  /* Start receiving thread. */
  pthread_create(&thr_pong, NULL, pong_listener, NULL);

  if (msg_size > 0)
  {
    if ((databuf = (char *) malloc(msg_size * sizeof(char))) == NULL)
    {
      printf("Could not allocate memory\n");
      exit(1);
    }
  }
  
  for (counter = 0; counter < nr_msgs; counter++)
  {
    tpl_create(&sendbuf);
    DEBUG(printf("Sending %d\n", my_rank);fflush(stdout);)
    tpl_pkchar(sendbuf, databuf, msg_size);
    tpl_send(&rank_ping, 1, MSG_PONG, sendbuf);
    DEBUG(printf("Sent %d\n", my_rank);fflush(stdout);)
  }

  if (msg_size > 0)
  {
    free(databuf);
  }

  DEBUG(printf("Joining %d\n", my_rank);fflush(stdout);)  
  pthread_join(thr_pong, NULL);
  DEBUG(printf("Joined %d\n", my_rank);fflush(stdout);)
}

int main(int argc, char **argv)
{
  int counter;
  int nr_procs = 2;

  tpl_initialize(&nr_procs, &my_rank, &argc, &argv);

  if (argc != 3)
  {
    printf("Usage: ping <nr_msgs> <msg_size>\n");
    tpl_deinitialize();
    exit(1);
  }

  nr_msgs = atoi(argv[1]);
  msg_size = atoi(argv[2]);

  switch(my_rank)
  {
    case RANK_PING:
      ping();
      break;
    
    case RANK_PONG:
      pong();
      break;
    
    default:
      printf("Unknown role\n");
      exit(1);
      break;
  }

  /* Terminate the program. */
  tpl_deinitialize();

  return(0);
}
