#include <math.h>
#include <stdio.h>
#include <stddef.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_multimin.h>
#include <gsl/gsl_sf.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_cdf.h>
#include <gsl/gsl_linalg.h>
#include <getopt.h>
#include <string.h>

#include "libhdr"
#include "Library_DFE_v1.8.h"

/******************************************************************************/
int get_average_egf_vecs_1_and_2(double s,int n1, int n2, double t2_real,
				 int t2_lower, int t2_upper, double *egf_vec,
				 char *buffer_p1_n2_t2_lower,
				 char *buffer_p2_n2_t2_lower,
				 char *buffer_p1_n2_t2_upper,
				 char *buffer_p2_n2_t2_upper)
{
  static double egf_vec1_lower[maxnd+1], egf_vec2_lower[maxnd+1],
    egf_vec1_upper[maxnd+1], egf_vec2_upper[maxnd+1],
    egf_vec1[maxnd+1], egf_vec2[maxnd+1];
  double total_density;
  int i, res;
  int n1d=2*n1;
  int n2d=2*n2;

  if (s<-1)
    /*
       The contribution of strongly selected mutations
       is assumed to be inversely proportional to s, as
       expected from mutation selection balance, with a
       contribution from phase 2 mutations only.
    */
    {
      for (i=0; i<=2*n2; i++)
	{
	  egf_vec1[i] = 0;
	  egf_vec2[i] = 0;
	}
      egf_vec2[1] = -2/s;
    }
  else
    {
      res =  get_binary_egf_vec(buffer_p1_n2_t2_lower, n2, t2_lower,
				s, egf_vec1_lower);
      if (res==0) return 0;
      res =  get_binary_egf_vec(buffer_p2_n2_t2_lower, n2, t2_lower,
				s, egf_vec2_lower);
      if (res==0) return 0;
      res =  get_binary_egf_vec(buffer_p1_n2_t2_upper, n2, t2_upper,
				s, egf_vec1_upper);
      if (res==0) return 0;
      res =  get_binary_egf_vec(buffer_p2_n2_t2_upper, n2, t2_upper,
				s, egf_vec2_upper);
      if (res==0) return 0;

      compute_weighted_average_egf_vec(t2_real, t2_lower, t2_upper,
				       egf_vec1_lower, egf_vec1_upper, egf_vec1, 2*n2);
      //   monitorinput();
      compute_weighted_average_egf_vec(t2_real, t2_lower, t2_upper,
				       egf_vec2_lower, egf_vec2_upper, egf_vec2, 2*n2);
    }

  total_density = compute_weighted_average(egf_vec, egf_vec1, egf_vec2, n1d, 2*n2);
  if (s==0)
    {
      total_density_s0 = total_density;
      //      printf("s=0: total_density %lf\n", total_density);
    }

  return (1);
}

/******************************************************************************/
get_upper_lower_int(double parm, int *lower, int *upper, int n_int_evaluated, int *int_evaluated_vec)
{
  int i;
  *lower = undefined_int;
  *upper = undefined_int;
  //   printf("get_upper_lower_int parm %lf n_int_evaluated %d\n",
  //      parm, n_int_evaluated);
  for (i=0; i<n_int_evaluated; i++)
    {
      //      printf("int_evaluated_vec[%d] %d\n", i, int_evaluated_vec[i]);
      if ((double)int_evaluated_vec[i] >= parm)
	{
	  *upper = int_evaluated_vec[i];
	  if (i != 0) *lower = int_evaluated_vec[i-1];
	  return;
	}
    }
}
/******************************************************************************/
get_upper_lower_double(double parm, double*lower, double *upper,
		       int n_int_evaluated, double *double_evaluated_vec)
{
  int i;
  *lower = undefined;
  *upper = undefined;
  //   printf("get_upper_lower_int parm %lf n_int_evaluated %d\n",
  //      parm, n_int_evaluated);
  for (i=0; i<n_int_evaluated; i++)
    {
      //      printf("int_evaluated_vec[%d] %d\n", i, int_evaluated_vec[i]);
      if (double_evaluated_vec[i] >= parm)
	{
	  *upper =double_evaluated_vec[i];
	  if (i != 0) *lower = double_evaluated_vec[i-1];
	  return;
	}
    }
}
/******************************************************************************/
get_s_ranges()
{
  int i, stat, ind, i1;
  double d1, d2, d3;
  FILE *inf, *fopen();
  inf = openforreadsilent(s_range_file);
  stat = fscanf(inf, "%d", &nranges);
  if (stat!=1) gabort("get_s_ranges: input error 1", 0);
  if (nranges >max_s_ranges)
    gabort("get_s_ranges: nranges >max_s_ranges", nranges);
  for (i=1; i<=nranges; i++)
    {
      stat = fscanf(inf, "%d %lf %lf %lf %d", &ind, &d1, &d2, &d3, &i1);
      if (stat!=5) gabort("get_s_ranges: input error 2", 0);
      if (ind!=i) gabort("get_s_ranges: ind!=i", i);
      s_ranges[i].s_lower = d1;
      s_ranges[i].s_upper = d2;
      s_ranges[i].s_step = d3;
      s_ranges[i].s_lower_index = i1;
    }

}
/******************************************************************************/
get_int_evaluated_vec(int *int_evaluated_vec, int *n_int_evaluated,
		      int *int_lower, int *int_step,
		      char *int_evaluated_vec_file)
{
  int i, stat;
  FILE *inf, *fopen();
  int int_input, first_int, second_int, int_upper;
  inf = openforreadsilent(int_evaluated_vec_file);
  i = 0;
  for (;;)
    {
      stat = fscanf(inf, "%d", &int_input);
      if (stat==EOF) break;
      if (stat!=1) gabort("get_int_evaluated_vec: read error 1", stat);
      if (*n_int_evaluated >= max_n_int_evaluated)
	gabort("Value of n_int_evaluated too large: get_int_evaluated_vec",
	       get_int_evaluated_vec);
      int_evaluated_vec[*n_int_evaluated] = int_input;
      //      printf("int_evaluated_vec[%d] %d\n",
      //         *n_int_evaluated, int_evaluated_vec[*n_int_evaluated]);
      (*n_int_evaluated)++;
      if (i==0)
	{
	  first_int = int_input;
	  *int_lower = first_int;
	}
      if (i==1)
	{
	  second_int = int_input;
	}
      int_upper = int_input;
      i++;
    }
  fclose (inf);
  *int_step = second_int - first_int;
  //printf("int_lower %d int_upper %d int_step %d\n",
  //  *int_lower, int_upper, *int_step);

}
/******************************************************************************/
scale_vector(double *vec, int n2, double total_density_s0)
{
  int i;
  double tot = 0;
  //   printf("scale_vector: n2 %d total_density_s0 %lf\n", n2, total_density_s0);
  //   monitorinput();
  for (i=1; i<n2; i++)
    {
      vec[i] /= total_density_s0;
      tot += vec[i];
    }
  vec[0] = 1.0 - tot;        // Put all "missing" density in the lost mutant class
}
/******************************************************************************/
int get_binary_egf_vec(char *buffer, int n, int t2, double s, double *egf_vec)
{
  int i, j, file_size_bytes, n_read, t2_read, offset, tab_num, tab_size,
    size_of_float;
  float s_read, freq_read;
  char *ptr, ch;
  // Each expected gene freq table has two integers (n, t2), a float (s), then
  //   2*n -1 floats (frequencies)
  tab_size = (2*sizeof(int)+sizeof(float) + sizeof(float)*(2*n -1));
  file_size_bytes = n_s_evaluated_file*tab_size;
  //   printf("tab_size %d file_size_bytes %d\n", tab_size, file_size_bytes);
  tab_num =  compute_s_tab_num(s);
  //   printf("tab_num assuming ranges %d\n", tab_num);
  //   monitorinput();
  if (tab_num>=n_s_evaluated_file) return 0;
  //   printf("s %lf s_lower %lf s_step %lf (s - s_lower)/s_step %lf tab_num %d\n",
  //        s, s_lower, s_step, (s - s_lower)/s_step, tab_num);
  //   monitorinput();
  offset = tab_num*tab_size;
  assign_int_from_read_buff(&n_read, buffer, offset);
  //   printf("n_read %d\n", n_read);
  if (n_read!=n)
    {
      // printf("n %d n_read %d\n", n, n_read);
      gabort("get_binary_egf_vec: read error 1", n_read);
    }
  offset += sizeof(int);
  assign_int_from_read_buff(&t2_read, buffer, offset);
  //printf("t2 %d t2_read %d\n", t2, t2_read);
  if (t2_read!=t2) gabort("get_binary_egf_vec: read error 2", t2_read);
  offset += sizeof(int);
  s_read = assign_float_from_read_buff(buffer, offset);
  //   printf("s_read %f\n", s_read);
  size_of_float = sizeof(float);
  for (i=1; i<2*n; i++)
    {
      offset += size_of_float;
      //    assign_float_from_read_buff - in line code for speed
      ptr = &freq_read;
      for (j=0; j<size_of_float; j++)
	{
	  //      printf("assign_float_from_read_buff: ch %d\n", ch);
	  *(ptr + j) = buffer[offset + j];
	}
      //      printf("freq_read %f\n", freq_read);
      egf_vec[i] = (double)freq_read;
    }
  //   dumpvector(egf_vec, 1, 2*n -1, "egf_vec:");
  //   monitorinput();
  return 1;
}
/******************************************************************************/
fold_vector_double(double *vec, int counter)
{
  int i, j;
  j = 1;
  for (i=counter - 1; i > counter/2; i--)
    {
      vec[j] += vec[i];
      vec[i] = 0;
      j++;

    }
}
/******************************************************************************/
fold_vector_int(int *vec, int counter)
{
  int i, j;
  j = 1;
  for (i=counter - 1; i > counter/2; i--)
    {
      vec[j] += vec[i];
      vec[i] = 0;
      j++;

    }
}
/******************************************************************************/
int nearest_n2_ind(int n2_real)
{
  int i, ind = undefined_int;
  double min = undefined, diff, res;
  for (i=0; i<n_n2_evaluated; i++)
    {
      diff = n2_evaluated_vec[i] - n2_real;
      if (diff < 0) diff = -diff;

      if (min==undefined)
	{
	  min = diff;
	  ind = i;
	}
      else
	{
	  if (min > diff)
	    {
	      min = diff;
	      ind = i;
	    }
	}
    }
  if (ind == undefined_int) gabort("nearest_n2: ind == undefined_int", ind);
  res = n2_evaluated_vec[ind];;
  //   printf("nearest_n2: n2_real %lf result %lf\n", n2_real, res);
  //   monitorinput();
  return ind;
}
/******************************************************************************/
read_phase1_phase2_file_into_buffer(int n1, int phase, int n2, int t2,
				    char *buffer,
				    int file_size_bytes)
{
  static char n1_str[max_str_len], n2_str[max_str_len], t2_str[max_str_len];
  char file_str[max_str_len];
  int n2d;
  FILE *inf, *fopen();

  //   printf("read_phase1_phase2_file_into_buffer: phase %d n2 %d t2 %d\n",
  //      phase, n2, t2);
  n2d = 2*n2;
  int_to_string(n1, n1_str, max_str_len);

  int_to_string(n2, n2_str, max_str_len);

  int_to_string(t2, t2_str, max_str_len);

  if (strlen(n1_str) + strlen(n2_str) + strlen(t2_str) >= max_str_len)
    gabort("read_phase1_phase2_file_into_buffer: string too long", max_str_len);
  if (phase==1)
    {
      strcpy(file_str, phase_1_dir);
      strcat(file_str, "/n1:");
      strcat(file_str, n1_str);
      strcat(file_str, ":");
    }
  else
    {
      strcpy(file_str, phase_2_dir);
      strcat(file_str, "/");
    }
  strcat(file_str, "n2:");
  strcat(file_str, n2_str);
  strcat(file_str, ":t2:");
  strcat(file_str, t2_str);

  inf = openforreadsilent(file_str);
  fread(buffer, 1, file_size_bytes, inf);
  fclose(inf);
}
/******************************************************************************/
read_const_pop_file_into_buffer(int n1, char *buffer, int file_size_bytes)
{
  char file_str[max_str_len];
  FILE *inf, *fopen();

  //   printf("read_const_pop_file_into_buffer: n1 %d file_size_bytes %d\n",
  //      n1, file_size_bytes);
  strcpy(file_str, const_pop_dir);
  strcat(file_str, const_pop_file);
  //   printf("read_const_pop_file_into_buffer: file_str %s\n", file_str);
  //   monitorinput();
  inf = openforreadsilent(file_str);
  fread(buffer, 1, file_size_bytes, inf);
  fclose(inf);
}
/******************************************************************************/
compute_weighted_average_egf_vec(double t2_real, int t2_lower, int t2_upper,
				 double *egf_vec_lower, double *egf_vec_upper,
				 double *egf_vec, int n2d)
{
  int i;
  for (i=0; i<=n2d; i++)
    {
      egf_vec[i] = egf_vec_lower[i] + (egf_vec_upper[i] - egf_vec_lower[i])*(t2_real - (double)t2_lower)/(double)(t2_upper - t2_lower);
    }
}
/******************************************************************************/
double compute_weighted_average(double *weighted_vec, double *phase1_mut_vec,
				double *phase2_mut_vec, int n1d, int n2d)
{
  int i;
  double total = 0;
  for (i=1; i<n2d; i++)
    {
      weighted_vec[i]= (n1d*phase1_mut_vec[i] + n2d*phase2_mut_vec[i])/
	(n1d + n2d);
      total += weighted_vec[i];
    }
  //   dumpvector(weighted_vec, 0, n2d, "weighted average vector:"); monitorinput();
  return(total);
}
/******************************************************************************/
int get_const_pop_egf_vec(double s, int n1, double *egf_vec, char *buffer_const_pop,
			  int assign_tds0_mode)
{
  double total_density;
  int i, res;
  if (s<-1)          // The contribution of strongly selected mutations
    // is assumed to be inversely proportional to s, as
    // expected from mutation selection balance, with a
    // contribution from phase 2 mutations only.
    {
      for (i=0; i<=2*n1; i++)
	{
	  egf_vec[i] = 0;
	}
      egf_vec[1] = -2/s;
      //      printf("get_const_pop_egf_vec s %lf\n", s);
      //      dumpvector(egf_vec, 0, 2*n1, "get_const_pop_egf_vec:");
      //      monitorinput();
    }
  else
    {
      res =  get_binary_egf_vec(buffer_const_pop, n1, 0, s, egf_vec);
      if (res==0) return 0;
      //      dumpvector(egf_vec, 0, 2*n1, "get_const_pop_egf_vec:");
      //      monitorinput();
    }
  total_density = compute_weighted_average(egf_vec, egf_vec, egf_vec, 2*n1, 2*n1);
  if (assign_tds0_mode)
    {
      total_density_s0 = total_density;
    }
  //   printf("assign_tds0_mode %d: total_density_s0 %lf\n", assign_tds0_mode, total_density_s0);
  return (1);
}
/******************************************************************************/
assign_int_from_read_buff(int *res, char *buffer, int start)
{
  int temp, i;
  char *ptr, ch;
  ptr = &temp;
  for (i=0; i<sizeof(int); i++)
    {
      ch = buffer[start+i];

      *(ptr + i) = ch;
    }

  *res = temp;
}
/******************************************************************************/
float assign_float_from_read_buff(char *buffer, int start)
{
  float temp;
  int i;
  char *ptr, ch;
  ptr = &temp;
  for (i=0; i<sizeof(float); i++)
    {
      ch = buffer[start+i];
      //      printf("assign_float_from_read_buff: ch %d\n", ch);
      *(ptr + i) = ch;
    }
  //   printf("assign_float_from_read_buff: temp %f\n", temp);
  return temp;
}
/******************************************************************************/
int compute_s_tab_num(double s)
{
  double s_lower, s_upper, s_step;
  int i, tab_num;
  for (i=1; i<=nranges; i++)
    {
      s_lower = s_ranges[i].s_lower;
      s_upper = s_ranges[i].s_upper;
      if ((s >= s_lower)&&(s <= s_upper))
	{
	  s_step = s_ranges[i].s_step;
	  tab_num = (int)(floor((s - s_lower)/s_step + 0.0000000001));
	  return(tab_num + s_ranges[i].s_lower_index - 1);
	}
    }
  gabort("compute_s_tab_num: failure to find tab_num", 0);
}
/******************************************************************************/
int compute_file_size_bytes(int n)
{
  int tab_size, file_size_bytes;
  tab_size = (2*sizeof(int)+sizeof(float) + sizeof(float)*(2*n - 1));
  file_size_bytes = n_s_evaluated_file*tab_size;
  return (file_size_bytes);
}
/******************************************************************************/
set_up_file_name(int n1, char *froot, char *file_name)
{
  int len;
  static char n1_str[max_str_len];
  int_to_string(n1, n1_str, max_str_len);
  //   printf("set_up_file_name n1 %d froot %s n1_str %s\n", n1, froot, n1_str);
  //   monitorinput();
  len = strlen(data_path);
  len += strlen(n1_str) + 1;
  len += strlen(froot);
  if (len >= max_str_len) gabort("set_up_file_name: string too long (1)", len);
  strcpy(file_name, data_path);
  strcat(file_name, "n1_");
  strcat(file_name, n1_str);
  strcat(file_name, "/");
  strcat(file_name, froot);
  //   printf("set_up_file_name n1 file_name %s\n", file_name);
  //   monitorinput();
}
/******************************************************************************/
get_data_path(char *data_path)
{
  FILE *inf, *fopen();
  int stat;
  inf = openforread("directory_config.dat");
  stat = fscanf(inf, "%s", data_path);
  if (stat!=1) gabort("get_data_path: read error", 1);
  //   printf("get_data_path: data_path %s\n", data_path); monitorinput();
}
/******************************************************************************/
get_s_evaluated_vec(double *s_evaluated_vec, int *n_s_evaluated,
		    int *n_s_evaluated_file, char *s_evaluated_vec_file)
{
  int stat;
  FILE *inf, *fopen();
  double s;
  *n_s_evaluated = 0;
  inf = openforreadsilent(s_evaluated_vec_file);
  s = -100.0;                                 // Set up large |s| values specially
  while (s < -1)
    {
      s_evaluated_vec[*n_s_evaluated] = s;
      (*n_s_evaluated)++;
      if (s > -20) s = s + 1.0; else s = s + 5.0;
    }
  *n_s_evaluated_file = 0;
  for (;;)
    {
      stat = fscanf(inf, "%lf", &s);
      if (stat==EOF) break;
      if (stat!=1) gabort("get_s_limits: read error 1", stat);
      if (*n_s_evaluated >= max_n_s_evaluated)
	gabort("Value of n_s_evaluated too large: get_s_evaluated_vec",
	       get_s_evaluated_vec);
      s_evaluated_vec[*n_s_evaluated] = s;
      (*n_s_evaluated)++;
      (*n_s_evaluated_file)++;
    }
  fclose (inf);
  //   printf("n_s_evaluated %d n_s_evaluated_file %d\n", *n_s_evaluated,
  //      *n_s_evaluated_file);
  //   monitorinput();
}
/******************************************************************************/
get_lower_upper(int ind, double *s_evaluated_vec, int n_s_evaluated,
		double *lower, double *upper)
{
  double x, temp_r;
  if ((ind<0)||(ind>=n_s_evaluated))
    gabort("get_lower_upper: ind out of range", ind);
  x = s_evaluated_vec[ind];
  if (ind<n_s_evaluated-1)
    *upper = (x + s_evaluated_vec[ind+1])/2.0;
  else
    *upper = x;
  if (ind!=0)
    *lower = (x + s_evaluated_vec[ind-1])/2.0;
  else
    *lower = x;
  if (x <=0.0)
    {
      temp_r = *lower;
      *lower = *upper;
      *upper = temp_r;
      if (*lower!=0.0) *lower = -*lower;
      if (*upper!=0.0) *upper = -*upper;
    }
}
/******************************************************************************/
dumpvector(double *v, int min, int max, char *s)
{
  int i, control = 0;
  //   printf("\n");
  printf("%s", s); printf("\n");
  for (i = min; i<max; i++)
    {
      printf("%3d %10.8lf ", i, v[i]);
      control++;
      if (control==5)
	{
	  control = 0;
	  printf("\n");
	}
    }
  printf("\n");
}
/******************************************************************************/
dumpmatrix(double **m, int s1,
	   int rows, int s2, int cols, char *s)
{
  int i, j;
  printf("\n%s\n", s);
  for (i = s1; i<=rows; i++)
    {
      for (j = s2; j<=cols; j++)
	{
	  printf(" %2.3f", m[i][j]);
	}
      printf("\n");
    }
  printf("\n");
}
/******************************************************************************/
double selfn(double q, double s, double h)
{
  double num, denom;
  num = s*q*(1-q)*(q + h*(1-2*q));
  denom = 1 + 2.0*h*s*q*(1-q) + s*q*q;
  if (denom== 0) denom = .00000000000001;
  return( q + (num/denom));
}
/******************************************************************************/
setuprow(double **a, double p, int row, int n)
{

  double z;         /*probability of zero failures*/
  double x, y, temp1, temp2;
  int j;

  temp1 = 1.0 - p;
  if (temp1 == 0.0)  temp1 = .000000000001;            /*prevent zero divide*/
  temp1 = p/temp1;       /*First prevent negative log*/
  if (temp1 <=0)   temp1 = .000000000001;
  x = log(temp1);    /*used for later multiplication of series*/
  temp2 = 1.0-p;
  if (temp2<=0)   temp2 = .000000000001;
  z = (double)n*log(temp2);
  if (z > 150)   z = 150;
  a[row][0] = exp(z);
  y = a[row][0];        /*add up for later*/
  for (j = 1;j<=n; j++)
    {

      z = z + log(((double)n + 1.0 - (double)j)/(double)j) + x;
      /*             cumulative binomial coeff.*/
      a[row][j] = exp(z);
      y = y + a[row][j];          /*add up*/
    }
  if (y == 0.0)   y = .0000000000001;
  for (j=0; j<=n; j++) a[row][j] = a[row][j]/y;
}
/******************************************************************************/
setupmatrix(double **a, int n1, int n2, double s, double h)
{
  int i;
  double p;

  for (i = 1; i<=n1 - 1; i++)             /*do variable rows*/
    {
      p = (double)i/(double)n1;
      p = selfn(p, s,h);
      setuprow(a, p, i, n2);
    }
  for (i = 0; i<=n2; i++)
    {              /*do constant edges*/
      a[0][i] = 0;
      a[n1][i] = 0;
    }
  a[0][0] = 1.0;                          /*do absorbing corners*/
  a[n1][n2] = 1.0;

}
/******************************************************************************/
matrixinvert(int N, double **a,double *inverter)
{

  int i, j, lotkin_signum;
  gsl_matrix  *lotkin_a,*lotkin_inv;
  gsl_vector *x, *lotkin_b, *lotkin_x;

  gsl_permutation *lotkin_perm;

  /* allocate a, x, b */
  lotkin_a = gsl_matrix_alloc(N-1, N-1);
  lotkin_inv = gsl_matrix_alloc(N-1, N-1);
  x = gsl_vector_alloc(N-1);
  lotkin_b = gsl_vector_alloc(N-1);

  lotkin_x = gsl_vector_alloc(N-1);

  gsl_matrix *identity=gsl_matrix_alloc(N-1, N-1);

  for (i = 0; i<=N-2; i++)
    {
      for (j = 0; j<=N-2; j++)
	{
          gsl_matrix_set( identity, i, i,1);
	}
      ;
    }

  for (i = 0; i<=N-2; i++)
    {
      for (j = 0; j<=N-2; j++)
	{
	  gsl_matrix_set(lotkin_a,i,j,gsl_matrix_get(identity,i,j)-a[i+1][j+1]);
	}

    }

  /* LU decomposition and forward&backward substition */

  //gsl_blas_dgemv(CblasNoTrans, 1.0, lotkin_a, x, 0.0, lotkin_b);

  lotkin_perm = gsl_permutation_alloc(N-1);
  gsl_linalg_LU_decomp(lotkin_a, lotkin_perm, &lotkin_signum);
  gsl_linalg_LU_solve(lotkin_a, lotkin_perm, lotkin_b, lotkin_x);
  gsl_linalg_LU_invert (lotkin_a,lotkin_perm, lotkin_inv);

  /*apparently the inversion adds +1 to the first element of the vector
    so I substract it*/

  gsl_matrix_set(lotkin_inv,0,0,gsl_matrix_get(lotkin_inv,0,0)-1);

  //printf("\n%s\n", "Equilibrium");

  for (i = 0; i<1; i++)
    {
      for (j = 0; j<=N-2; j++)
	{

	  // printf("%d: %2.3f  ",j,gsl_matrix_get(lotkin_inv, i, j));
	}
      //printf("\n");
    }

  for (i=0;i<=N-2;i++){

    inverter[i+1]=gsl_matrix_get(lotkin_inv,0,i);

  }

  gsl_matrix_free(lotkin_a);
  gsl_matrix_free(lotkin_inv);
  gsl_vector_free(x);
  gsl_vector_free(lotkin_b);
  gsl_vector_free(lotkin_x);
  gsl_matrix_free(identity);
  gsl_permutation_free(lotkin_perm);
}
/******************************************************************************/
tmiterate(double **a,double *mut1,
	  int t, int n1, int n2,int decay,double *sumf)
{

  int k, i, j = 0;
  double z;

  double * steadystatefreq = (double*) calloc (maxnd+1, sizeof(double));
  double * mut2 = (double*) calloc (maxnd+1, sizeof(double));

  for (k =1 ;k<=t; k++)
    {                 /*t iterations*/

      /* Increment steadystatefreq. distribution*/
      for (i = 0; i<=n1; i++)
	{
	  steadystatefreq[i] = steadystatefreq[i] + mut1[i];
	}

      /* Perform matrix multiplication*/
      for (i = 0; i<=n2; i++)
	{
	  z = 0;
	  for (j = 0; j<=n1; j++)
	    {
	      z = z + a[j][i]*mut1[j];
	    }
	  mut2[i] = z;
	}

      /* Copy result of multiplication*/
      for (i=0; i<=n2; i++) {mut1[i] = mut2[i];

	if (decay==0){
	  sumf[i]+=mut1[i];}
	if(decay==1){sumf[i]=mut1[i];}
      }

    }

  free(steadystatefreq);
  free(mut2);
}
/******************************************************************************/
tmiterate2(double **a,double *mut1,
	   int t, int n1, int n2,int decay,double **a1)
{

  double *  steadystatefreq = (double*) calloc (maxnd+1, sizeof(double));
  double * mut2 = (double*) calloc (maxnd+1, sizeof(double));

  int k, i, j = 0;
  double z;

  for (k = 1; k<=t; k++){                 /*t iterations*/

    /* Increment steadystatefreq. distribution*/
    for (i = 0; i<=n1; i++)
      {
	steadystatefreq[i] = steadystatefreq[i] + mut1[i];
      }

    /* Perform matrix multiplication*/
    for (i = 0; i<=n2; i++)
      {
	z = 0;
	for (j = 0; j<=n1; j++)
	  {
            z = z + a[j][i]*mut1[j];
	  }
	mut2[i] = z;
      }

    /* Copy result of multiplication*/
    for (i=0; i<=n2; i++) {mut1[i] = mut2[i];

      if (decay==0){

	a1[k][i]+=mut1[i];
	if (k>1){a1[k][i]+=a1[k-1][i];}
      }

      if(decay==1){

	a1[k][i]=mut1[i];
	//if (k>0){ a1[k][i]+=a1[k-1][i];}
      }
    }

    //dumpvector(a1[k],0,n2,"matrix");
  }

  free(steadystatefreq);
  free(mut2);

}
/******************************************************************************/
eqf_using_matrix_inversion(int n1,double s,double **a,double *egf_out)
{
  setupmatrix(a, n1, n1, s, H);
  matrixinvert(n1,a,egf_out);
}
/******************************************************************************/
vector_average(int n1,int n2,double *fv1,double *fv2, double *fv_out)
{
  int i;
  for (i=0;i<n2;i++)
    {

      fv_out[i]=((n1*fv1[i])+(n2*fv2[i]))/(n1+n2);

    }
}
/******************************************************************************/
vector_s_average(int n,double P1,double *fv1,double *fv2, double *fv_out)
{
  int i;
  for (i=0;i<n;i++)
    {

      fv_out[i]=(P1*fv1[i])+((1-P1)*fv2[i]);

    }
}
/******************************************************************************/
output_sfs_to_file_thanasis_format(int n, int *sfs1,int *sfs2,char *sfs_filename)
{

  FILE *file= fopen(sfs_filename, "w" );

  int i;

  //neutral sites first

  for (i=0;i<n;i++){

    fprintf(file,"%d ",sfs1[i]);

  }

  fprintf(file,"\n");

  //selected sites second

  for (i=0;i<n;i++){

    fprintf(file,"%d ",sfs2[i]);

  }

  fclose(file);

}
/******************************************************************************/
output_sfs_to_file_peter_format(int n,int sample1, int sample2, int *sfs1,int *sfs2,char *sfs_filename)
{

  FILE *file= fopen(sfs_filename, "w" );
  int i;
  fprintf(file,"1\n");
  fprintf(file,"%d\n",n);

  //selected sites first
  fprintf(file,"%d\n",sample2);

  for (i=1;i<n;i++){

    fprintf(file,"%d ",sfs2[i]);

  }

  fprintf(file,"\n");
  //neutral sites second

  fprintf(file,"%d\n",sample1);
  for (i=1;i<n;i++){

    fprintf(file,"%d ",sfs1[i]);

  }

  fclose(file);

}
/******************************************************************************/
output_sfs_to_file_peter_format2(int n,int sample1, int sample2, int *sfs1,int *sfs2,char *sfs_filename)
{

  FILE *file= fopen(sfs_filename, "w" );
  int i;
  fprintf(file,"name\n");

  fprintf(file,"0 0\n0 0\n");
  fprintf(file,"%d\n",n);
  for (i=0;i<n;i++){

    fprintf(file,"%d ",sfs1[i]);

  }

  fprintf(file,"0\n");
  //fprintf(file,"%d\n",sample1);
  for (i=0;i<n;i++){

    fprintf(file,"%d ",sfs2[i]);

  }
fprintf(file,"0");
  fclose(file);

}
/******************************************************************************/

get_sfs(int alleles,int *par,char *sfs_filename)
{

  int i;
  int read;
  // char line[80];
  FILE *file= fopen(sfs_filename, "r" );

  for (i =0 ;i<alleles; i++)
    {
      fscanf (file, "%d", &read);
      par[i]=read;
    }

  fclose(file);

}
/******************************************************************************/
get_sfs_peter1(int *total_neu,int *total_sel,int *nalleles,int *sfs_sel,int *sfs_neu,char *sfs_filename)
{
  int i;
  int gene_number,sum;

  FILE *file= fopen(sfs_filename, "r" );
  fscanf (file, "%d", &gene_number);
  fscanf (file, "%d", nalleles);

  //printf("%d %d\n",read_sfs,*nalleles);

  fscanf (file, "%d", total_sel);
  sum=0;
  for (i =1 ;i<*nalleles; i++)
    {
      fscanf (file, "%d", &sfs_sel[i]);
      //printf("%d\n",sfs_sel[i]);
      sum+=sfs_sel[i];
    }
  sfs_sel[0]=*total_sel-sum;
  sfs_sel[*nalleles]=0;
  fscanf (file, "%d", total_neu);
  sum=0;
  for (i =1 ;i<*nalleles; i++)
    {
      fscanf (file, "%d", &sfs_neu[i]);
      //printf("%d\n",sfs_neu[i]);
      sum+=sfs_neu[i];
    }
  sfs_neu[0]=*total_neu-sum;
  sfs_neu[*nalleles]=0;
  //printf("%d\n",sfs_neu[*nalleles]);
  //monitorinput();
  fclose(file);
}
/******************************************************************************/
get_sfs_peter2(int *total_neu,int *total_sel,int *nalleles,int *sfs_sel,int *sfs_neu,char *sfs_filename,int *sel_sites,int *sel_diff,int *neu_sites,int *neu_diff)
{

  int i;
  int gene_number,sum;
  char name[100];
  FILE *file= fopen(sfs_filename, "r" );

  fscanf (file, "%s", name);
  fscanf (file, "%d", sel_sites);
  fscanf (file, "%d", sel_diff);
  fscanf (file, "%d", neu_sites);
  fscanf (file, "%d", neu_diff);
  fscanf (file, "%d", nalleles);

  //printf("%s %d %d %d %d %d\n",name,sel_sites,sel_diff,neu_sites,neu_diff,*nalleles);

  fscanf (file, "%d", &sfs_sel[0]);
  //printf("%d\n",sfs_sel[0]);
  sum=0;
  for (i =1 ;i<=*nalleles; i++)
    {
      fscanf (file, "%d", &sfs_sel[i]);
      //printf("%d\n",sfs_sel[i]);
      sum+=sfs_sel[i];
    }

//printf("next\n");
  fscanf (file, "%d", &sfs_neu[0]);
  //printf("%d\n",sfs_neu[0]);
  sum=0;
  for (i =1 ;i<=*nalleles; i++)
    {
      fscanf (file, "%d", &sfs_neu[i]);
      //printf("%d\n",sfs_neu[i]);
      sum+=sfs_neu[i];
    }

  //monitorinput();
  fclose(file);
}

/******************************************************************************/
double sfsfold_f(int nalleles,double *sfs,double *sfs_folded,int index)
{

int i;

      for (i=0;i<=nalleles/2;i++)
	{
	  sfs_folded[index+i]=sfs[i]+sfs[index-i];

	  if(i==nalleles/2&&(!odd(nalleles)))
	  {
	  sfs_folded[index+i]=sfs[i];
	  }

	}

}

/******************************************************************************/
int binarysearchcumulative_from_zero(double *array, int size)
{
  int cur, pre, post;
  double r;
  r = uniform();
  /*   printf("Uniform %f\n", r);*/
  pre = 0;
  post = size;
  cur = (size + 1)/2;
  do
    {
      if (array[cur] < r) pre = cur; else post = cur;
      cur = (pre + post)/2;
      if ((cur==post) || (cur==pre))
	{
	  if (r < array[pre])
	    {
	      /*            printf("Returning pre %d\n", pre);*/
	      return(pre);
	    }
	  else
	    {
	      /*            printf("Returning post %d\n", post);*/
	      return(post);
	    }
	}
    } while (size > 0);
}
/******************************************************************************/
double randomly_select_allele_frequency(double *egf_vec, double *temp1, int nd)
{
  double u, freq;
  int ind;
  ind = binarysearchcumulative_from_zero(egf_vec, nd-1);
  //   printf("ind %d\n", ind);
  (temp1[ind])++;
  freq = (double)ind/(double)nd;
  //   printf("ind %d freq %lf\n", ind, freq); monitorinput();
  return(freq);
}
/******************************************************************************/
set_up_cumulative_allele_freq_vec(int nd, double *egf_vec, double *cum_vec,
				  char *str)
{
  int i;
  for (i=0; i<nd; i++)
    {
      cum_vec[i] = egf_vec[i];
    }
  for (i=1; i<=nd; i++)
    {
      cum_vec[i] += cum_vec[i-1];
    }
}
/******************************************************************************/
generate_simulated_data(double *egf_vec, int nsites, int sample_size, int nd,
			int *sfs)
{
  int i, n_seg;
  double allele_freq;
  struct acc a;
  initacc(&a);
  for (i=1; i<=nsites; i++)
    {
      allele_freq =
	randomly_select_allele_frequency(egf_vec, allele_freq_sampled, nd);
      n_seg = bnldev(allele_freq, sample_size);
      accum(&a, (double)n_seg);
      (sfs[n_seg])++;
    }
}

/******************************************************************************/
egf_scaling_s(int n1,double f0,double *v_s_in, double *v_s_out)
{

  double sumx,sumy;
  int i;
  sumx=0;
  sumy=0;

  for (i=1;i<n1;i++){

    sumx+=v_s_in[i];

  }

  for (i=1;i<n1;i++){
    v_s_out[i]/=sumx;
    sumy+=v_s_out[i];
  }

  v_s_out[0]=(1-sumy);
}
/******************************************************************************/
egf_scaling_f0(int n1,double f0,double *fv)
{

  int i;
  for (i=0;i<n1;i++){
    fv[i]*=(1-f0);
  }
  fv[0]+=f0;
}
/******************************************************************************/
egf_scaling(int N,double f0,double *FV0,double *FVS)
{
  int n2d=N*2;
  egf_scaling_s(n2d,f0,FV0, FVS);
  egf_scaling_f0(n2d,f0,FVS);
}
/******************************************************************************/
binomial_sampling(int n1,int alleles,int sample, double *invy,int *discrete)
{

  int s1,success;
  int i;
  int n1d=2*n1;
  double prob;

  const gsl_rng_type * T;
  T = gsl_rng_taus;
  gsl_rng  *rgen =gsl_rng_alloc(T);

  /*
    based on equilibrium frequency vector (use it as a probability vector),
    I randomly generate numbers from 0 to n1d-1 for sample sites.
  */

  for (i=0;i<sample;i++)
    {

      gsl_ran_discrete_t *r= gsl_ran_discrete_preproc (n1d,invy);
      s1=gsl_ran_discrete (rgen, r);
      prob=(double)(s1)/n1d;

      success=gsl_ran_binomial(rgen,prob,alleles-1);
      discrete[success]++;

      gsl_ran_discrete_free(r);
    }
  gsl_rng_free (rgen);

}
/******************************************************************************/
int set_up_density_vec_equal_effects(double s, int n_s_evaluated, double *gamma_density_vec)
{
   double s_lower, s_upper, p_upper;
   int i, s_ind;
   s_lower = undefined;
   s_upper = undefined;
//   printf("set_up_gamma_density_vec_equal_effects s %lf\n", s);
//   monitorinput();
   for (i=0; i<n_s_evaluated; i++)
   {
      gamma_density_vec[i] = 0.0;
   }
   for (i=1; i<n_s_evaluated; i++)
   {
//      printf("%d %lf  ", i, s_evaluated_vec[i]);
//      printf("\n");
      if (s_evaluated_vec[i] >= s)
      {
         s_upper = s_evaluated_vec[i];
         s_ind = i;
         if (i != 0 ) s_lower = s_evaluated_vec[i-1];
         break;
      }
   }
   if ((s_lower==undefined)||(s_upper==undefined)) return 0;
   p_upper = (s - s_lower)/(s_upper - s_lower);
//   printf("set_up_gamma_density_vec_equal_effects: s_lower %lf s_upper %lf\n",
//      s_lower, s_upper);
//   printf("p_upper %lf\n", p_upper);
//   monitorinput();
   gamma_density_vec[s_ind] = p_upper;
   if (s_ind==1)
      gamma_density_vec[s_ind] += (1.0 - p_upper);
   else
      gamma_density_vec[s_ind-1] = (1.0 - p_upper);
//   for (i=1; i<n_s_evaluated; i++)
//   {
//      printf("s %lf gamma_density_vec[%d] %lf\n", s_evaluated_vec[i], i,
//         gamma_density_vec[i]);
//   }
//   monitorinput();
return (1);
}
/******************************************************************************/
int set_up_density_vec_step_effects(double alpha, double beta,
				    int n_s_evaluated, double *gamma_density_vec)
{
 if (alpha<0){alpha=-alpha;}
  if (beta<0){beta=-beta;}
  int i,s_ind;
  double lower, upper, area, total_area = 0.0, x, diff, step, inc_x, inc_y, gf;
  double mean;

  for (i=0; i<n_s_evaluated; i++)
    {

      x = s_evaluated_vec[i];

      get_lower_upper(i, s_evaluated_vec, n_s_evaluated, &lower, &upper);

      /*calculate cumulative uniform distribution for lower and upper*/
      inc_x=gsl_cdf_flat_P(lower,alpha,beta);
      inc_y=gsl_cdf_flat_P(upper,alpha,beta);

      area=inc_y-inc_x;
      gamma_density_vec[i] = area;
      total_area+=area;
      //printf("x %f inc_x %f inc_y %f area %f\n",x,inc_x,inc_y,area);
    }
//printf("alpha:%f beta:%f total_area:%f\n",alpha,beta, total_area);
//monitorinput();
// area =gsl_cdf_flat_P(-s_evaluated_vec[0], alpha, beta);
// gamma_density_vec[0]+=1.0-area;

 //if (gamma_density_vec[0]>0){printf("%f\n",gamma_density_vec[0]);}
  return (1);
}
/******************************************************************************/
double compute_beta_densities(double alpha, double *s_evaluated_vec,
			      int n_s_evaluated,
			      double *gamma_density_vec, double beta)
{
  int i;
  double lower, upper, area, total_area = 0.0, x, diff, step, inc_x, inc_y, gf;

  if (alpha <= 0) return undefined;
  if (beta <= 0) return undefined;

  for (i=0; i<n_s_evaluated; i++)
    {
      x = s_evaluated_vec[i];
      get_lower_upper(i, s_evaluated_vec, n_s_evaluated, &lower, &upper);

      inc_x = gsl_cdf_beta_P (lower, alpha, beta);
      inc_y = gsl_cdf_beta_P (upper, alpha, beta);

      area = (inc_y - inc_x);
      gamma_density_vec[i] = area;

      //printf("x %f inc_x %f inc_y %f area %f\n",x,inc_x,inc_y,area);

    }
   area = gsl_cdf_beta_P(-s_evaluated_vec[0], alpha, beta);
   gamma_density_vec[0]+=1.0-area;

  return (1);
}
/******************************************************************************/
double compute_gamma_densities(double alpha, double *s_evaluated_vec,
			       int n_s_evaluated, double *gamma_density_vec, double beta)
{
  int i;
  double lower, upper, area, total_area = 0.0, x, diff, step, inc_x, inc_y, gf;

  if (alpha <= 0) return undefined;

  for (i=0; i<n_s_evaluated; i++)
    {
      x = s_evaluated_vec[i];
      get_lower_upper(i, s_evaluated_vec, n_s_evaluated, &lower, &upper);

      inc_x = gsl_cdf_gamma_P(lower, beta,1/alpha);
      inc_y = gsl_cdf_gamma_P(upper, beta,1/alpha);

      area = (inc_y - inc_x);
      gamma_density_vec[i] = area;
      total_area+=area;
      //printf("x%f 1/alpha %f beta %f area %f\n",x,1/alpha,beta,area);

    }

  area = gsl_cdf_gamma_P(-s_evaluated_vec[0], beta,1/alpha);
  gamma_density_vec[0]+=1.0-area;
  //printf("x%f inc_x %f inc_y %f area %f\n",x,inc_x,inc_y,area);

  return (1);

}

/******************************************************************************/
double calculate_ne(int n1,int n2,int t)
{
double w1,w2,n_e;
double N1=(double)n1;
double N2=(double)n2;
double t2=(double)t;
/*compute weighted Ne*/
w1=N1*pow((1-1/(2*N2)),t2);
w2=N2*(1-exp(-t2/(2*N2)));

n_e=(N1*w1+N2*w2)/(w1+w2);

return n_e;
}


/*
 * Copyright 2020 Makani Technologies LLC
 *
 * 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.
 */

#include "common/c_math/linalg.h"

#include <assert.h>
#include <float.h>
#include <math.h>
#include <stddef.h>
#include <stdint.h>

#include <gsl/gsl_cblas.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_linalg.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_permute_vector.h>
#include <gsl/gsl_vector.h>

#include "common/c_math/gsl_linalg_extra.h"
#include "common/c_math/linalg_common.h"

#if defined(__arm__)
// Turn off any error handling by the GSL library.  The standard way
// of doing this is with gsl_set_error_handler_off; however this
// method does not require bringing in more components of GSL.
void gsl_error(const char *reason, const char *file, int line, int gsl_errno) {
  (void)reason;
  (void)file;
  (void)line;
  (void)gsl_errno;
}
#endif  // defined(__arm__)

// Operations on vectors with any number of elements.

bool VecIsSize(const Vec *v, int32_t l) { return v != NULL && l == v->length; }

double *VecPtr(const Vec *v, int32_t i) {
  assert(0 <= i && i < v->length);
  return &v->d[i];
}

double VecGet(const Vec *v, int32_t i) {
  assert(0 <= i && i < v->length);
  return v->d[i];
}

// v = [0, 0, ... ]
const Vec *VecZero(Vec *v) {
  for (int32_t i = 0; i < v->length; ++i) {
    v->d[i] = 0.0;
  }
  return v;
}

// v_out = a*v + v_out
const Vec *VecAxpy(double a, const Vec *v, Vec *v_out) {
  assert(v->length == v_out->length);
  for (int32_t i = 0; i < v->length; ++i) {
    v_out->d[i] += a * v->d[i];
  }
  return v_out;
}

// Changes the length of a vector.  Returns true on success, false
// otherwise.
bool VecResize(int32_t length, Vec *v) {
  if (v->length == length) {
    return true;
  } else if (v->variable_len && 0 <= length && length <= v->max_length) {
    v->length = length;
    return true;
  }
  assert(false);
  return false;
}

// v_out = v_in
const Vec *VecCopy(const Vec *v_in, Vec *v_out) {
  VecResize(v_in->length, v_out);
  for (int32_t i = 0; i < v_out->length; ++i) {
    v_out->d[i] = v_in->d[i];
  }
  return v_out;
}

// Copy the contents of an array in a Vec.
const Vec *VecInit(const double a_in[], int32_t length, Vec *v_out) {
  VecResize(length, v_out);
  for (int32_t i = 0; i < v_out->length; ++i) {
    v_out->d[i] = a_in[i];
  }
  return v_out;
}

// v_out = scale * v_in
const Vec *VecScale(const Vec *v_in, double scale, Vec *v_out) {
  VecResize(v_in->length, v_out);
  for (int32_t i = 0; i < v_in->length; ++i) {
    v_out->d[i] = scale * v_in->d[i];
  }
  return v_out;
}

// v_out = v0 + v1
const Vec *VecAdd(const Vec *v0, const Vec *v1, Vec *v_out) {
  assert(v0->length == v1->length);
  VecResize(v0->length, v_out);
  for (int32_t i = 0; i < v_out->length; ++i) {
    v_out->d[i] = v0->d[i] + v1->d[i];
  }
  return v_out;
}

// v_out = v0 + v1 + v2
const Vec *VecAdd3(const Vec *v0, const Vec *v1, const Vec *v2, Vec *v_out) {
  assert(v0->length == v1->length && v0->length == v2->length);
  VecResize(v0->length, v_out);
  for (int32_t i = 0; i < v_out->length; ++i) {
    v_out->d[i] = v0->d[i] + v1->d[i] + v2->d[i];
  }
  return v_out;
}

// v_out = v0 - v1
const Vec *VecSub(const Vec *v0, const Vec *v1, Vec *v_out) {
  assert(v0->length == v1->length);
  VecResize(v0->length, v_out);
  for (int32_t i = 0; i < v_out->length; ++i) {
    v_out->d[i] = v0->d[i] - v1->d[i];
  }
  return v_out;
}

// v_out = c0 * v0 + c1 * v1
const Vec *VecLinComb(double c0, const Vec *v0, double c1, const Vec *v1,
                      Vec *v_out) {
  assert(v0->length == v1->length);
  VecResize(v0->length, v_out);
  for (int32_t i = 0; i < v_out->length; ++i) {
    v_out->d[i] = c0 * v0->d[i] + c1 * v1->d[i];
  }
  return v_out;
}

// v_out = c0 * v0 + c1 * v1 + c2 * v2
const Vec *VecLinComb3(double c0, const Vec *v0, double c1, const Vec *v1,
                       double c2, const Vec *v2, Vec *v_out) {
  assert(v0->length == v1->length && v0->length == v2->length);
  VecResize(v0->length, v_out);
  for (int32_t i = 0; i < v_out->length; ++i) {
    v_out->d[i] = c0 * v0->d[i] + c1 * v1->d[i] + c2 * v2->d[i];
  }
  return v_out;
}

// v_out = v0 .* v1
const Vec *VecMult(const Vec *v0, const Vec *v1, Vec *v_out) {
  assert(v0->length == v1->length);
  VecResize(v0->length, v_out);
  for (int32_t i = 0; i < v_out->length; ++i) {
    v_out->d[i] = v0->d[i] * v1->d[i];
  }
  return v_out;
}

// v0^T * v1
double VecDot(const Vec *v0, const Vec *v1) {
  assert(v0->length == v1->length);
  double tmp = 0.0;
  for (int32_t i = 0; i < v0->length; ++i) {
    tmp += v0->d[i] * v1->d[i];
  }
  return tmp;
}

// Computes the 2-norm of a vector.
//
// TODO: This does not yet handle overflow and underflow
// gracefully.
double VecNorm(const Vec *v) { return sqrt(VecNormSquared(v)); }

double VecNormBound(const Vec *v, double low) { return fmax(low, VecNorm(v)); }

double VecNormSquared(const Vec *v) { return VecDot(v, v); }

const Vec *VecNormalize(const Vec *v_in, Vec *v_out) {
  return VecScale(v_in, 1.0 / VecNormBound(v_in, 1e-9), v_out);
}

// Return the elements of v_in where inds != 0. The MATLAB logical indexing
// equivalent is: v_out = v_in(inds).
const Vec *VecSlice(const Vec *v_in, const int32_t inds[], Vec *v_out) {
  int32_t i_out = 0;
  for (int32_t i = 0; i < v_in->length; ++i) {
    if (inds[i]) {
      assert(i_out < v_out->max_length);
      // Don't use VecPtr here because we haven't set v_out->length yet.
      v_out->d[i_out] = VecGet(v_in, i);
      i_out++;
    }
  }

  // Do not resize v_out until the end of the function if you want to
  // be able to reuse and input as an output.
  VecResize(i_out, v_out);
  return v_out;
}

// Assigns elements of a larger vector with those from a smaller one. The
// MATLAB logical indexing equivalent is: v_out(inds) = v_in. The sum(inds != 0)
// must be equal to length(v_in). Warning: there is no check that
// length(inds) == length(v_out).
const Vec *VecSliceSet(const Vec *v_in, const int32_t inds[], Vec *v_out) {
  assert(v_in != v_out);

  int32_t i_in = 0;
  for (int32_t i = 0; i < v_out->length; ++i) {
    if (inds[i]) {
      *VecPtr(v_out, i) = VecGet(v_in, i_in);
      i_in++;
    }
  }

  // Assert that the sum(inds != 0) is equal to length(v_in).
  assert(v_in->length == i_in);

  return v_out;
}

// Operations on arbitrary matrices.

bool MatIsSize(const Mat *m, int32_t nr, int32_t nc) {
  return m != NULL && m->nr == nr && m->nc == nc;
}

double *MatPtr(const Mat *m, int32_t i, int32_t j) {
  assert(i < m->nr && j < m->nc);
  return &m->d[i * m->nc + j];
}

double MatGet(const Mat *m, int32_t i, int32_t j) {
  assert(i < m->nr && j < m->nc);
  return m->d[i * m->nc + j];
}

bool MatResize(int32_t nr, int32_t nc, Mat *m) {
  if (m->nr == nr && m->nc == nc) {
    return true;
  } else if (m->variable_dim && nr * nc <= m->max_size) {
    m->nr = nr;
    m->nc = nc;
    return true;
  }
  assert(false);
  return false;
}

// Copy the contents of an array into a Mat.
const Mat *MatInit(const double *d, int32_t nr, int32_t nc, Mat *m) {
  MatResize(nr, nc, m);
  for (int32_t i = 0; i < m->nr; ++i) {
    for (int32_t j = 0; j < m->nc; ++j) {
      *MatPtr(m, i, j) = d[i * m->nc + j];
    }
  }
  return m;
}

// m_out = scale * m_in
const Mat *MatScale(const Mat *m_in, double scale, Mat *m_out) {
  MatResize(m_in->nr, m_in->nc, m_out);
  for (int32_t i = 0; i < m_out->nr; ++i) {
    for (int32_t j = 0; j < m_out->nc; ++j) {
      *MatPtr(m_out, i, j) = scale * MatGet(m_in, i, j);
    }
  }
  return m_out;
}

const Mat *MatZero(Mat *m) {
  MatArrZero(m->nr, m->nc, m->d);
  return m;
}

const Mat *MatCopy(const Mat *a, Mat *b) {
  if (a != b) {  // No work to do if the Mat pointers are equal.
    MatResize(a->nr, a->nc, b);
    // This is still safe if the data pointers are equal.
    MatArrCopy(a->d, a->nr, a->nc, b->d);
  }
  return b;
}

// Copies the sub-matrix src(start_row + (1:num_rows), start_col + (1:num_cols))
// into dest(dest_row + (1:num_rows), dest_col + (1:num_cols)).
//
// It is not safe for src and dest to share any memory.
//
// Returns:
//   A pointer to dest.
const Mat *MatSubmatSet(const Mat *src, int32_t start_row, int32_t start_col,
                        int32_t num_rows, int32_t num_cols, int32_t dest_row,
                        int32_t dest_col, Mat *dest) {
  assert(src != dest && src->d != dest->d);
  assert(start_row >= 0 && start_col >= 0);
  assert(num_rows >= 0 && num_cols >= 0);
  assert(dest_row >= 0 && dest_col >= 0);
  // Make sure there is sufficient data and sufficient room.
  assert(src->nr >= start_row + num_rows && src->nc >= start_col + num_cols);
  assert(dest->nr >= dest_row + num_rows && dest->nc >= dest_col + num_cols);

  for (int32_t i = 0; i < num_rows; ++i) {
    for (int32_t j = 0; j < num_cols; ++j) {
      *MatPtr(dest, dest_row + i, dest_col + j) =
          MatGet(src, i + start_row, j + start_col);
    }
  }
  return dest;
}

const Mat *MatI(int32_t n, Mat *m) {
  MatResize(n, n, m);
  MatArrI(n, m->d);
  return m;
}

const Mat *MatMult(const Mat *m1, const Mat *m2, Mat *m_out) {
  assert(m1->nc == m2->nr);
  MatResize(m1->nr, m2->nc, m_out);
  MatArrMult(m1->d, m1->nr, m1->nc, m2->d, m2->nc, m_out->d);
  return m_out;
}

// Generalized matrix-vector multiply.
//
//   z = alpha*op(a)*x + beta*y.
//
// Args:
//   ta: Indicates if a should be transposed.
//   alpha, a, x, beta: See formula above.
//   y: Can be NULL if beta = 0.0.
//   z: Resulting vector.  It is safe to reuse y as z.
const Vec *MatVecGenMult(TransposeType ta, double alpha, const Mat *a,
                         const Vec *x, double beta, const Vec *y, Vec *z) {
  assert(a != NULL && x != NULL && z != NULL);
  assert((ta == kNoTrans ? a->nc : a->nr) == x->length);
  const int32_t zl = (ta == kNoTrans) ? a->nr : a->nc;
  assert(y == NULL || y->length == zl);
  assert(y != NULL || beta == 0.0);
  if (y == NULL) beta = 0.0;

  if (y == NULL || beta == 0.0) {
    VecResize(zl, z);
  } else {
    VecCopy(y, z);
  }

  MatArrGemv(ta, alpha, a->d, a->nr, a->nc, x->d, beta, z->d);
  return z;
}

// Generalized matrix-matrix multiply.
//
//   d = alpha*op(a)*op(b) + beta*c.
//
// Args:
//   ta: Indicates if a should be transposed.
//   tb: Indicates if b should be transposed.
//   alpha, a, b, beta: See above formula.
//   c: Can be NULL if beta = 0.0.
//   d: Resulting matrix.  It is safe to reuse c as d.
const Mat *MatGenMult(TransposeType ta, TransposeType tb, double alpha,
                      const Mat *a, const Mat *b, double beta, const Mat *c,
                      Mat *d) {
  assert(a != NULL && b != NULL && d != NULL);
  assert((ta == kNoTrans ? a->nc : a->nr) == (tb == kNoTrans ? b->nr : b->nc));
  const int32_t nr_d = (ta == kNoTrans ? a->nr : a->nc);
  const int32_t nc_d = (tb == kNoTrans ? b->nc : b->nr);
  assert(c == NULL || (c->nr == nr_d && c->nc == nc_d));
  assert(c != NULL || beta == 0.0);
  if (c == NULL) beta = 0.0;

  if (c == NULL || beta == 0.0)
    MatResize(nr_d, nc_d, d);
  else
    MatCopy(c, d);

  MatArrGemm(ta, tb, alpha, a->d, a->nr, a->nc, b->d, nc_d, beta, d->d);
  return d;
}

const Mat *MatMult3(const Mat *m0, const Mat *m1, const Mat *m2, Mat *m_out) {
  assert(m0->nc == m1->nr && m1->nc == m2->nr);
  MatResize(m0->nr, m2->nc, m_out);
  MAT(m0->nr, m1->nc, m0m1);
  return MatMult(MatMult(m0, m1, &m0m1), m2, m_out);
}

const Vec *MatVecMult(const Mat *m, const Vec *v_in, Vec *v_out) {
  assert(m->nc == v_in->length);
  VecResize(m->nr, v_out);
  MatArrMult(m->d, m->nr, m->nc, v_in->d, 1, v_out->d);
  return v_out;
}

const Vec *MatTransVecMult(const Mat *m, const Vec *v_in, Vec *v_out) {
  assert(m->nr == v_in->length);
  VecResize(m->nc, v_out);
  MatArrGemm(kTrans, kNoTrans, 1.0, m->d, m->nr, m->nc, v_in->d, 1, 0.0,
             v_out->d);
  return v_out;
}

const Mat *MatTrans(const Mat *m, Mat *m_out) {
  MatResize(m->nc, m->nr, m_out);
  MatArrTrans(m->d, m->nr, m->nc, m_out->d);
  return m_out;
}

const Mat *MatAdd(const Mat *m0, const Mat *m1, Mat *m_out) {
  assert(m0->nr == m1->nr && m0->nc == m1->nc);
  MatResize(m0->nr, m0->nc, m_out);
  for (int32_t i = 0; i < m0->nr; ++i) {
    for (int32_t j = 0; j < m0->nc; ++j) {
      *MatPtr(m_out, i, j) = MatGet(m0, i, j) + MatGet(m1, i, j);
    }
  }
  return m_out;
}

const Mat *MatSub(const Mat *m0, const Mat *m1, Mat *m_out) {
  assert(m0->nr == m1->nr && m0->nc == m1->nc);
  MatResize(m0->nr, m0->nc, m_out);
  for (int32_t i = 0; i < m0->nr; ++i) {
    for (int32_t j = 0; j < m0->nc; ++j) {
      *MatPtr(m_out, i, j) = MatGet(m0, i, j) - MatGet(m1, i, j);
    }
  }
  return m_out;
}

void MatQrDecomp(const Mat *A, Mat *Q, Mat *R) {
  MatResize(A->nr, A->nc, R);
  double *Qd = NULL;
  if (Q != NULL) {
    MatResize(A->nr, A->nr, Q);
    Qd = Q->d;
  }
  MatArrQrDecomp(A->d, A->nr, A->nc, Qd, R->d);
}

// Calculates the "thin" singular value decomposition of an m x n
// matrix A, defined by:
//
//   A = U * S * V'
//
// where U is an m x n column-orthogonal matrix (i.e. U'*U = I), S is
// an n x n diagonal matrix with non-decreasing elements, and V is an
// n x n orthogonal matrix (i.e. V'*V = V*V' = I).  The "thin" aspect
// is equivalent to MATLAB's "economy" SVD.
//
// Args:
//   A: Input m x n matrix with m >= n.
//   U: Output m x n column-orthogonal matrix.
//   s: Output n length vector of the diagonal elements of S.
//   V: Output n x n orthogonal matrix.
//
// Returns:
//   kLinalgErrorNone on success and kLinalgErrorMaxIter if GSL's
//   solver reached its maximum number of iterations.
int32_t MatThinSvDecomp(const Mat *A, Mat *U, Vec *s, Mat *V) {
  assert(A != NULL && U != NULL && s != NULL && V != NULL);
  assert(A->nc <= VEC_MAX_ELEMENTS);
  assert(A->nr >= A->nc);

  MatCopy(A, U);
  VecResize(A->nc, s);
  MatResize(A->nc, A->nc, V);

  double work_data[VEC_MAX_ELEMENTS];
  gsl_matrix_view U_gsl =
      gsl_matrix_view_array(U->d, (size_t)U->nr, (size_t)U->nc);
  gsl_matrix_view V_gsl =
      gsl_matrix_view_array(V->d, (size_t)V->nr, (size_t)V->nc);
  gsl_vector_view s_gsl = gsl_vector_view_array(s->d, (size_t)s->length);
  gsl_vector_view work = gsl_vector_view_array(work_data, (size_t)s->length);
  int32_t ret = gsl_linalg_SV_decomp(&U_gsl.matrix, &V_gsl.matrix,
                                     &s_gsl.vector, &work.vector);
  // These are the only errors that gsl_linalg_SV_decomp should ever
  // return, assuming we're passing in matrices and vectors of the
  // appropriate sizes.
  assert(ret == GSL_SUCCESS || ret == GSL_EMAXITER);

  if (ret == GSL_EMAXITER) {
    assert(false);
    return kLinalgErrorMaxIter;
  }
  return kLinalgErrorNone;
}

// Calculates the rank of a matrix.  Performs a QR decomposition (with
// pivoting) of the input matrix and finds the number of diagonal
// elements of R with magnitude greater than a specified tolerance.
//
// Args:
//   A: Input matrix.
//   tol: Tolerance for checking linear independence of vectors.
//
// Returns:
//   Rank of the matrix.
int32_t MatRank(const Mat *A, double tol) {
  assert(A != NULL);
  assert(A->nr * A->nc <= MAT_MAX_ELEMENTS);
  assert(A->nr <= VEC_MAX_ELEMENTS && A->nc <= VEC_MAX_ELEMENTS);
  assert(tol >= DBL_EPSILON);

  int32_t min_dim = (A->nr < A->nc) ? A->nr : A->nc;
  if (min_dim == 0) return 0;

  double qr_data[MAT_MAX_ELEMENTS];
  double tau_data[VEC_MAX_ELEMENTS];
  double norm_data[VEC_MAX_ELEMENTS];
  size_t perm_data[VEC_MAX_ELEMENTS];
  int signum;
  gsl_matrix_const_view A_gsl =
      gsl_matrix_const_view_array(A->d, (size_t)A->nr, (size_t)A->nc);
  gsl_matrix_view QR =
      gsl_matrix_view_array(qr_data, (size_t)A->nr, (size_t)A->nc);
  gsl_vector_view tau = gsl_vector_view_array(tau_data, (size_t)min_dim);
  gsl_permutation p = {(size_t)A->nc, perm_data};
  gsl_vector_view norm = gsl_vector_view_array(norm_data, (size_t)A->nc);

  gsl_matrix_memcpy(&QR.matrix, &A_gsl.matrix);
  int32_t ret = gsl_linalg_QRPT_decomp(&QR.matrix, &tau.vector, &p, &signum,
                                       &norm.vector);
  assert(ret == GSL_SUCCESS);
  (void)ret;

  int32_t rank = 0;
  while (rank < min_dim) {
    if (fabs(gsl_matrix_get(&QR.matrix, (size_t)rank, (size_t)rank)) <= tol) {
      break;
    }
    rank++;
  }

  return rank;
}

// Solves A*x = b for x similar to x = A \ b.
int32_t MatVecLeftDivide(const Mat *A, const Vec *b, Vec *x) {
  assert(A != NULL && b != NULL && x != NULL);
  assert(b != x);
  assert(A->nr == b->length && A->nr > 0);
  VecResize(A->nc, x);
  if (A->nc == 0) return kLinalgErrorNone;

  size_t m = (size_t)A->nr, n = (size_t)A->nc;
  gsl_matrix_const_view A_gsl = gsl_matrix_const_view_array(A->d, m, n);
  gsl_matrix_const_view b_gsl = gsl_matrix_const_view_array(b->d, m, 1U);
  gsl_matrix_view x_gsl = gsl_matrix_view_array(x->d, n, 1U);
  int32_t ret =
      GslMatrixDivide(CblasLeft, &A_gsl.matrix, &b_gsl.matrix, &x_gsl.matrix);
  assert(ret == GSL_SUCCESS || ret == GSL_ESING);
  return (ret == GSL_ESING) ? kLinalgErrorSingularMat : kLinalgErrorNone;
}

// Solves x*A = b for x similar to x = b / A.
int32_t MatVecRightDivide(const Mat *A, const Vec *b, Vec *x) {
  assert(A != NULL && b != NULL && x != NULL);
  assert(b != x);
  assert(A->nc == b->length && A->nc > 0);
  VecResize(A->nr, x);
  if (A->nr == 0) return kLinalgErrorNone;

  size_t m = (size_t)A->nr, n = (size_t)A->nc;
  gsl_matrix_const_view A_gsl = gsl_matrix_const_view_array(A->d, m, n);
  gsl_matrix_const_view b_gsl = gsl_matrix_const_view_array(b->d, 1U, n);
  gsl_matrix_view x_gsl = gsl_matrix_view_array(x->d, 1U, m);
  int32_t ret =
      GslMatrixDivide(CblasRight, &A_gsl.matrix, &b_gsl.matrix, &x_gsl.matrix);
  assert(ret == GSL_SUCCESS || ret == GSL_ESING);
  return (ret == GSL_ESING) ? kLinalgErrorSingularMat : kLinalgErrorNone;
}

// Solves A*X = b for X similar to X = A \ B.
int32_t MatMatLeftDivide(const Mat *A, const Mat *B, Mat *X) {
  assert(A != NULL && B != NULL && X != NULL);
  assert(A != X && B != X);
  assert(A->nr == B->nr && A->nr > 0);
  MatResize(A->nc, B->nc, X);
  if (A->nc == 0) return kLinalgErrorNone;

  size_t m = (size_t)A->nr, n = (size_t)A->nc, k = (size_t)B->nc;
  gsl_matrix_const_view A_gsl = gsl_matrix_const_view_array(A->d, m, n);
  gsl_matrix_const_view B_gsl = gsl_matrix_const_view_array(B->d, m, k);
  gsl_matrix_view X_gsl = gsl_matrix_view_array(X->d, n, k);
  int32_t ret =
      GslMatrixDivide(CblasLeft, &A_gsl.matrix, &B_gsl.matrix, &X_gsl.matrix);
  assert(ret == GSL_SUCCESS || ret == GSL_ESING);
  return (ret == GSL_ESING) ? kLinalgErrorSingularMat : kLinalgErrorNone;
}

// Solves X*A = B for X similar to X = B / A.
int32_t MatMatRightDivide(const Mat *A, const Mat *B, Mat *X) {
  assert(A != NULL && B != NULL && X != NULL);
  assert(A != X && B != X);
  assert(A->nc == B->nc && A->nc > 0);
  MatResize(B->nr, A->nr, X);
  if (A->nr == 0) return kLinalgErrorNone;

  size_t m = (size_t)A->nr, n = (size_t)A->nc, k = (size_t)B->nr;
  gsl_matrix_const_view A_gsl = gsl_matrix_const_view_array(A->d, m, n);
  gsl_matrix_const_view B_gsl = gsl_matrix_const_view_array(B->d, k, n);
  gsl_matrix_view X_gsl = gsl_matrix_view_array(X->d, k, m);
  int32_t ret =
      GslMatrixDivide(CblasRight, &A_gsl.matrix, &B_gsl.matrix, &X_gsl.matrix);
  assert(ret == GSL_SUCCESS || ret == GSL_ESING);
  return (ret == GSL_ESING) ? kLinalgErrorSingularMat : kLinalgErrorNone;
}

bool MatIsUpperTriangular(const Mat *r) {
  return MatArrIsUpperTriangular(r->d, r->nr, r->nc);
}

bool MatIsLowerTriangular(const Mat *l) {
  return MatArrIsLowerTriangular(l->d, l->nr, l->nc);
}

bool MatHasNonnegDiag(const Mat *m) {
  assert(m != NULL);
  for (int32_t i = 0; i < m->nc && i < m->nr; ++i) {
    if (MatGet(m, i, i) < 0.0) return false;
  }
  return true;
}

// Calculates the matrix square root of P = A'*A + B'*B by calculating a QR
// decomposition, Q*R = [A ; B].  It is assumed that the matrix [A ; B] is full
// row rank (i.e. has linearly independent columns).
//
// Q is an orthogonal matrix and R "trapezoidal matrix", i.e. R = [sqrt_P ; 0]
// where sqrt_P is upper triangular.  Thus:
//   P = A'*A + B'*B = [A ; B]'*[A ; B] = [sqrT_P ; 0]'*Q'*Q*[sqrt_P ; 0].
//
// Args:
//   A: na-by-nc matrix.
//   B: nb-by-nc matrix.
//   sqrt_P: Upper-triangular matrix with positive diagonal elements such that
//     sqrt_P'*sqrt_P = P.
//
// Note: it is safe to reuse A or B as sqrt_P (assuming shapes match).
void MatSqrtSum(const Mat *A, const Mat *B, Mat *sqrt_P) {
  assert(A != NULL && B != NULL && sqrt_P != NULL);
  assert(A->nc == B->nc);
  assert(A->nr + B->nr >= A->nc);  // This is required for full row rank.
  MatResize(A->nc, A->nc, sqrt_P);

  // Define a temporary matrix, stack = [A ; B].
  MAT(A->nr + B->nr, A->nc, stack);
  MatSubmatSet(A, 0, 0, A->nr, A->nc, 0, 0, &stack);
  MatSubmatSet(B, 0, 0, B->nr, B->nc, A->nr, 0, &stack);
  MatQrDecomp(&stack, NULL, &stack);
  MatSubmatSet(&stack, 0, 0, A->nc, A->nc, 0, 0, sqrt_P);
  // Next, enforce a non-negative diagonal.
  for (int32_t i = 0; i < sqrt_P->nr; ++i) {
    if (MatGet(sqrt_P, i, i) < 0.0) {
      for (int32_t j = i; j < sqrt_P->nc; ++j) {
        *MatPtr(sqrt_P, i, j) *= -1.0;
      }
    }
  }
}

int32_t MatMatBackSub(const Mat *R, const Mat *b, Mat *x) {
  assert(b->nr == R->nr && b->nc == x->nc);
  assert(x->nr == R->nc);
  return MatArrBackSub(R->d, R->nr, R->nc, b->d, b->nc, x->d);
}

int32_t MatMatForwardSub(const Mat *L, const Mat *b, Mat *x) {
  assert(b->nr == L->nr && b->nc == x->nc);
  assert(x->nr == L->nc);
  return MatArrForwardSub(L->d, L->nr, L->nc, b->d, b->nc, x->d);
}

int32_t MatVecBackSub(const Mat *R, const Vec *b, Vec *x) {
  assert(b->length == R->nr);
  assert(x->length == R->nc);
  return MatArrBackSub(R->d, R->nr, R->nc, b->d, 1, x->d);
}

int32_t MatVecForwardSub(const Mat *L, const Vec *b, Vec *x) {
  assert(b->length == L->nr);
  assert(x->length == L->nc);
  return MatArrForwardSub(L->d, L->nr, L->nc, b->d, 1, x->d);
}

// Returns the rows and columns of m where rows != 0 and cols != 0
const Mat *MatSlice(const Mat *m, const int32_t rows[], const int32_t cols[],
                    Mat *m_out) {
  int32_t nr_out = 0, nc_out = 0;
  if (rows == NULL) {
    nr_out = m->nr;
  } else {
    for (int32_t i = 0; i < m->nr; ++i)
      if (rows[i]) nr_out++;
  }
  if (cols == NULL) {
    nc_out = m->nc;
  } else {
    for (int32_t i = 0; i < m->nc; ++i)
      if (cols[i]) nc_out++;
  }

  if (m_out->variable_dim) {
    assert(nr_out * nc_out <= m_out->max_size);
  } else {
    assert(m_out->nr == nr_out && m_out->nc == nc_out);
  }

  int32_t i_out = 0;
  for (int32_t i = 0; i < m->nr; ++i) {
    if (rows == NULL || rows[i]) {
      int32_t j_out = 0;
      for (int32_t j = 0; j < m->nc; ++j) {
        if (cols == NULL || cols[j]) {
          // Don't use MatPtr here because we haven't set m_out->nc yet.
          m_out->d[i_out * nc_out + j_out] = MatGet(m, i, j);
          j_out++;
        }
      }
      i_out++;
    }
  }

  // Do not resize m until the end of the function if you want to be
  // able to reuse the input as the output.  Otherwise, the matrix
  // indexing in MatGet will be off.
  MatResize(nr_out, nc_out, m_out);
  return m_out;
}

const double *MatArrMult(const double *m1, int32_t nr1, int32_t nc1,
                         const double *m2, int32_t nc2, double *m_out) {
  assert(m_out != m1 && m_out != m2);
  return MatArrGemm(kNoTrans, kNoTrans, 1.0, m1, nr1, nc1, m2, nc2, 0.0, m_out);
}

// This is modeled after the BLAS function GEMV, which calculates:
//     y <-- alpha*op(A)*x + beta*y,
// where op either returns the matrix or its transpose.
//
// Args:
//   ta: Whether to transpose A.
//   alpha: See above formula.
//   a: Data for A.
//   nra: Number of rows of A without transpose applied.
//   nca: Number of columns of A without transpose applied.
//   x: Data for x.
//   beta: See above formula.
//   c: Data for y.
const double *MatArrGemv(TransposeType ta, double alpha, const double *a,
                         int32_t nra, int32_t nca, const double *x, double beta,
                         double *y) {
  assert(x != y);
  int32_t ly = ((ta == kNoTrans) ? nra : nca);
  int32_t lx = ((ta == kNoTrans) ? nca : nra);
  double aij, axj;
  for (int32_t i = 0; i < ly; ++i) {  // For each entry of y.
    axj = 0.0;
    for (int32_t j = 0; j < lx; ++j) {  // Walk down the row of op(A).
      aij = (ta == kNoTrans) ? a[i * nca + j] : a[j * nca + i];
      axj += aij * x[j];
    }
    if (beta == 0.0) {
      y[i] = alpha * axj;
    } else {
      y[i] = alpha * axj + beta * y[i];
    }
  }
  return y;
}

// This is modeled after the BLAS function GEMM, which calculates:
//     C <-- alpha*op(A)*op(B) + beta*C,
// where op either returns the matrix or its transpose.
//
// Args:
//   ta: Whether to transpose A.
//   tb: Whether to transpose B.
//   alpha: See above formula.
//   a: Data for A.
//   nra: Number of rows of A without transpose applied.
//   nca: Number of columns of A without transpose applied.
//   b: Data for B.
//   ncc: Number of columns of C.
//   beta: See above formula.
//   c: Data for C.
const double *MatArrGemm(TransposeType ta, TransposeType tb, double alpha,
                         const double *a, int32_t nra, int32_t nca,
                         const double *b, int32_t ncc, double beta, double *c) {
  assert(a != c && b != c);
  int32_t nrc = (ta == kNoTrans) ? nra : nca;
  int32_t nk = (ta == kNoTrans) ? nca : nra;
  double ABij, aik, bkj;
  for (int32_t i = 0; i < nrc; ++i) {
    for (int32_t j = 0; j < ncc; ++j) {
      ABij = 0.0;
      for (int32_t k = 0; k < nk; ++k) {
        aik = (ta == kNoTrans) ? a[i * nca + k] : a[k * nca + i];
        bkj = (tb == kNoTrans) ? b[k * ncc + j] : b[j * nk + k];
        ABij += aik * bkj;
      }
      if (beta == 0.0) {
        c[i * ncc + j] = alpha * ABij;
      } else {
        c[i * ncc + j] = alpha * ABij + beta * c[i * ncc + j];
      }
    }
  }
  return c;
}

const double *MatArrZero(int32_t nr, int32_t nc, double *m_out) {
  for (int32_t i = 0; i < nr; ++i) {
    for (int32_t j = 0; j < nc; ++j) {
      m_out[i * nc + j] = 0.0;
    }
  }
  return m_out;
}

void MatArrCopy(const double *d_in, int32_t nr, int32_t nc, double *d_out) {
  for (int32_t i = 0; i < nr; ++i) {
    for (int32_t j = 0; j < nc; ++j) {
      d_out[i * nc + j] = d_in[i * nc + j];
    }
  }
}

void MatArrTrans(const double *m_in, int32_t nr, int32_t nc, double *m_out) {
  assert(m_out != m_in && m_out != NULL);
  for (int32_t i = 0; i < nr; ++i) {
    for (int32_t j = 0; j < nc; ++j) {
      m_out[j * nr + i] = m_in[i * nc + j];
    }
  }
}

const double *MatArrI(int32_t nr, double *m_out) {
  MatArrZero(nr, nr, m_out);
  for (int32_t i = 0; i < nr; ++i) {
    m_out[i * nr + i] = 1.0;
  }
  return m_out;
}

// QR decomposition.  Decompose a matrix into orthogonal and
// upper-right-triangular components.
//
// You may reuse a as r, and you may skip calculating q by passing a
// NULL pointer.
void MatArrQrDecomp(const double *a, int32_t nr, int32_t nc, double *q,
                    double *r) {
  assert(q != a && r != q && r != NULL);
  assert(nr <= VEC_MAX_ELEMENTS && nc <= VEC_MAX_ELEMENTS);

  double tau_data[VEC_MAX_ELEMENTS];
  gsl_matrix_const_view A =
      gsl_matrix_const_view_array(a, (size_t)nr, (size_t)nc);
  gsl_matrix_view QR = gsl_matrix_view_array(r, (size_t)nr, (size_t)nc);
  gsl_vector_view tau =
      gsl_vector_view_array(tau_data, (nr < nc) ? (size_t)nr : (size_t)nc);
  gsl_matrix_memcpy(&QR.matrix, &A.matrix);
  gsl_linalg_QR_decomp(&QR.matrix, &tau.vector);

  if (q != NULL) {
    gsl_matrix_view Q = gsl_matrix_view_array(q, (size_t)nr, (size_t)nr);
    gsl_matrix_view R = gsl_matrix_view_array(r, (size_t)nr, (size_t)nc);
    gsl_linalg_QR_unpack(&QR.matrix, &tau.vector, &Q.matrix, &R.matrix);
  } else {
    for (int32_t i = 1; i < nr; ++i) {
      for (int32_t j = 0; j < i && j < nc; ++j) {
        r[i * nc + j] = 0.0;
      }
    }
  }
}

bool MatArrIsUpperTriangular(const double *r, int32_t nr, int32_t nc) {
  for (int32_t i = 1; i < nr; ++i) {
    for (int32_t j = 0; j < i && j < nc; ++j) {
      if (fabs(r[i * nc + j]) >= DBL_EPSILON) {
        return false;
      }
    }
  }
  return true;
}

bool MatArrIsLowerTriangular(const double *l, int32_t nr, int32_t nc) {
  for (int32_t i = 0; i < nr; ++i) {
    for (int32_t j = i + 1; j < nc; ++j) {
      if (fabs(l[i * nc + j]) >= DBL_EPSILON) {
        return false;
      }
    }
  }
  return true;
}

// Solves R*X = B using back substitution where R is an upper
// triangular matrix.  If the triangular matrix contains a zero
// diagonal element (i.e. the matrix is singular), we set the
// corresponding x element to zero.
//
// Returns kLinalgErrorNone on success and kLinalgErrorSingularMat
// if the matrix was singular.
//
// Note: Confirmed safe to reuse b as x.
int32_t MatArrBackSub(const double *r, int32_t nr, int32_t nc, const double *b,
                      int32_t nc_b, double *x) {
  assert(r != x);
  assert(nr >= nc);
  assert(nr >= 0 && nc >= 0 && nc_b >= 0);
  assert(MatArrIsUpperTriangular(r, nr, nc));
  (void)nr;

  int32_t err = kLinalgErrorNone;
  // Solve for the k-th column of x.
  for (int32_t k = 0; k < nc_b; ++k) {
    // Iterate over the rows of r.
    for (int32_t i = nc - 1; i >= 0; --i) {
      if (fabs(r[i * nc + i]) <= DBL_EPSILON) {
        err = kLinalgErrorSingularMat;
        x[i * nc_b + k] = 0.0;
      } else {
        double sum = 0.0;
        for (int32_t j = nc - 1; j > i; --j) {
          sum += x[j * nc_b + k] * r[i * nc + j];
        }
        x[i * nc_b + k] = (b[i * nc_b + k] - sum) / r[i * nc + i];
      }
    }
    // TODO: If nr > nc, then the bottom of r is all zeros.
    // Should we test if b has non-zero elements for these rows?
  }
  return err;
}

// Solves L*X = B using forward substitution where L is a lower
// triangular matrix.  If the triangular matrix contains a zero
// diagonal element (i.e. the matrix is singular), we set the
// corresponding x element to zero.
//
// Returns kLinalgErrorNone on success and kLinalgErrorSingularMat
// if the matrix is singular.
//
// Note: Confirmed safe to reuse b as x.
int32_t MatArrForwardSub(const double *l, int32_t nr, int32_t nc,
                         const double *b, int32_t nc_b, double *x) {
  assert(nr <= nc);
  assert(MatArrIsLowerTriangular(l, nr, nc));
  int32_t err = kLinalgErrorNone;
  for (int32_t k = 0; k < nc_b; ++k) {  // Solve for the k-th column of x.
    for (int32_t i = 0; i < nr; ++i) {  // Iterate over the rows of l.
      if (fabs(l[i * nc + i]) <= DBL_EPSILON) {
        err = kLinalgErrorSingularMat;
        x[i * nc_b + k] = 0.0;
      } else {
        double sum = 0.0;
        for (int32_t j = 0; j < i; ++j) {
          sum += x[j * nc_b + k] * l[i * nc + j];
        }
        x[i * nc_b + k] = (b[i * nc_b + k] - sum) / l[i * nc + i];
      }
    }
    for (int32_t i = nr; i < nc; ++i) {
      x[i * nc_b + k] = 0.0;
    }
  }
  return err;
}


#include "globals.h"
#include "peano.h"
#include <gsl/gsl_heapsort.h>

peanoKey Peano_Key ( const double x, const double y, const double z );
static void reorder_particles();

void Print_Int_Bits128 ( const peanoKey val )
{
    for ( int i = 127; i >= 0; i-- ) {

        printf ( "%llu", ( long long ) ( ( val & ( ( peanoKey ) 1 << i ) ) >> i ) );

        if ( i % 3 == 0 && i != 0 ) {
            printf ( "." );
        }
    }
    printf ( "\n" );
    fflush ( stdout );

    return ;
}

void Print_Int_Bits128r ( const peanoKey val )
{
    for ( int i = 127; i >= 0; i-- ) {

        printf ( "%llu", ( long long ) ( ( val & ( ( peanoKey ) 1 << i ) ) >> i ) );

        if ( i % 3 - 2 == 0 && i != 0 ) {
            printf ( "." );
        }
    }
    printf ( "\n" );
    fflush ( stdout );

    return ;
}

int compare_peanoKeys ( const void *a, const void *b )
{
    const peanoKey *x = ( const peanoKey * ) a;
    const peanoKey *y = ( const peanoKey * ) b;

    return ( int ) ( *x > *y ) - ( *x < *y );
}

static peanoKey *Keys = NULL;
static size_t *Idx = NULL;

void Sort_Particles_By_Peano_Key()
{
    if ( Keys == NULL ) {
        Keys = malloc ( Param.Npart * sizeof ( *Keys ) );
    } else {
        memset ( Keys, 0, Param.Npart * sizeof ( *Keys ) );
    }

    if ( Idx == NULL ) {
        Idx = malloc ( Param.Npart * sizeof ( *Idx ) );
    } else {
        memset ( Idx, 0, Param.Npart * sizeof ( *Idx ) );
    }

    #pragma omp parallel for
    for ( int ipart = 0; ipart < Param.Npart; ipart++ ) {

        double px = P[ipart].Pos[0] / Problem.Boxsize[0];
        double py = P[ipart].Pos[1] / Problem.Boxsize[1];
        double pz = P[ipart].Pos[2] / Problem.Boxsize[2];

        P[ipart].Key = Keys[ipart] = Peano_Key ( px, py, pz );
    }

    gsl_heapsort_index ( Idx, Keys, Param.Npart, sizeof ( *Keys ), &compare_peanoKeys );

    reorder_particles();

    return ;
}

static void reorder_particles()
{
    for ( int i = 0; i < Param.Npart; i++ ) {

        if ( Idx[i] == i ) {
            continue;
        }

        int dest = i;

        struct ParticleData Ptmp = P[i];
        struct GasParticleData Sphtmp = SphP[i];

        int src = Idx[i];

        for ( ;; ) {

            P[dest] = P[src];
            SphP[dest] = SphP[src];

            Idx[dest] = dest;

            dest = src;

            src = Idx[dest];

            if ( src == i ) {
                break;
            }
        }

        P[dest] = Ptmp;
        SphP[dest] = Sphtmp;

        Idx[dest] = dest;

    } // for i

    return ;
}

peanoKey Peano_Key ( const double x, const double y, const double z )
{
    Assert ( x >= 0 && x <= 1, "X coordinate of out range [0,1] have %g", x );
    Assert ( y >= 0 && y <= 1, "Y coordinate of out range [0,1] have %g", y );
    Assert ( z >= 0 && z <= 1, "Z coordinate of out range [0,1] have %g", z );

    const uint64_t m = 1UL << 63; // = 2^63;

    uint64_t X[3] = { y * m, z * m, x * m };

    /* Inverse undo */

    for ( uint64_t q = m; q > 1; q >>= 1 ) {

        uint64_t P = q - 1;

        if ( X[0] & q ) {
            X[0] ^= P;    // invert
        }

        for ( int i = 1; i < 3; i++ ) {

            if ( X[i] & q ) {

                X[0] ^= P; // invert

            } else {

                uint64_t t = ( X[0] ^ X[i] ) & P;

                X[0] ^= t;
                X[i] ^= t;

            } // exchange
        }
    }

    /* Gray encode (inverse of decode) */

    for ( int i = 1; i < 3; i++ ) {
        X[i] ^= X[i - 1];
    }

    uint64_t t = X[2];

    for ( int i = 1; i < 64; i <<= 1 ) {
        X[2] ^= X[2] >> i;
    }

    t ^= X[2];

    for ( int i = 1; i >= 0; i-- ) {
        X[i] ^= t;
    }

    /* branch free bit interleave of transpose array X into key */

    peanoKey key = 0;

    X[1] >>= 1;
    X[2] >>= 2;    // lowest bits not important

    for ( int i = 0; i < N_PEANO_TRIPLETS + 1; i++ ) {

        uint64_t col = ( ( X[0] & 0x8000000000000000 )
                         | ( X[1] & 0x4000000000000000 )
                         | ( X[2] & 0x2000000000000000 ) ) >> 61;

        key <<= 3;

        X[0] <<= 1;
        X[1] <<= 1;
        X[2] <<= 1;

        key |= col;
    }

    key <<= 2;

    return key;
}

/* This constructs the peano key with reversed triplet order. The order in the
 * triplets however is the same ! Also level zero is carried explicitely
 * to ease tree construction. */

peanoKey Reversed_Peano_Key ( const double x, const double y, const double z )
{
    Assert ( x >= 0 && x <= 1, "X coordinate of out range [0,1] have %g", x );
    Assert ( y >= 0 && y <= 1, "Y coordinate of out range [0,1] have %g", y );
    Assert ( z >= 0 && z <= 1, "Z coordinate of out range [0,1] have %g", z );

    const uint64_t m = 1UL << 63; // = 2^63;

    uint64_t X[3] = { y * m, z * m, x * m };

    /* Inverse undo */

    for ( uint64_t q = m; q > 1; q >>= 1 ) {

        uint64_t P = q - 1;

        if ( X[0] & q ) {
            X[0] ^= P;    // invert
        }

        for ( int i = 1; i < 3; i++ ) {

            if ( X[i] & q ) {

                X[0] ^= P; // invert

            } else {

                uint64_t t = ( X[0] ^ X[i] ) & P;

                X[0] ^= t;
                X[i] ^= t;

            } // exchange
        }
    }

    /* Gray encode (inverse of decode) */

    for ( int i = 1; i < 3; i++ ) {
        X[i] ^= X[i - 1];
    }

    uint64_t t = X[2];

    for ( int i = 1; i < 64; i <<= 1 ) {
        X[2] ^= X[2] >> i;
    }

    t ^= X[2];

    for ( int i = 1; i >= 0; i-- ) {
        X[i] ^= t;
    }

    /* branch free reversed (!) bit interleave of transpose array X into key */

    peanoKey key = 0;

    X[0] >>= 18;
    X[1] >>= 19;
    X[2] >>= 20;    // lowest bits not important

    for ( int i = 0; i < N_PEANO_TRIPLETS + 1; i++ ) {

        uint64_t col = ( ( X[0] & 0x4 ) | ( X[1] & 0x2 ) | ( X[2] & 0x1 ) );

        key <<= 3;

        key |= col;

        X[0] >>= 1;
        X[1] >>= 1;
        X[2] >>= 1;
    }

    key <<= 3; // include level 0

    return key;
}

void test_peanokey()
{
    const double box[3]  = { 1.0, 1, 1};
    double a[3] = { 0 };
    int order = 1;
    float delta = 1 / pow ( 2.0, order );
    int n = roundf ( 1 / delta );

    for ( int i = 0; i < n; i++ ) {
        for ( int j = 0; j < n; j++ ) {
            for ( int k = 0; k < n; k++ ) {

                a[0] = ( i + 0.5 ) * delta / box[0];
                a[1] = ( j + 0.5 ) * delta / box[1];
                a[2] = ( k + 0.5 ) * delta / box[2];

                peanoKey stdkey =  Peano_Key ( a[0], a[1], a[2] );

                printf ( "%g %g %g %llu  \n", a[0], a[1], a[2],
                         ( long long ) stdkey );

                Print_Int_Bits128 ( stdkey );

                printf ( "\n" );
            }
        }
    }

    return ;
}


 // MIT License

 // Copyright (c) [2017] [Vinay Yuvashankar]

 // Permission is hereby granted, free of charge, to any person obtaining a copy
 // of this software and associated documentation files (the "Software"), to deal
 // in the Software without restriction, including without limitation the rights
 // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 // copies of the Software, and to permit persons to whom the Software is
 // furnished to do so, subject to the following conditions:

 // The above copyright notice and this permission notice shall be included in all
 // copies or substantial portions of the Software.

 // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 // SOFTWARE.

/** \file ERSP.cc

	\brief This file contains all of the function required to generate the Event Related Spectral Pertubation of EEG signals.
*/
#include "wavelet.h"
#include <gsl/gsl_statistics.h>

#define NORMALIZATION_FACTOR 0.375402913609157562

int ERSP (double * raw_data, double* scales, const int sampling_frequency, const int n,
	const int J, int const trials, const int padding_type,
	double * output)
{
	int i, j, x;

	//Calculate the necessary constants for the Continuous Wavelet Transform.
    const int    PADDED_SIZE = CalculatePaddingSize(n, padding_type);
    const int    m           = PRE_EVENT_TIME * sampling_frequency;
    const double dw          = (2 * M_PI * sampling_frequency)/(PADDED_SIZE); //NOT IN RAD/SEC in Hz

    fftw_init_threads();
    #pragma omp parallel private(i, j, x) shared(raw_data, output, scales) default(none)
    {
    	//Array Inits
    	double * pre_stimulus, *wavelet_out, *baseline_out;
    	fftw_complex *data_in, *fft_data, *filter_convolution, *fftw_result;
    	fftw_plan plan_forward, plan_backward;

	    //Memory Allocations
	    wavelet_out  = (double*) malloc( n * J * sizeof(double) );
	    baseline_out = (double*) malloc( n * J * sizeof(double) );
	    pre_stimulus = (double*) malloc( m     * sizeof(double) );

	    //FFTW Memory Allocations
	    data_in            = (fftw_complex *) fftw_malloc( sizeof( fftw_complex ) * PADDED_SIZE );
		fft_data           = (fftw_complex *) fftw_malloc( sizeof( fftw_complex ) * PADDED_SIZE );
		filter_convolution = (fftw_complex *) fftw_malloc( sizeof( fftw_complex ) * PADDED_SIZE );
		fftw_result        = (fftw_complex *) fftw_malloc( sizeof( fftw_complex ) * PADDED_SIZE );

		#pragma omp critical (make_plan)
		{
			fftw_plan_with_nthreads(1);
			if (fftw_import_wisdom_from_filename("FFTW_Plan.wise") == 0)
			{
				printf("No FFTW Plan, using an unoptimized method\n");
				plan_forward  = fftw_plan_dft_1d(PADDED_SIZE, data_in,            fft_data,    FFTW_FORWARD,  FFTW_ESTIMATE);
				plan_backward = fftw_plan_dft_1d(PADDED_SIZE, filter_convolution, fftw_result, FFTW_BACKWARD, FFTW_ESTIMATE);
			}
			else
			{
				plan_forward  = fftw_plan_dft_1d(PADDED_SIZE, data_in,            fft_data,    FFTW_FORWARD,  FFTW_EXHAUSTIVE);
				plan_backward = fftw_plan_dft_1d(PADDED_SIZE, filter_convolution, fftw_result, FFTW_BACKWARD, FFTW_EXHAUSTIVE);
			}
		}

		/*Begin ERSP*/
		#pragma omp for
		for ( x = 0; x < trials; ++x)
		{
			memset(data_in,            0.0, sizeof( fftw_complex ) * PADDED_SIZE);
			memset(fft_data,           0.0, sizeof( fftw_complex ) * PADDED_SIZE);
			memset(filter_convolution, 0.0, sizeof( fftw_complex ) * PADDED_SIZE);
			memset(fftw_result,        0.0, sizeof( fftw_complex ) * PADDED_SIZE);

			/*Begin Wavelet Analysis*/
			PopulateDataArray(raw_data, n, x,
							  PADDED_SIZE, padding_type, data_in);
			fftw_execute(plan_forward);

			for (i = 0; i < J; ++i)
			{
				FrequencyMultiply(fft_data, PADDED_SIZE, scales[i], dw,
								  filter_convolution);

				//Take the inverse FFT and store it in fftw_result
				fftw_execute(plan_backward);

				//Calculate the power and store it in result this may need to be changed to accomodate for phase
				for (j = 0; j < n; ++j)
				{
					wavelet_out[i * n + j] = MAGNITUDE(fftw_result[j][0], fftw_result[j][1]) / (NORMALIZATION_FACTOR * sqrt(scales[i]) );
				}
			}
			/*End Wavelet Analysis*/

			//Remove the baseline
			RemoveBaseline(pre_stimulus, wavelet_out,
						   n, J, m,
						   baseline_out);

			for ( i = 0; i < n * J; ++i)
			{
				output[i] += fabs(baseline_out[i]);
			}
		}

		#pragma omp for simd
		for ( i = 0; i < n * J; ++i)
		{
			output[i] /= trials;
		}
		/*End ERSP*/

		//Sanitation Engineering
		fftw_destroy_plan(plan_forward); fftw_destroy_plan(plan_backward);
		fftw_free(data_in); fftw_free(fft_data); fftw_free(filter_convolution); fftw_free(fftw_result);
		free(pre_stimulus); free(baseline_out); free(wavelet_out);

    }/*End of OpenMP*/
	return(0);
}

int RemoveBaseline(double* pre_stimulus, double* pre_baseline_array,
	const int n, const int J, const int m,
	double* output)
{
	double value;
	double mean, sDeviation;
	const int stride = 1;
	int i, j;

	for ( i = 0; i < J; ++i)
	{
		//Copy the pre trial results from each frequency block into pre_stimulus.
		for ( j = 0; j < m; ++j)
		{
			pre_stimulus[j] = pre_baseline_array[i * n + j];
		}

		//Calculate mean and standard deviation
		mean       = gsl_stats_mean(pre_stimulus, stride, m);
    	sDeviation = gsl_stats_sd_m(pre_stimulus, stride, m, mean);

    	//Remove the Baseline
	    for ( j = 0; j < n; ++j)
	    {
	    	value = pre_baseline_array[i * n + j] * pre_baseline_array[i * n + j];
	        output[i * n + j] = (fabs(value) - mean) / sDeviation;
	        // output[i * n + j] = ( pre_baseline_array[i * n + j] - mean ) / sDeviation;
	    }
	}

	return(0);
}

int Generate_FFTW_Wisdom(int padded_size)
{
	int success_flag = 1;

	//Array Inits
	fftw_complex *data_in, *fft_data, *filter_convolution, *fftw_result;
	fftw_plan plan_forward, plan_backward;

	fftw_init_threads();

	//FFTW Memory Allocations
    data_in            = (fftw_complex *) fftw_malloc( sizeof( fftw_complex ) * padded_size );
	fft_data           = (fftw_complex *) fftw_malloc( sizeof( fftw_complex ) * padded_size );
	filter_convolution = (fftw_complex *) fftw_malloc( sizeof( fftw_complex ) * padded_size );
	fftw_result        = (fftw_complex *) fftw_malloc( sizeof( fftw_complex ) * padded_size );

	//FFTW Planning
	fftw_plan_with_nthreads(1);
	printf("Generating an Exhaustive FFTW Plan\n");
	plan_forward  = fftw_plan_dft_1d(padded_size, data_in,            fft_data,    FFTW_FORWARD,  FFTW_EXHAUSTIVE);
	plan_backward = fftw_plan_dft_1d(padded_size, filter_convolution, fftw_result, FFTW_BACKWARD, FFTW_EXHAUSTIVE);

	printf("Writing FFTW plan to FFTW_Plan.wise\n");
	if (fftw_export_wisdom_to_filename("FFTW_Plan.wise") != 0)
	{
		success_flag = 0;
	}

	fftw_destroy_plan(plan_forward); fftw_destroy_plan(plan_backward);
	fftw_free(data_in); fftw_free(fft_data); fftw_free(filter_convolution); fftw_free(fftw_result);

	return(success_flag);
}

int FrequencyMultiply(fftw_complex* fft_data,
	const int data_size, const double scale, const double dw,
	fftw_complex* filter_convolution)
{
	int j;
	double value;
	//Compute the Fourier Morlet at 0 and N/2
	double norm = sqrt(scale);
	value = CompleteFourierMorlet(0.0, scale, norm);

	filter_convolution[0][0] = (fft_data[0][0]/data_size) * value;
	filter_convolution[0][1] = (fft_data[0][1]/data_size) * value;

	filter_convolution[data_size/2][0] = 0.0;
	filter_convolution[data_size/2][1] = 0.0;

	//Compute the Fourier Morlet Convolution in between
	for (j = 1; j < data_size/2 - 1; ++j)
	{
		value = CompleteFourierMorlet( j * dw , scale, norm);
		filter_convolution[j][0] = (fft_data[j][0]/data_size) * value;
		filter_convolution[j][1] = (fft_data[j][1]/data_size) * value;

		filter_convolution[data_size- j][0] = 0.0;
		filter_convolution[data_size- j][1] = 0.0;
	}

	return(0);
}

int PopulateDataArray(double* input_data, const int data_size, const int trial_number,
	const int padded_size, const int padding_type,
	fftw_complex* output_data)
{
	const double ramp = 2.0/data_size; // = 1.0/ n / 2
	double gain;
	int i;

	int output_counter = 0;
	int input_counter = 0;

	switch(padding_type)
	{
		case 0: //No Padding what so ever
			//populate the FFTW data vector.
			for (i = 0; i < data_size; ++i)
		    {
		    	output_data[i][0] = input_data[trial_number * data_size + i];
		    	output_data[i][1] = 0.0;
		    }
		    break;

		case 1: //Zero - Padding
			//populate the FFTW data vector.
			for (i = 0; i < data_size; ++i)
		    {
		    	output_data[i][0] = input_data[trial_number * data_size + i];
		    	output_data[i][1] = 0.0;
		    }

		    //Horse the rest of the data vector to zero just in case
		    for (i = data_size; i < padded_size; ++i)
		    {
		    	output_data[i][0] = 0.0;
		    	output_data[i][1] = 0.0;
		    }
		    break;

		case 2: //Duplicate array and ramp up and ramp down output
			for (i = 0; i < data_size; ++i)
		    {
		    	output_data[i][0] = input_data[trial_number * data_size + i];
		    	output_data[i][1] = 0.0;
		    }

		    for (i = 0; i < (int) (0.5* data_size); ++i)
		    {
		    	output_counter = data_size + i;
		    	input_counter = (int) (trial_number * data_size + (0.5*data_size) + i);
		    	gain = i * ramp;

		    	output_data[output_counter][0] = (1.0 - gain) * input_data[input_counter];
		    	output_data[output_counter][1] = 0.0;

		    	output_data[output_counter + (int) (0.5 * data_size)][0] = gain * input_data[trial_number * data_size + i];
		    	output_data[output_counter + (int) (0.5 * data_size)][1] = 0.0;
		    }
			break;

		default: //Return just the array no padding.
			//populate the FFTW data vector.
			for (i = 0; i < data_size; ++i)
		    {
		    	output_data[i][0] = input_data[trial_number * data_size + i];
		    	output_data[i][1] = 0.0;
		    }
		    break;
	}
    return(padding_type);
}


#pragma once
#include <gsl/gsl_vector.h>
#include "BooleanNodes.h"
#include "BooleanDAG.h"
#include "NodeVisitor.h"
#include "VarStore.h"
#include <map>
#include "FloatSupport.h"
#include <glpk.h>

#include <iostream>
#include <fstream>

using namespace std;

class Line { // m(x - x0) + d < 0 (x0 is fixed for a particular evaluator, hence not stored in the line)
public:
	double d;
	gsl_vector* m;
	Line(double _d, gsl_vector* _m): d(_d), m(_m) {}
	~Line(void) {
		delete m;
	}

	double getOffset(const gsl_vector* x0) { // m.x0 - d
		double r = 0.0;
		for (int i = 0; i < x0->size; i++) {
			r += gsl_vector_get(x0, i) * gsl_vector_get(m, i);
		}
		return r - d;
	}

	double getCoeff(int i) {
		return gsl_vector_get(m, i);
	}

	double getDist() {
		double r = 0;
		for (int i = 0; i < m->size; i++) {
			r += gsl_vector_get(m, i) * gsl_vector_get(m, i);
		}
		return d/sqrt(r);
	}

	string print() {
		stringstream str;
		str << "<";
		for (int i = 0; i < m->size; i++) {
			str << gsl_vector_get(m, i) << " ";
		}
		str << d << ">";
		return str.str();
	}
};

class Region { // Intersection of lines
public:
	vector<Line*> lines;
	int size() {
		return lines.size();
	}
	Line* get(int i) {
		return lines[i];
	}
	void addLine(Line* l) {
		if (find(lines.begin(), lines.end(), l) == lines.end()) {
			lines.push_back(l);
		}
	}

	static Region* intersect(Region* r1, Region* r2) {
		Region* r = new Region();
		for (int i = 0; i < r1->size(); i++) {
			r->addLine(r1->get(i));
		}
		for (int i = 0; i < r2->size(); i++) {
			r->addLine(r2->get(i));
		}
		return r;
	}

	static Region* copy(Region* r1) {
		Region* r = new Region();
		for (int i = 0; i < r1->size(); i++) {
			r->addLine(r1->get(i));
		}
		return r;
	}

	double getMaxDist() {
		double d = - numeric_limits<double>::max();
		for (int i = 0; i < lines.size(); i++) {
			if (lines[i]->getDist() > d) {
				d = lines[i]->getDist();
			}
		}
		return d;
	}

	gsl_vector* getMaxGrad() {
		double d = - numeric_limits<double>::max();
		gsl_vector* g;
		for (int i = 0; i < lines.size(); i++) {
			if (lines[i]->d > d) {
				d = lines[i]->d;
				g = lines[i]->m;
			}
		}
		return g;

	}

	~Region(void) {
		for (int i = 0; i < lines.size(); i++) {
			delete lines[i];
		}
	}

	string print() {
		stringstream str;
		str << "{";
		for (int i = 0; i < lines.size(); i++) {
			str << lines[i]->print() << " ";
		}
		str << "}";
		return str.str();
	}

};

class Value {
public:
	double val;
	gsl_vector* grad;
	Value(double v, gsl_vector* g): val(v), grad(g) {}
	~Value(void) {
		delete grad;
	}

	double getOffset(const gsl_vector* x0) { // d - m.x0
		double r = 0.0;
		for (int i = 0; i < x0->size; i++) {
			r += gsl_vector_get(x0, i) * gsl_vector_get(grad, i);
		}
		return val - r;
	}

	double getCoeff(int i) {
		return gsl_vector_get(grad, i);
	}

	static Value* add(Value* v1, Value* v2);
	static Value* mult(Value* v1, Value* v2);
	static Value* div(Value* v1, Value* v2);
	static Value* neg(Value* v1);
	static Value* copy(Value* v1);
	static Value* arctan(Value* v1);
	static Value* tan(Value* v1);
	static Value* cos(Value* v1);
	static Value* sin(Value* v1);
	static Value* sqrt(Value* v1);
	static Value* square(Value* v1);

	string print() {
		stringstream str;
		str << "<";
		for (int i = 0; i < grad->size; i++) {
			str << gsl_vector_get(grad, i) << " ";
		}
		str << val << ">";
		return str.str();
	}
};

class RegionValuePair {
public:
	Region* r;
	Value* v;
	RegionValuePair(Region* _r, Value* _v): r(_r), v(_v) {}

	string print() {
		stringstream str;
		str << r->print() << " -> " << v->print();
		return str.str();
	}
};

typedef vector<RegionValuePair*> ValuesList;

class LP {
	glp_prob* lp;
	int ia[1+1000], ja[1+1000];
	double ar[1+1000];
	double LOWBND = -32.0; // TODO: don;t hardcode
	double HIGHBND = 32.0;
	double PRECISION = 1e-5;

public:
	double optimize(RegionValuePair* rv, const gsl_vector* x0) {
		lp = glp_create_prob();
		glp_set_obj_dir(lp, GLP_MIN);
		int numrows = rv->r->size();
		if (numrows > 0) {
			glp_add_rows(lp, numrows);
			for (int i = 0; i < numrows; i++) {
				glp_set_row_name(lp, (i+1), "r" + (i+1));
				double c = rv->r->lines[i]->getOffset(x0);
				glp_set_row_bnds(lp, (i+1), GLP_UP, 0.0, c - PRECISION);
			}
		}
		int numcols = x0->size;
		glp_add_cols(lp, numcols);
		for (int i = 0; i < numcols; i++) {
			glp_set_col_name(lp, (i+1), "x" + (i+1));
			glp_set_col_bnds(lp, (i+1), GLP_DB, LOWBND, HIGHBND);
			glp_set_obj_coef(lp, (i+1), rv->v->getCoeff(i));
		}
		int counter = 1;
		for (int i = 0; i < numrows; i++) {
			for (int j = 0; j < numcols; j++) {
				ia[counter] = (i+1);
				ja[counter] = (j+1);
				ar[counter] = rv->r->lines[i]->getCoeff(j);
				counter++;
			}
		}
		Assert(counter-1 <= 1000, "Out of bounds");

		glp_load_matrix(lp, counter-1, ia, ja, ar);
		glp_simplex(lp, NULL);
		int status = glp_get_prim_stat(lp);
		cout << "status: " << status << endl;
		if (status == GLP_FEAS) {
			double z = glp_get_obj_val(lp);
			cout << "z = " << z;
			for (int i = 0; i < numcols; i++) {
				cout << " x" << i << " = " << glp_get_col_prim(lp, (i+1));
			}
			glp_delete_prob(lp);
			return z + rv->v->getOffset(x0);
		} else {
			return 1000;
		}

	}

	bool check(Region* r, const gsl_vector* x0) {
		int numrows = r->size();
		if (numrows == 0) return true;

		lp = glp_create_prob();
		glp_set_obj_dir(lp, GLP_MIN);
		glp_add_rows(lp, numrows);
		for (int i = 0; i < numrows; i++) {
			stringstream str;
			str << "r" << (i+1);
			glp_set_row_name(lp, (i+1), str.str().c_str());
			double c = r->lines[i]->getOffset(x0);
			glp_set_row_bnds(lp, (i+1), GLP_UP, 0.0, c-PRECISION);
		}
		int numcols = x0->size;
		glp_add_cols(lp, numcols);
		for (int i = 0; i < numcols; i++) {
			stringstream str;
			str << "x" << (i+1);
			glp_set_col_name(lp, (i+1), str.str().c_str());
			glp_set_col_bnds(lp, (i+1), GLP_DB, LOWBND, HIGHBND);
			glp_set_obj_coef(lp, (i+1), 0.0);
		}
		int counter = 1;
		for (int i = 0; i < numrows; i++) {
			for (int j = 0; j < numcols; j++) {
				ia[counter] = (i+1);
				ja[counter] = (j+1);
				ar[counter] = r->lines[i]->getCoeff(j);
				counter++;
			}
		}

		glp_load_matrix(lp, counter-1, ia, ja, ar);
		glp_smcp param;
		glp_init_smcp(&param);
		param.msg_lev = GLP_MSG_ERR;
		glp_simplex(lp, &param);
		int status = glp_get_prim_stat(lp);
		glp_delete_prob(lp);
		if (status == GLP_FEAS) {
			return true;
		} else {
			return false;
		}
	}

};

class GlobalEvaluator: NodeVisitor
{
	FloatManager& floats;
	BooleanDAG& bdag;
  map<string, int> floatCtrls; // Maps float ctrl names to indices with grad vectors
  int nctrls; // number of float ctrls
	gsl_vector* ctrls; // Ctrl values
	vector<ValuesList> values; // Keeps track of values for each node for different regions
	map<int, int> inputValues; // Maps node id to values set by the SAT solver
	double error = 0.0;
	gsl_vector* errorGrad;

	int DEFAULT_INP = -1;
	static gsl_vector* tmp;
	LP* lp;

	int MAX_REGIONS = 16;
	int ctr = 0;
public:
  GlobalEvaluator(BooleanDAG& bdag_p, FloatManager& _floats, const map<string, int>& floatCtrls_p);
  ~GlobalEvaluator(void);

  virtual void visit( SRC_node& node );
  virtual void visit( DST_node& node );
  virtual void visit( CTRL_node& node );
  virtual void visit( PLUS_node& node );
  virtual void visit( TIMES_node& node );
  virtual void visit( ARRACC_node& node );
  virtual void visit( DIV_node& node );
  virtual void visit( MOD_node& node );
  virtual void visit( NEG_node& node );
  virtual void visit( CONST_node& node );
  virtual void visit( LT_node& node );
  virtual void visit( EQ_node& node );
	virtual void visit( AND_node& node );
	virtual void visit( OR_node& node );
	virtual void visit( NOT_node& node );
  virtual void visit( ARRASS_node& node );
  virtual void visit( UFUN_node& node );
  virtual void visit( TUPLE_R_node& node );
	virtual void visit( ASSERT_node& node );

  double run(const gsl_vector* ctrls_p, map<int, int>& inputValues_p, gsl_vector* errorGrad_p);
	void truncate(ValuesList& vl, const gsl_vector* x0, int numRegions);
	void truncate(vector<Region*>& vl, const gsl_vector* x0, int numRegions);
	void truncateUnfeasibleRegions(ValuesList& vl, const gsl_vector* x0);
	void truncateUnfeasibleRegions(vector<Region*>& vl, const gsl_vector* x0);
	void truncateFarRegions(ValuesList& vl, const gsl_vector* x0, int numRegions);
	void truncateFarRegions(vector<Region*>& vl, const gsl_vector* x0, int numRegions);

	void setvalues(bool_node& bn, ValuesList& v) {
		values[bn.id] = v;
	}

	ValuesList& getvalues(bool_node& bn) {
		return values[bn.id];
	}

	ValuesList& getvalues(bool_node* bn) {
		return getvalues(*bn);
	}

	void print() {
    for (int i = 0; i < bdag.size(); i++) {
      cout << bdag[i]->lprint() << endl;
		}
  }

	bool isFloat(bool_node& bn) {
		return (bn.getOtype() == OutType::FLOAT);
	}

	bool isFloat(bool_node* bn) {
		return (bn->getOtype() == OutType::FLOAT);
	}

	int getInputValue(bool_node& bn) {
		if (inputValues.find(bn.id) != inputValues.end()) {
			int val = inputValues[bn.id];
			Assert(val == 0 || val == 1, "NYI: Integer values");
			return val;
		} else {
			return DEFAULT_INP;
		}
	}

	int getInputValue(bool_node* bn) {
		return getInputValue(*bn);
	}

	static gsl_vector* default_grad(int nctrls) {
		gsl_vector* g = gsl_vector_alloc(nctrls);
		for (int i = 0; i < nctrls; i++) {
			gsl_vector_set(g, i, 0);
		}
		return g;
	}

	static gsl_vector* add_grad(gsl_vector* g1, gsl_vector* g2);
	static gsl_vector* mult_grad(double v1, double v2, gsl_vector* g1, gsl_vector* g2);
	static gsl_vector* div_grad(double v1, double v2, gsl_vector* g1, gsl_vector* g2);
	static gsl_vector* neg_grad(gsl_vector* g1);
	static gsl_vector* copy_grad(gsl_vector* g1);
	static gsl_vector* sub_grad(gsl_vector* g1, gsl_vector* g2);
	static gsl_vector* arctan_grad(double v1, gsl_vector* g1);
	static gsl_vector* tan_grad(double v1, gsl_vector* g1);
	static gsl_vector* cos_grad(double v1, gsl_vector* g1);
	static gsl_vector* sin_grad(double v1, gsl_vector* g1);
	static gsl_vector* sqrt_grad(double v1, gsl_vector* g1);
	static gsl_vector* square_grad(double v1, gsl_vector* g1);

};