#define VERSION "0.1"

#include <math.h>
#include <complex.h>
#include <fftw3.h>
#include <signal.h>
#include <string.h>
#include <stdlib.h>
#include <stdbool.h>
#include <stdint.h>
#include <errno.h>

#include "constants.cpp"
#include "driving.h"
#include "init.h"
#include "int_tools.h"
#include "io.h"
#include "navier-stokes.h"


// structure to store parameters, to make it easier to pass parameters to CLI functions
typedef struct nstrophy_parameters {
  double init_en; // initial value for the energy for ins and enstrophy for rns
  bool irreversible;
  int K1;
  int K2;
  int N1;
  int N2;
  unsigned int nsteps;
  double nu;
  double delta;
  double L;
  uint64_t print_freq;
  int seed;
  uint64_t starting_time;
  unsigned int driving;
  unsigned int init;
  unsigned int algorithm;
  FILE* initfile;
  FILE* drivingfile;
} nstrophy_parameters;

// usage message
int print_usage();
// print parameters
int print_params(nstrophy_parameters parameters, char* initfile_str, char* drivingfile_str, FILE* file);

// read command line arguments
int read_args(int argc, const char* argv[], char** params, unsigned int* command, unsigned int* nthreads, char** savefile_str);
int read_params(char* param_str, nstrophy_parameters* parameters, char** initfile_str, char** drivingfile_str);
int set_parameter(char* lhs, char* rhs, nstrophy_parameters* parameters, bool* setN1, bool* setN2, char** initfile_str, char** drivingfile_str);

// set driving force
_Complex double* set_driving(nstrophy_parameters parameters);
// set initial condition
_Complex double* set_init(nstrophy_parameters parameters);

// signal handler
void sig_handler (int signo);

// global variable to handle interrupts
volatile bool g_abort = false;
// signal handler
void sig_handler (int signo){
  if (signo == SIGINT){
    g_abort = true;
  }
}


int main (
  int argc,
  const char* argv[]
){
  char* param_str=NULL;
  nstrophy_parameters parameters;
  int ret;
  unsigned int command;
  unsigned int nthreads=1;
  _Complex double* u0;
  _Complex double *g;
  char* savefile_str=NULL;
  char* initfile_str=NULL;
  char* drivingfile_str=NULL;
  FILE* savefile=NULL;

  command=0;

  // read command line arguments
  ret=read_args(argc, argv, &param_str, &command, &nthreads, &savefile_str);
  if(ret<0){
    return(-1);
  }

  // read params
  ret=read_params(param_str, &parameters, &initfile_str, &drivingfile_str);
  if(ret<0){
    return(-1);
  }

  // open initfile
  if(initfile_str!=NULL){
    parameters.initfile=fopen(initfile_str,"r");
    if(parameters.initfile==NULL){
      fprintf(stderr,"Error opening file '%s' for reading: %s\n", initfile_str, strerror(errno));
      return(-1);
    }
  }
  // open drivingfile
  if(drivingfile_str!=NULL){
    parameters.drivingfile=fopen(drivingfile_str,"r");
    if(parameters.drivingfile==NULL){
      fprintf(stderr,"Error opening file '%s' for reading: %s\n", drivingfile_str, strerror(errno));
      return(-1);
    }
  }

  // set driving force
  g=set_driving(parameters);
  // set initial condition
  u0=set_init(parameters);

  // close initfile (do this early, so that it is possible to use the same file for init and save)
  if (parameters.initfile!=NULL){
    fclose(parameters.initfile);
  }
  // close drivingfile
  if (parameters.drivingfile!=NULL){
    fclose(parameters.drivingfile);
  }

  // open savefile (do this after closing init file)
  if(savefile_str!=NULL){
    savefile=fopen(savefile_str,"w");
    if(savefile==NULL){
      fprintf(stderr,"Error opening file '%s' for writing: %s\n", savefile_str, strerror(errno));
      return(-1);
    }
  }

  // print parameters
  print_params(parameters, initfile_str, drivingfile_str, stderr);
  print_params(parameters, initfile_str, drivingfile_str, stdout);

  // free initfile_str
  if(initfile_str!=NULL){
    free(initfile_str);
  }
  // free drivingfile_str
  if(drivingfile_str!=NULL){
    free(drivingfile_str);
  }

  // run command
  if (command==COMMAND_UK){
    uk(parameters.K1, parameters.K2, parameters.N1, parameters.N2, parameters.nsteps, parameters.nu, parameters.delta, parameters.L, u0, g, parameters.irreversible, parameters.algorithm, parameters.print_freq, parameters.starting_time, nthreads, savefile);
  }
  else if(command==COMMAND_ENSTROPHY){
    // register signal handler to handle aborts
    signal(SIGINT, sig_handler);
    enstrophy(parameters.K1, parameters.K2, parameters.N1, parameters.N2, parameters.nsteps, parameters.nu, parameters.delta, parameters.L, u0, g, parameters.irreversible, parameters.algorithm, parameters.print_freq, parameters.starting_time, nthreads, savefile, (char*)argv[0], param_str, savefile_str);
  }
  else if(command==COMMAND_QUIET){
    quiet(parameters.K1, parameters.K2, parameters.N1, parameters.N2, parameters.nsteps, parameters.nu, parameters.delta, parameters.L, parameters.starting_time, u0, g, parameters.irreversible, parameters.algorithm, nthreads, savefile);
  }
  else if(command==0){
    fprintf(stderr, "error: no command specified\n");
    print_usage();
  }

  free(g);
  free(u0);

  // close savefile
  if (savefile!=NULL){
    fclose(savefile);
  }

  return(0);
}

// usage message
int print_usage(){
  fprintf(stderr, "usage:\n       nstrophy [-t nthreads] [-p parameters] [-s savefile] <command>\n\n");
  return(0);
}

// print parameters
int print_params(
  nstrophy_parameters parameters,
  char* initfile_str,
  char* drivingfile_str,
  FILE* file
){
  fprintf(file,"# ");

  if (parameters.irreversible){
    fprintf(file,"equation=irreversible");
  } else {
    fprintf(file,"equation=reversible");
  }

  fprintf(file,", K1=%d, K2=%d, N1=%d, N2=%d, nu=%.15e, delta=%.15e, L=%.15e, init_en=%.15e", parameters.K1, parameters.K2, parameters.N1, parameters.N2, parameters.nu, parameters.delta, parameters.L, parameters.init_en);

  switch(parameters.driving){
  case DRIVING_TEST:
    fprintf(file,", driving=test");
    break;
  case DRIVING_ZERO:
    fprintf(file,", driving=zero");
    break;
  case DRIVING_FILE:
    fprintf(file,", driving=file:%s", drivingfile_str);
    break;
  case DRIVING_FILE_TXT:
    fprintf(file,", driving=file_txt:%s", drivingfile_str);
    break;
  default:
    fprintf(file,", driving=test");
    break;
  }

  switch(parameters.init){
  case INIT_RANDOM:
    fprintf(file,", init=random");
    break;
  case INIT_GAUSSIAN:
    fprintf(file,", init=gaussian");
    break;
  case INIT_FILE:
    fprintf(file,", init=file:%s", initfile_str);
    break;
  case INIT_FILE_TXT:
    fprintf(file,", init=file_txt:%s", initfile_str);
    break;
  default:
    fprintf(file,", init=gaussian");
    break;
  }

  switch(parameters.algorithm){
  case ALGORITHM_RK4:
    fprintf(file,", algorithm=RK4");
    break;
  case ALGORITHM_RK2:
    fprintf(file,", algorithm=RK2");
    break;
  default:
    fprintf(file,", algorithm=RK4");
    break;
  }

  fprintf(file,"\n");

  return 0;
}


// read command line arguments
#define CP_FLAG_PARAMS 1
#define CP_FLAG_NTHREADS 2
#define CP_FLAG_SAVEFILE 3
int read_args(
  int argc,
  const char* argv[],
  char** params,
  unsigned int* command,
  unsigned int* nthreads,
  char** savefile_str
){
  int i;
  int ret;
  // pointers
  char* ptr;
  // flag that indicates what argument is being read
  int flag=0;

  // loop over arguments
  for(i=1;i<argc;i++){
    // flag
    if(argv[i][0]=='-'){
      for(ptr=((char*)argv[i])+1;*ptr!='\0';ptr++){
	switch(*ptr){
	case 'p':
	  flag=CP_FLAG_PARAMS;
	  break;
	case 't':
	  flag=CP_FLAG_NTHREADS;
	  break;
	case 's':
	  flag=CP_FLAG_SAVEFILE;
	  break;
	default:
	  fprintf(stderr, "unrecognized option '-%c'\n", *ptr);
	  print_usage();
	  return(-1);
	  break;
	}
      }
    }
    // params
    else if(flag==CP_FLAG_PARAMS){
      *params=(char*)argv[i];
      flag=0;
    }
    // nthreads
    else if(flag==CP_FLAG_NTHREADS){
      ret=sscanf(argv[i],"%u",nthreads);
      if(ret!=1){
	fprintf(stderr, "error: '-t' should be followed by an unsigned integer\n       got '%s'\n",argv[i]);
	return(-1);
      }
      flag=0;
    }
    // savefile
    else if(flag==CP_FLAG_SAVEFILE){
      *savefile_str=(char*)argv[i];
      flag=0;
    }
    // computation to run
    else{
      if(strcmp(argv[i], "uk")==0){
	*command=COMMAND_UK;
      }
      else if(strcmp(argv[i], "enstrophy")==0){
	*command=COMMAND_ENSTROPHY;
      }
      else if(strcmp(argv[i], "quiet")==0){
	*command=COMMAND_QUIET;
      }
      else{
	fprintf(stderr, "error: unrecognized command: '%s'\n",argv[i]);
	return(-1);
      }
      flag=0;
    }
  }

  return(0);
}

// read parameters string
int read_params(
  char* param_str,
  nstrophy_parameters* parameters,
  char** initfile_str,
  char** drivingfile_str
){
  int ret;
  // pointer in params
  char* ptr;
  // buffer and associated pointer
  char *buffer_lhs, *lhs_ptr;
  char *buffer_rhs, *rhs_ptr;
  // whether N was set explicitly
  bool setN1=false;
  bool setN2=false;
  // whether lhs (false is rhs)
  bool lhs=true;

  // defaults
  parameters->init_en=1.54511597324389e+02;
  parameters->irreversible=true;
  parameters->K1=16;
  parameters->K2=parameters->K1;
  //delta=2^-13
  parameters->delta=0.0001220703125;
  //nu=2^-11
  parameters->nu=0.00048828125;
  parameters->L=2*M_PI;
  parameters->nsteps=10000000;
  parameters->print_freq=1000;
  parameters->starting_time=0;
  parameters->seed=17;
  parameters->driving=DRIVING_TEST;
  parameters->drivingfile=NULL;
  parameters->init=INIT_GAUSSIAN;
  parameters->initfile=NULL;
  parameters->algorithm=ALGORITHM_RK4;

  if (param_str!=NULL){
    // init
    buffer_lhs=calloc(sizeof(char),strlen(param_str));
    lhs_ptr=buffer_lhs;
    *lhs_ptr='\0';
    buffer_rhs=calloc(sizeof(char),strlen(param_str));
    rhs_ptr=buffer_rhs;
    *rhs_ptr='\0';

    for(ptr=param_str;*ptr!='\0';ptr++){
      switch(*ptr){
      case '=':
	// reset buffer
	rhs_ptr=buffer_rhs;
	*rhs_ptr='\0';
	lhs=false;
	break;
      case ';':
	//set parameter
	ret=set_parameter(buffer_lhs, buffer_rhs, parameters, &setN1, &setN2, initfile_str, drivingfile_str);
	if(ret<0){
	  return ret;
	}
	// reset buffer
	lhs_ptr=buffer_lhs;
	*lhs_ptr='\0';
	lhs=true;
	break;
      default:
	// add to buffer
	if (lhs){
	  *lhs_ptr=*ptr;
	  lhs_ptr++;
	  *lhs_ptr='\0';
	}
	else{
	  *rhs_ptr=*ptr;
	  rhs_ptr++;
	  *rhs_ptr='\0';
	}
	break;
      }
    }

    // set last param
    if (*param_str!='\0'){
      ret=set_parameter(buffer_lhs, buffer_rhs, parameters, &setN1, &setN2, initfile_str, drivingfile_str);
      if(ret<0){
	return ret;
      }
    }

    // free vects
    free(buffer_lhs);
    free(buffer_rhs);
  }

  // if N not set explicitly, set it to the smallest power of 2 that is >3*K+1 (the fft is faster on vectors whose length is a power of 2)
  if (!setN1){
    parameters->N1=smallest_pow2(3*(parameters->K1));
  }
  if (!setN2){
    parameters->N2=smallest_pow2(3*(parameters->K2));
  }

  return(0);
}


// set a parameter from the parameter string
int set_parameter(
  char* lhs,
  char* rhs,
  nstrophy_parameters* parameters,
  bool* setN1,
  bool* setN2,
  char** initfile_str,
  char** drivingfile_str
){
  int ret;

  if (strcmp(lhs,"equation")==0){
    if (strcmp(rhs,"irreversible")==0){
      parameters->irreversible=true;
    }
    else if (strcmp(rhs,"reversible")==0){
      parameters->irreversible=false;
    }
    else {
      fprintf(stderr, "error: 'equation' should be 'irreversible' or 'reversible'\n       got '%s'\n",rhs);
      return(-1);
    }
  }
  else if (strcmp(lhs,"init_en")==0){
    ret=sscanf(rhs,"%lf",&(parameters->init_en));
    if(ret!=1){
      fprintf(stderr, "error: parameter 'init_en' should be a double\n       got '%s'\n",rhs);
      return(-1);
    }
  }
  else if (strcmp(lhs,"K1")==0){
    ret=sscanf(rhs,"%d",&(parameters->K1));
    if(ret!=1){
      fprintf(stderr, "error: parameter 'K1' should be an integer\n       got '%s'\n",rhs);
      return(-1);
    }
  }
  else if (strcmp(lhs,"K2")==0){
    ret=sscanf(rhs,"%d",&(parameters->K2));
    if(ret!=1){
      fprintf(stderr, "error: parameter 'K2' should be an integer\n       got '%s'\n",rhs);
      return(-1);
    }
  }
  else if (strcmp(lhs,"K")==0){
    ret=sscanf(rhs,"%d",&(parameters->K1));
    if(ret!=1){
      fprintf(stderr, "error: parameter 'K' should be an integer\n       got '%s'\n",rhs);
      return(-1);
    }
    parameters->K2=parameters->K1;
  }
  else if (strcmp(lhs,"N1")==0){
    ret=sscanf(rhs,"%d",&(parameters->N1));
    if(ret!=1){
      fprintf(stderr, "error: parameter 'N1' should be an integer\n       got '%s'\n",rhs);
      return(-1);
    }
    *setN1=true;
  }
  else if (strcmp(lhs,"N2")==0){
    ret=sscanf(rhs,"%d",&(parameters->N2));
    if(ret!=1){
      fprintf(stderr, "error: parameter 'N2' should be an integer\n       got '%s'\n",rhs);
      return(-1);
    }
    *setN2=true;
  }
  else if (strcmp(lhs,"N")==0){
    ret=sscanf(rhs,"%d",&(parameters->N1));
    if(ret!=1){
      fprintf(stderr, "error: parameter 'N' should be an integer\n       got '%s'\n",rhs);
      return(-1);
    }
    parameters->N2=parameters->N1;
    *setN1=true;
    *setN2=true;
  }
  else if (strcmp(lhs,"nsteps")==0){
    ret=sscanf(rhs,"%u",&(parameters->nsteps));
    if(ret!=1){
      fprintf(stderr, "error: parameter 'nsteps' should be an unsigned integer\n       got '%s'\n",rhs);
      return(-1);
    }
  }
  else if (strcmp(lhs,"nu")==0){
    ret=sscanf(rhs,"%lf",&(parameters->nu));
    if(ret!=1){
      fprintf(stderr, "error: parameter 'nu' should be a double\n       got '%s'\n",rhs);
      return(-1);
    }
  }
  else if (strcmp(lhs,"delta")==0){
    ret=sscanf(rhs,"%lf",&(parameters->delta));
    if(ret!=1){
      fprintf(stderr, "error: parameter 'delta' should be a double\n       got '%s'\n",rhs);
      return(-1);
    }
  }
  else if (strcmp(lhs,"L")==0){
    ret=sscanf(rhs,"%lf",&(parameters->L));
    if(ret!=1){
      fprintf(stderr, "error: parameter 'L' should be a double\n       got '%s'\n",rhs);
      return(-1);
    }
  }
  else if (strcmp(lhs,"print_freq")==0){
    ret=sscanf(rhs,"%lu",&(parameters->print_freq));
    if(ret!=1){
      fprintf(stderr, "error: parameter 'print_freq' should be an unsigned integer\n       got '%s'\n",rhs);
      return(-1);
    }
  }
  else if (strcmp(lhs,"random_seed")==0){
    ret=sscanf(rhs,"%d",&(parameters->seed));
    if(ret!=1){
      fprintf(stderr, "error: parameter 'random_seed' should be an integer\n       got '%s'\n",rhs);
      return(-1);
    }
  }
  else if (strcmp(lhs,"starting_time")==0){
    ret=sscanf(rhs,"%lu",&(parameters->starting_time));
    if(ret!=1){
      fprintf(stderr, "error: parameter 'starting_time' should be an unsigned integer\n       got '%s'\n",rhs);
      return(-1);
    }
  }
  else if (strcmp(lhs,"driving")==0){
    if (strcmp(rhs,"zero")==0){
      parameters->driving=DRIVING_ZERO;
    }
    else if (strcmp(rhs,"test")==0){
      parameters->driving=DRIVING_TEST;
    }
    // matches any argument that starts with 'file:'
    else if (strncmp(rhs,"file:",5)==0){
      parameters->driving=DRIVING_FILE;
      *drivingfile_str=calloc(sizeof(char), strlen(rhs)-5+1);
      strcpy(*drivingfile_str, rhs+5);
    }
    // matches any argument that starts with 'file_txt:'
    else if (strncmp(rhs,"file_txt:",9)==0){
      parameters->driving=DRIVING_FILE_TXT;
      *drivingfile_str=calloc(sizeof(char), strlen(rhs)-9+1);
      strcpy(*drivingfile_str, rhs+9);
    }
    else{
      fprintf(stderr, "error: unrecognized driving force '%s'\n",rhs);
      return(-1);
    }
  }
  // initial condition
  else if (strcmp(lhs,"init")==0){
    if (strcmp(rhs,"random")==0){
      parameters->init=INIT_RANDOM;
    }
    else if (strcmp(rhs,"gaussian")==0){
      parameters->init=INIT_GAUSSIAN;
    }
    // matches any argument that starts with 'file:'
    else if (strncmp(rhs,"file:",5)==0){
      parameters->init=INIT_FILE;
      *initfile_str=calloc(sizeof(char), strlen(rhs)-5+1);
      strcpy(*initfile_str, rhs+5);
    }
    // matches any argument that starts with 'file_txt:'
    else if (strncmp(rhs,"file_txt:",9)==0){
      parameters->init=INIT_FILE_TXT;
      *initfile_str=calloc(sizeof(char), strlen(rhs)-9+1);
      strcpy(*initfile_str, rhs+9);
    }
    else{
      fprintf(stderr, "error: unrecognized initial condition '%s'\n",rhs);
      return(-1);
    }
  }
  // algorithm
  else if (strcmp(lhs,"algorithm")==0){
    if (strcmp(rhs,"RK4")==0){
      parameters->algorithm=ALGORITHM_RK4;
    }
    else if (strcmp(rhs,"RK2")==0){
      parameters->algorithm=ALGORITHM_RK2;
    }
    else{
      fprintf(stderr, "error: unrecognized algorithm '%s'\n",rhs);
      return(-1);
    }
  }
  else{
    fprintf(stderr, "error: unrecognized parameter '%s'\n",lhs);
    return(-1);
  }

  return(0);
}

// set driving force
_Complex double* set_driving(
  nstrophy_parameters parameters
){
  _Complex double* g=calloc(sizeof(_Complex double),parameters.K1*(2*parameters.K2+1)+parameters.K2);

  switch(parameters.driving){
  case DRIVING_ZERO:
    g_zero(g, parameters.K1, parameters.K2);
    break;
  case DRIVING_TEST:
    g_test(g, parameters.K1, parameters.K2);
    break;

  case DRIVING_FILE:
    init_file_bin(g, parameters.K1, parameters.K2, parameters.drivingfile);
    break;

  case DRIVING_FILE_TXT:
    init_file_txt(g, parameters.K1, parameters.K2, parameters.drivingfile);
    break;

  default:
    g_test(g, parameters.K1, parameters.K2);
    break;
  }
   return g;
}

// set initial condition
_Complex double* set_init(
  nstrophy_parameters parameters
){
  _Complex double* u0=calloc(sizeof(_Complex double),parameters.K1*(2*parameters.K2+1)+parameters.K2);

  switch(parameters.init){
  case INIT_RANDOM:
    init_random(u0, parameters.init_en, parameters.K1, parameters.K2, parameters.L, parameters.seed, parameters.irreversible);
    break;

  case INIT_GAUSSIAN:
    init_gaussian(u0, parameters.init_en, parameters.K1, parameters.K2, parameters.L, parameters.irreversible);
    break;

  case INIT_FILE:
    init_file_bin(u0, parameters.K1, parameters.K2, parameters.initfile);
    break;

  case INIT_FILE_TXT:
    init_file_txt(u0, parameters.K1, parameters.K2, parameters.initfile);
    break;

  default:
    init_gaussian(u0, parameters.init_en, parameters.K1, parameters.K2, parameters.L, parameters.irreversible);
    break;
  }
  
  return u0;
}