#include <getopt.h>
#include <mpi.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>

#include "configuration.h"

#define MAX_ITERATIONS 100

const char *verbosity_levels[] = {"OFF", "INFO", "DEBUG", "TRACE"};

const char *communication_computation_modes[] = {
   "No overlap of communication and computation",
   "Overlapping communication and computation"
};

const char *transmission_modes[] = {
   "Sparse collective - MPI_Dist_graph_create_adjacent / MPI_[N/In]eighbor_alltoallw",
   "Point-to-point - MPI_ISend / MPI_IRecv / MPI_Waitall",
   "Persistent request - MPI_Send_init / MPI_Recv_init / MPI_Startall / MPI_Waitall"
};

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

void conf_set_or_validate_nprocs(conf_t *c);

long parse_long(const char *str);

int log_enabled(const conf_t *c, int lvl);

void chop_wi_extension(char *filename);

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

void conf_init_default(conf_t *c)
{
   c->verbosity_level = INFO;
   c->transmission_mode = SPARSE_COLLECTIVE;
   c->communication_computation_mode = OVERLAP;

   c->nprocs[0] = 0;
   c->nprocs[1] = 0;

   // Default values:
   c->n_iterations = 1;
   c->n_generations_per_iteration = 5000;

   c->file_basename[0] = '\0';
}

void conf_init_from_args(conf_t *c, int argc, char* argv[])
{
   struct option long_options[] = {
      {"verbose", required_argument, NULL, 'v'},

      // Transmission modes
      {"sparse",  no_argument, &c->transmission_mode, SPARSE_COLLECTIVE},
      {"p2p",     no_argument, &c->transmission_mode, POINT_TO_POINT},
      {"persist", no_argument, &c->transmission_mode, PERSISTENT_REQUEST},

      // Communication computation overlap
      {"no-overlap", no_argument, &c->communication_computation_mode, NO_OVERLAP},
      {"overlap",    no_argument, &c->communication_computation_mode, OVERLAP},

      {"nprocs-x", required_argument, NULL, 'x'},
      {"nprocs-y", required_argument, NULL, 'y'},

      {"iterations", required_argument, NULL, 'i'},
      {"generations-per-iteration", required_argument, NULL, 'g'},

      {NULL, 0, NULL, 0}
   };

   conf_init_default(c);

   int option;
   while((option = getopt_long(argc, argv, "v:x:y:i:g:", long_options, NULL)) >= 0) {
      switch(option) {
         case 'v': c->verbosity_level = atoi(optarg); break;
         case 'x': c->nprocs[0] = atoi(optarg); break;
         case 'y': c->nprocs[1] = atoi(optarg); break;
         case 'i': c->n_iterations = atoi(optarg); break;
         case 'g': c->n_generations_per_iteration = parse_long(optarg); break;
      }
   }

   if(c->n_iterations > MAX_ITERATIONS) {
      fprintf(stderr,
         "Specified number of iterations, %d, exceeds maximum of %d.\n"
         "(One output file will be generated for each iteration)\n",
         c->n_iterations, MAX_ITERATIONS
      );
      MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);
   }

   if(optind > argc) {
      fprintf(stderr, "Expected base name of input file (omitting '%s' file extension)\n", FILE_EXT);
      MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);
   }

   if(strlen(argv[optind]) > sizeof(c->file_basename)-1) {
      fprintf(stderr, "Input filename too long.\n");
      MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);
   }

   strncpy(c->file_basename, argv[optind], sizeof(c->file_basename));
   chop_wi_extension(c->file_basename);

   conf_set_or_validate_nprocs(c);
}

void conf_set_or_validate_nprocs(conf_t *c)
{
   int nprocs, rc;

   MPI_Comm_size(MPI_COMM_WORLD, &nprocs);
   rc = MPI_Dims_create(nprocs, 2, c->nprocs);

   if(rc != MPI_SUCCESS) {
      fprintf(stderr, "Cannot create 2D Cartesian grid with specified number of processes in x or y direction.\n");
      MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);
   }
}

void conf_print(const conf_t *c, FILE *f)
{
   int i = c->verbosity_level;
   if(i < 0) i = 0; else if(i > 3) i = 3;

   fprintf(f, "Configuration:\n");
   fprintf(f, " * Verbosity level:           %s (%d)\n", verbosity_levels[i], c->verbosity_level);
   fprintf(f, " * Input file:                %s%s\n", c->file_basename, FILE_EXT);
   fprintf(f, " * Transmission mode:         %s\n", transmission_modes[c->transmission_mode]);
   fprintf(f, " * Overlap mode:              %s communication and computation\n", c->communication_computation_mode == NO_OVERLAP ? "No overlap of" : "Overlapping");
   fprintf(f, " * Grid of processes:         %d x %d\n", c->nprocs[0], c->nprocs[1]);
   fprintf(f, " * Number of iterations:      %d\n", c->n_iterations);
   fprintf(f, " * Generations per iteration: %ld\n", c->n_generations_per_iteration);
   fprintf(f, "\n");
}

int log_enabled(const conf_t *c, int lvl)
{
   return c->verbosity_level >= lvl;
}

int info_enabled(const conf_t *c)
{
   static int rank = -1;
   if(rank < 0) {
      MPI_Comm_rank(MPI_COMM_WORLD, &rank);
   }

   return log_enabled(c, INFO) && rank == 0;
}

int debug_enabled(const conf_t *c)
{
   static int rank = -1;
   if(rank < 0) {
      MPI_Comm_rank(MPI_COMM_WORLD, &rank);
   }

   return log_enabled(c, DEBUG) && rank == 0;
}

int trace_enabled(const conf_t *c)
{
   return log_enabled(c, TRACE);
}

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

long parse_long(const char *str)
{
   char* p;
   long result = strtol(str, &p, 10);
   if(*p == 'k') {
      result *= 1000;
   } else if(*p == 'M') {
      result *= 1000000;
   }
   return result;
}

void chop_wi_extension(char *filename)
{
   char *extension = strrchr(filename, '.');
   if(extension != NULL) {
      if(strcmp(extension, FILE_EXT) == 0) {
         *extension = '\0';
      } else {
         fprintf(stderr,
            "Input file '%s' does not have expected file extension '%s'.\n",
            filename, FILE_EXT
         );
         MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);
      }
   }
}