#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 <errno.h>
#include "navier-stokes.h"
#include "driving.h"
#include "init.h"
#include "int_tools.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;
  unsigned int print_freq;
  unsigned int avg_window;
  int seed;
  unsigned int starting_time;
  unsigned int driving;
  unsigned int init;
  FILE* initfile;
} nstrophy_parameters;

// usage message
int print_usage();
// print parameters
int print_params(nstrophy_parameters parameters, char* initfile_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);
int set_parameter(char* lhs, char* rhs, nstrophy_parameters* parameters, bool* setN1, bool* setN2, bool* setavg_window, char** initfile_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);

#define COMMAND_UK 1
#define COMMAND_EEA 2
#define COMMAND_QUIET 3

#define DRIVING_ZERO 1
#define DRIVING_TEST 2

#define INIT_RANDOM 1
#define INIT_GAUSSIAN 2
#define INIT_FILE 3


// 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;
  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);
  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);
    }
  }

  // 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);
  }

  // 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, stderr);
  print_params(parameters, initfile_str, stdout);

  // 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.print_freq, parameters.starting_time, nthreads, savefile);
  }
  else if(command==COMMAND_EEA){
    // register signal handler to handle aborts
    signal(SIGINT, sig_handler);
    eea(parameters.K1, parameters.K2, parameters.N1, parameters.N2, parameters.nsteps, parameters.nu, parameters.delta, parameters.L, u0, g, parameters.irreversible, parameters.print_freq, parameters.avg_window, parameters.starting_time, nthreads, savefile);
  }
  else if(command==COMMAND_QUIET){
    quiet(parameters.K1, parameters.K2, parameters.N1, parameters.N2, parameters.nsteps, parameters.nu, parameters.delta, parameters.L, u0, g, parameters.irreversible, 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,
  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;
  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;
  default:
    fprintf(file,", init=gaussian");
    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], "energy")==0){
	*command=COMMAND_EEA;
      }
      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
){
  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 avg_window was set explicitly
  bool setavg_window=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->init=INIT_GAUSSIAN;
  parameters->initfile=NULL;

  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, &setavg_window, initfile_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, &setavg_window, initfile_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));
  }
  // if avg_window is not set explicitly, set it to print_freq
  if (!setavg_window){
    parameters->avg_window=parameters->print_freq;
  }

  return(0);
}


// set a parameter from the parameter string
int set_parameter(
  char* lhs,
  char* rhs,
  nstrophy_parameters* parameters,
  bool* setN1,
  bool* setN2,
  bool* setavg_window,
  char** initfile_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,"%u",&(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,"avg_window")==0){
    ret=sscanf(rhs,"%u",&(parameters->avg_window));
    if(ret!=1){
      fprintf(stderr, "error: parameter 'avg_window' should be an integer\n       got '%s'\n",rhs);
      return(-1);
    }
    *setavg_window=true;
  }
  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,"%u",&(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;
    }
    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=(char*)rhs+5;
    }
    else{
      fprintf(stderr, "error: unrecognized initial condition '%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),(2*parameters.K1+1)*(2*parameters.K2+1));

  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;

  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),(2*parameters.K1+1)*(2*parameters.K2+1));

  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(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;
}