#define VERSION "0.1"

#include <math.h>
#include <complex.h>
#include <fftw3.h>
#include <string.h>
#include <stdlib.h>
#include <stdbool.h>
#include "navier-stokes.h"
#include "driving.h"
#include "init.h"

// structure to store parameters, to make it easier to pass parameters to CLI functions
typedef struct nstrophy_parameters {
  int K1;
  int K2;
  int N1;
  int N2;
  unsigned int nsteps;
  double nu;
  double delta;
  double L;
  unsigned int print_freq;
  int seed;
} nstrophy_parameters;

// usage message
int print_usage();

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

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

#define COMMAND_UK 1
#define COMMAND_ALPHA 2
#define COMMAND_QUIET 3
#define COMMAND_ENERGY 4

#define DRIVING_ZERO 1
#define DRIVING_TEST 2

#define INIT_RANDOM 1
#define INIT_GAUSSIAN 2


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

  command=0;
  driving=0;
  init=0;

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

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

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

  // 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.print_freq, nthreads);
  }
  else if (command==COMMAND_ENERGY){
    energy(parameters.K1, parameters.K2, parameters.N1, parameters.N2, parameters.nsteps, parameters.nu, parameters.delta, parameters.L, u0, g, parameters.print_freq, nthreads);
  }
  else if(command==COMMAND_ALPHA){
    alpha(parameters.K1, parameters.K2, parameters.N1, parameters.N2, parameters.nsteps, parameters.nu, parameters.delta, parameters.L, u0, g, parameters.print_freq, nthreads);
  }
  else if(command==COMMAND_QUIET){
    quiet(parameters.K1, parameters.K2, parameters.N1, parameters.N2, parameters.nsteps, parameters.nu, parameters.delta, parameters.L, u0, g, nthreads);
  }
  else if(command==0){
    fprintf(stderr, "error: no command specified\n");
    print_usage();
  }

  free(g);
  free(u0);

  return(0);
}

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

// read command line arguments
#define CP_FLAG_PARAMS 1
#define CP_FLAG_DRIVING 2
#define CP_FLAG_NTHREADS 3
#define CP_FLAG_INIT 4
int read_args(
  int argc,
  const char* argv[],
  char** params,
  unsigned int* driving_force,
  unsigned int* init,
  unsigned int* command,
  unsigned int* nthreads
){
  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 'g':
	  flag=CP_FLAG_DRIVING;
	  break;
	case 't':
	  flag=CP_FLAG_NTHREADS;
	  break;
	case 'i':
	  flag=CP_FLAG_INIT;
	  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;
    }
    // driving force
    else if(flag==CP_FLAG_DRIVING){
      if (strcmp(argv[i],"zero")==0){
	*driving_force=DRIVING_ZERO;
      }
      else if (strcmp(argv[i],"test")==0){
	*driving_force=DRIVING_TEST;
      }
      else{
	fprintf(stderr, "error: unrecognized driving force '%s'\n",argv[i]);
	return(-1);
      }
      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;
    }
    // initial condition
    else if(flag==CP_FLAG_INIT){
      if (strcmp(argv[i],"random")==0){
	*init=INIT_RANDOM;
      }
      else if (strcmp(argv[i],"gaussian")==0){
	*init=INIT_GAUSSIAN;
      }
      else{
	fprintf(stderr, "error: unrecognized initial condition '%s'\n",argv[i]);
	return(-1);
      }
      flag=0;
    }
    // computation to run
    else{
      if(strcmp(argv[i], "uk")==0){
	*command=COMMAND_UK;
      }
      else if(strcmp(argv[i], "alpha")==0){
	*command=COMMAND_ALPHA;
      }
      else if(strcmp(argv[i], "quiet")==0){
	*command=COMMAND_QUIET;
      }
      else if(strcmp(argv[i], "energy")==0){
	*command=COMMAND_ENERGY;
      }
      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
){
  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->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->seed=17;

  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);
	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);
      if(ret<0){
	return ret;
      }
    }

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

  // if N not set explicitly, set it
  if (!setN1){
    parameters->N1=4*(parameters->K1)+1;
  }
  if (!setN2){
    parameters->N2=4*(parameters->K2)+1;
  }

  return(0);
}


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

  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,"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{
    fprintf(stderr, "error: unrecognized parameter '%s'\n",lhs);
    return(-1);
  }

  return(0);
}

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

  switch(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(
  unsigned int init,
  nstrophy_parameters parameters
){
  _Complex double* u0=calloc(sizeof(_Complex double),(2*parameters.K1+1)*(2*parameters.K2+1));

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

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

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