/*  utility.c contains utility codes to be called from other c programs
 *
 *  Created by Surya Tokdar (Sep 2010).
 *  This code can be modified and distributed with proper
 *  acknowledgement given to the original creator.
 *
 */


#include "standard.h"
#include "utility.h"
#define min(a,b) a < b ? a : b

//-- sigmoid function ---
double sigmoid(double x){
  return 1.0 / (1.0 + exp(-x));
}


//--- sigmoid function remapped to (-1,1) 
double sigmoid2(double x){
  return 2.0 / (1.0 + exp(-x)) - 1.0;
}


//--- Solve a * x^2 + b * x + c = 0 ---
void quadsolve(double a, double b, double c, double *x, int code){
	double discr = b * b - 4.0 * a * c;
	if(discr < 0)
		Rprintf("Negative discriminant, complex solutions\n");
	else{
		discr = sqrt(discr);
		if(code == 0)
			x[0] = (-b + discr) / (2.0 * a);
		else if(code == 1)
			x[0] = (-b - discr) / (2.0 * a);
		else{
			x[0] = (-b + discr) / (2.0 * a);
			x[1] = (-b - discr) / (2.0 * a);
		}
	}
}

//--- extract fractional part of a double ---
double fracf(double x){
	double ix;
	modf(x, &ix);
	return x - ix;
}

//--- fold a number to (0,1) by boundary reflection ---
double fold(double x){
	double ix;
	modf(x, &ix);
	x = fabs(x - ix);
	if((int)ix % 2)
		x = 1.0 - x;
	return x;
}

//--- add in the log scale ---
double logadd(double lx, double ly){
	double val;
	if(lx < ly)
		val = ly + log1p(exp(lx - ly));
	else 
		val = lx + log1p(exp(ly - lx));
	
	return val;
}

//--- swap ---
void swap(double *x, int i, int j){
	double v = x[i];
	x[i] = x[j];
	x[j] = v;
}

//---- Sum ----
double sum(double *x, int n){

  int i;
  double s;
  for(s = 0.0, i = 0; i < n; i++)
    s += x[i];

  return s;
}

//---- Sum of logs ----
double sumlog(double *x, int n){
	
  int i;
  double s;
  for(s = 0.0, i = 0; i < n; i++)
    s += log(x[i]);
	
  return s;
}

//--- Sum of squares ---
double ss(double *vec, int len){
  int i;
  double ss = 0.0;         
  for(i = 0; i < len; i++) 
    ss += vec[i] * vec[i]; 
  return ss;
}

//--- Finds the position of the minimum ---
int fwhich_min(double *x, int n){
  int i = 0, j = 0;
  double xmin = x[0];
  for(i = 1; i < n; i++){
    if(x[i] < xmin){
      xmin = x[i];
      j = i;
    }
  }

  return j;
}

//--- Finds the maximum ---
double f_max(double *x, int n){
  int i;
  double v;
  for(v = x[0], i = 1; i < n; i++)
    if(x[i] > v)
      v = x[i];

  return v;
}

//--- Finds the minimum ---
double f_min(double *x, int n){
  int i;
  double v;
  for(v = x[0], i = 1; i < n; i++)
    if(x[i] < v)
      v = x[i];

  return v;
}


//--- Random number generator from a discrete probability vector ---
int rdisc(int n, double *prob, int islog){
  int i, i0;
  double u, probmax, csump[n];

  if(islog){
		probmax = f_max(prob, n);
    csump[0] = exp(prob[0] - probmax);
    for(i = 1; i < n; i++) 
			csump[i] = csump[i - 1] + exp(prob[i] - probmax);
  } else{
    csump[0] = prob[0];
    for(i = 1;i < n;i++) csump[i] = csump[i - 1] + prob[i];
  }	
  if(csump[n - 1] == 0){
    i0 = 0;
  } else{
    u = runif(0.0,1.0);
    i = 0;
    while(csump[i] <= u * csump[n - 1]) i++;
    i0 = i;
  }
  return i0;
}

//--- Random permutation of a string of integers ---
void rperm(int size, int *seq, int n){
  int i, j, temp;
  
  for(i = 0; i < size; i++){                
    j = i + floor((n - i) * runif(0.0,1.0));
    temp = *(seq + i);      
    *(seq + i) = *(seq + j);
    *(seq + j) = temp;     
  }
}

//--- Generating from Dirichlet(k, rep(shape, k)) ---
void r_simple_dirichlet(double *x, int k, double *shape){

  int i;
  double xtot;

  for(i = 0; i < k; i++)
    x[i] = rgamma(shape[i], 1.0);    
  xtot = sum(x, k);
  if(xtot == 0.0){
    x[0] = 1.0;
    xtot = 1.0;
  }
  for(i = 0; i < k; i++)
    x[i] /= xtot;
}


//--- L-2 norm ---
double norm2(double *vec, int len){
  int i;
  double c_norm = 0.0;       
  for(i = 0; i < len; i++) 
    c_norm += pow(vec[i],2);
  return sqrt(c_norm);      
}

/********************************************************************
 ==================================================================== 
 Extreme gamma distribution
 
 All calculations are based on the following definition
   
   x ~ ExGam(shape, scale)
   
   IFF
   
   x = scale * log(1 + z)
   z ~ Gam(shape, 1)
 ====================================================================
*********************************************************************/

//--- density function ---
double ldexgam(double x, double shape, double scale){
  double z, lf;
  z = x / scale;
  lf = -log(scale) + z + dgamma(expm1(z), shape, 1.0, 1); 
  return lf;
}

//--- CDF ---
double pexgam(double x,double shape,double scale){
  double z, p;
  z = x/scale;
  p = pgamma(expm1(z),shape,1,1,0);
  return p;
}

//--- Quantile ---
double qexgam(double p,double shape,double scale){
  return scale * log(1 + qgamma(p, shape, 1, 1, 0));
}

//--- Random sample ---
double rexgam(double shape, double scale){
  return scale * log(1 + rgamma(shape, 1.0));
}


//--- Cholesky Factorization of symmetric pd matrix ---
void chfact(double **A, int p){

  // This rewrites the upper triangular part of A
  // with its cholesky factor. A[i][j] for j >= i
	// are the cholesky elements, whereas for j < i
	// they're the original elements of A
  
  int i, j, k;
  for(i = 0; i < p; i++){
    for(k = 0; k < i; k++)
      A[i][i] -= pow(A[k][i], 2.0);
		if(A[i][i] < 0){
			Rprintf("A[i][i] is negative, quitting\n");
			exit(EXIT_FAILURE);
		}
		A[i][i] = sqrt(A[i][i]);
		
		for(j = i + 1; j < p; j++){
			for(k = 0; k < i; k++)
				A[i][j] -= A[k][i] * A[k][j];
			A[i][j] /= A[i][i];
		}  
  }
}


//--- Backsolving Ax = b where A is upper triangular ---
void Usolve(double **A, int p, double *b, double *x, int transpose){
	
  int i, j;
  if(transpose){
    for(j = 0; j < p; j++){
			for(x[j] = b[j], i = 0; i < j; i++)
				x[j] -= x[i] * A[i][j];
			x[j] /= A[j][j];
		}
  } else {	
    for(j = p - 1; j >= 0; j--){
      for(x[j] = b[j], i = j + 1; i < p; i++) 
				x[j] -= A[j][i] * x[i];
			x[j] /= A[j][j];
		}
  }	  
}



//--- Compute b = Ax where A is upper traingular ---
void Uprod(double **A, int n, double *x, double *b, int transpose){

  int i, j;
  if(transpose){
    for(i = 0; i < n; i++)
	  for(b[i] = 0.0, j = 0; j <= i; j++)
	    b[i] += A[j][i] * x[j];
  } else{
    for(i = 0; i < n; i++)
      for(b[i] = 0.0, j = i; j < n; j++)
        b[i] += A[i][j] * x[j];
  }
}

//--- transpose of a matrix ---
void transpose(double **A, int m, int n, double **B){
	
	int i, j;
	for(i = 0; i < n; i++)
		for(j = 0; j < m; j++)
			B[i][j] = A[j][i];
}

//--- crossproduct t(A)b ---
void cprod(double **A, double *b, int m, int n, double *x){
	int i, j;
	for(i = 0; i < n; i++){
		for(x[i] = 0.0, j = 0; j < m; j++){
			x[i] += A[j][i] * b[j];	
		}
	}
}


//--- trace of a matrix ---
double trace(double **a, int n, int is_log){
	int i;
	double ds = 0.0;
	if(is_log){
		for(i = 0; i < n; i++)
			ds += log(a[i][i]);
	} else {
		for(i = 0; i < n; i++)
			ds += a[i][i];
	}
	
	return ds;
}

//--- inner product of two vectors ---
double inprod(double *v1, double *v2, int d){  
	
  int i;
  double val = 0.0;
  for(i = 0; i < d; i++) 
    val += v1[i] * v2[i];
  return val;
}


//--- Numerical integration via Simpson's method --- 
double int_simp(double *f, double h, int n){
  int i;
  double sum_odd = 0.0, sum_even = 0.0, out;
  for(i = 0; i < n; i++){
    sum_even += *(f + 2*i);
    sum_odd += *(f + 2*i + 1);
  }
  out = (h/3) * ( *(f + 2*n) - *f + 2*sum_even + 4*sum_odd);
  return out;
}

//--- Numerical integration via Romberg's method ---
double int_romberg(double *f){
  int l, k;
  double *t_seq, val, width = 1.0, coeff = 0.0;

  t_seq = (double *)calloc(order + 1, sizeof(double));

	//	Create the first column of the Romberg array
	//	this is simply the trapezoidal approximations
	//	in successive order of refinement 
  

  t_seq[0] = width * (f[0] + f[twopow[order]])/2.0;

  for(l = 1; l <= order; l++){
    width = width/2.0;
    for(t_seq[l] = t_seq[l - 1]/2, k = 0; k < twopow[l - 1]; k++) 
      t_seq[l] += width * f[(2 * k + 1) * twopow[order - l]];
  }

  // Fill in the other entries of the Romberg array

  for(l = 1; l <= order; l++){
    coeff = 1.0 / (pow((double)twopow[l],2.0) - 1.0);
    for(k = order; k >= l; k--)
      t_seq[k] -= coeff * (t_seq[k] - t_seq[k - 1]);
  }

  val = t_seq[order];
  free(t_seq);

  return val;
}

//--- Trapezoidal integration ---
double int_trape(double *f){

  int l;
  double val = 0.0;

  for(l = 0; l < grid_len; l++)
    val += f[l];
  val -= 0.5 * (f[0] + f[grid_len - 1]);
  return grid_incr * val;
}


//-- print a double vector ---

void Rprintvec(char *a, double *x, int len){

  int i;
	Rprintf("%s", a);
  for(i = 0; i < len; i++)
    Rprintf("%g ", x[i]);

  Rprintf("\n");
}

//-- print an int vector ---

void Rprintveci(char *a, int *x, int len){
	
  int i;
	Rprintf("%s", a);
  for(i = 0; i < len; i++)
    Rprintf("%d ", x[i]);
	
  Rprintf("\n");
}

//-- print a double matrix ---

void Rprintmat(char *a, double **m, int d1, int d2, int iscol){

  int i, j;
	Rprintf("%s :\n", a);
  if(iscol){
    for(j = 0; j < d2; j++){
      for(i = 0; i < d1; i++)
	Rprintf("%7.3g ", m[i][j]);
      Rprintf("\n");
    }
  } else{
    for(i = 0; i < d1; i++){
      for(j = 0; j < d2; j++)
	Rprintf("%7.3g ", m[i][j]);
      Rprintf("\n");
    }
  }

}

//-- print a int matrix ---

void Rprintmati(char *a, int **m, int d1, int d2, int iscol){
	
  int i, j;
	Rprintf("%s :\n", a);
  if(iscol){
    for(j = 0; j < d2; j++){
      for(i = 0; i < d1; i++)
				Rprintf("%d\t", m[i][j]);
      Rprintf("\n");
    }
  } else{
    for(i = 0; i < d1; i++){
      for(j = 0; j < d2; j++)
				Rprintf("%d\t", m[i][j]);
      Rprintf("\n");
    }
  }
	
}


//--- Memory handling for double arrays (up to 4 dims) ---	
//--- consider using R_alloc instead

void *setmem(int ndim, ...){
	
	va_list ap;
	
	va_start(ap, ndim);	
	if(ndim == 1){
		
		int n = va_arg(ap, int);
		double *a = (double *)calloc(n, sizeof(double));
		va_end(ap);
		return a;
		
	} else if(ndim == 2){
		
		int i, n1, n2;
		n1 = va_arg(ap, int);
		n2 = va_arg(ap, int);
		double **a = (double **)calloc(n1, sizeof(double *));
		for(i = 0; i < n1; i++)
			a[i] = (double *)calloc(n2, sizeof(double));
		va_end(ap);
		return a;
		
	} else if(ndim == 3){
		
		int i, j, n1, n2, n3;
		
		n1 = va_arg(ap, int);
		n2 = va_arg(ap, int);
		n3 = va_arg(ap, int);
		
		double ***a = (double ***)calloc(n1, sizeof(double **));
		for(i = 0; i < n1; i++){
			a[i] = (double **)calloc(n2, sizeof(double *));
			for(j = 0; j < n2; j++)
				a[i][j] = (double *)calloc(n3, sizeof(double));
		}
		va_end(ap);
		return a;
		
	} else if(ndim == 4){
		
		int i, j, k, n1, n2, n3, n4;
		
		n1 = va_arg(ap, int);
		n2 = va_arg(ap, int);
		n3 = va_arg(ap, int);
		n4 = va_arg(ap, int);
		
		double ****a = (double ****)calloc(n1, sizeof(double ***));
		for(i = 0; i < n1; i++){
			a[i] = (double ***)calloc(n2, sizeof(double **));
			for(j = 0; j < n2; j++){
				a[i][j] = (double **)calloc(n3, sizeof(double *));
				for(k = 0; k < n3; k++)
					a[i][j][k] = (double *)calloc(n4, sizeof(double));
			}
		}
		va_end(ap);
		return a;
		
	} else {
	  
	  Rprintf("Cannot set memory for arrays with more than 4 dimensions\n");
		va_end(ap);
	  
	}
}

void freemem(int ndim, ...){
	
	va_list ap;
	
	va_start(ap, ndim);	
	if(ndim == 1){
		double *a = va_arg(ap, double *);
		free(a);
		va_end(ap);
	} else if(ndim == 2){
		
		int i, n1;
		double **a = va_arg(ap, double **);
		n1 = va_arg(ap, int);
		for(i = 0; i < n1; i++)
			free(a[i]);
		free(a);
		va_end(ap);
		
	} else if(ndim == 3){
		
		int i, j, n1, n2;
		double ***a = va_arg(ap, double ***);
		n1 = va_arg(ap, int);
		n2 = va_arg(ap, int);
		for(i = 0; i < n1; i++){
			for(j = 0; j < n2; j++)
				free(a[i][j]);
			free(a[i]);
		}
		free(a);
		va_end(ap);
		
	} else if(ndim == 4){
		
		int i, j, k, n1, n2, n3, n4;
		double ****a = va_arg(ap, double ****);		
		n1 = va_arg(ap, int);
		n2 = va_arg(ap, int);
		n3 = va_arg(ap, int);
		
		for(i = 0; i < n1; i++){
			for(j = 0; j < n2; j++){
				for(k = 0; k < n3; k++)
					free(a[i][j][k]);
				free(a[i][j]);
			}
			free(a[i]);
		}
		free(a);
		va_end(ap);
		
	} else {
	  
	  Rprintf("Cannot free memory for arrays with more than 4 dimensions\n");
		va_end(ap);
	  
	}
	
	
}

//--- Memory handling for int arrays (up to 4 dims) ---	
//--- consider using R_alloc instead ---

void *isetmem(int ndim, ...){
	
	va_list ap;
	
	va_start(ap, ndim);	
	if(ndim == 1){
		
		int n = va_arg(ap, int);
		int *a = (int *)calloc(n, sizeof(int));
		va_end(ap);
		return a;
		
	} else if(ndim == 2){
		
		int i, n1, n2;
		n1 = va_arg(ap, int);
		n2 = va_arg(ap, int);
		int **a = (int **)calloc(n1, sizeof(int *));
		for(i = 0; i < n1; i++)
			a[i] = (int *)calloc(n2, sizeof(int));
		va_end(ap);
		return a;
		
	} else if(ndim == 3){
		
		int i, j, n1, n2, n3;
		
		n1 = va_arg(ap, int);
		n2 = va_arg(ap, int);
		n3 = va_arg(ap, int);
		
		int ***a = (int ***)calloc(n1, sizeof(int **));
		for(i = 0; i < n1; i++){
			a[i] = (int **)calloc(n2, sizeof(int *));
			for(j = 0; j < n2; j++)
				a[i][j] = (int *)calloc(n3, sizeof(int));
		}
		va_end(ap);
		return a;
		
	} else if(ndim == 4){
		
		int i, j, k, n1, n2, n3, n4;
		
		n1 = va_arg(ap, int);
		n2 = va_arg(ap, int);
		n3 = va_arg(ap, int);
		n4 = va_arg(ap, int);
		
		int ****a = (int ****)calloc(n1, sizeof(int ***));
		for(i = 0; i < n1; i++){
			a[i] = (int ***)calloc(n2, sizeof(int **));
			for(j = 0; j < n2; j++){
				a[i][j] = (int **)calloc(n3, sizeof(int *));
				for(k = 0; k < n3; k++)
					a[i][j][k] = (int *)calloc(n4, sizeof(int));
			}
		}
		va_end(ap);
		return a;
		
	} else {
	  
	  Rprintf("Cannot set memory for arrays with more than 4 dimensions\n");
		va_end(ap);
	  
	}
	
	
	
}


void ifreemem(int ndim, ...){
	
	va_list ap;
	
	va_start(ap, ndim);	
	if(ndim == 1){
		int *a = va_arg(ap, int *);
		free(a);
		va_end(ap);
	} else if(ndim == 2){
		
		int i, n1;
		int **a = va_arg(ap, int **);
		n1 = va_arg(ap, int);
		for(i = 0; i < n1; i++)
			free(a[i]);
		free(a);
		va_end(ap);
		
	} else if(ndim == 3){
		
		int i, j, n1, n2;
		int ***a = va_arg(ap, int ***);
		n1 = va_arg(ap, int);
		n2 = va_arg(ap, int);
		for(i = 0; i < n1; i++){
			for(j = 0; j < n2; j++)
				free(a[i][j]);
			free(a[i]);
		}
		free(a);
		va_end(ap);
		
	} else if(ndim == 4){
		
		int i, j, k, n1, n2, n3, n4;
		int ****a = va_arg(ap, int ****);		
		n1 = va_arg(ap, int);
		n2 = va_arg(ap, int);
		n3 = va_arg(ap, int);
		
		for(i = 0; i < n1; i++){
			for(j = 0; j < n2; j++){
				for(k = 0; k < n3; k++)
					free(a[i][j][k]);
				free(a[i][j]);
			}
			free(a[i]);
		}
		free(a);
		va_end(ap);
		
	} else {
	  
	  Rprintf("Cannot free memory for arrays with more than 4 dimensions\n");
		va_end(ap);
	  
	}
	
	
}