/*
Copyright 2017-2023 Ian Jauslin

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

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