#include "R.h"
#include "Rmath.h"
#include "R_ext/Applic.h"
#include "blas1_d.h"
#include "linpack_d.h"

int n, ngrid, toprint, nperm, **perm;
double *z, *theta, *r, frac;
double *fmix, *fmix0, *fnew, *gqwt, *par_mv, *L1, *A1;

double Fp(double x){
//	return 0.5 + 0.5 * plogis(x, 0.0, 1.0, 1, 0);
	return plogis(x, 0.0, 1.0, 1, 0);
}

double fp(double x){
//	return log(0.5) + dlogis(x, 0.0, 1.0, 1);
	return dlogis(x, 0.0, 1.0, 1);
}

double psi(double x, double mu){
	return dnorm(x, mu, 1.0, 0);
}

static double * vect(int n){
	return (double *)R_alloc(n, sizeof(double));
}

static double ** mymatrix(int nr, int nc){
	int   i;
	double **m;
	
	m = (double **) R_alloc(nr, sizeof(double *));
	for (i = 0; i < nr; i++)
		m[i] = (double *) R_alloc(nc, sizeof(double));
	return m;
}

static int * ivect(int n){
	return (int *)R_alloc(n, sizeof(int));
}

static int ** imymatrix(int nr, int nc){
	int   i;
	int **m;
	
	m = (int **) R_alloc(nr, sizeof(int *));
	for (i = 0; i < nr; i++)
		m[i] = (int *) R_alloc(nc, sizeof(int));
	return m;
}

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

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

double int_gq(double *f, double *x, double *w, int n){
	double v = 0.0;
	int i = 0;
	for(i = 0; i < n; i++)
		v += w[i] * f[i];
	return v;
}

double int_trape(double *f, double *x, int n){
	double v = 0.0;
	int i = 0;
	for(i = 1; i < n; i++)
		v += 0.5 * (x[i] - x[i - 1]) * (f[i] + f[i - 1]);
	return v;
}

double inprod(double *x, double *y, int n){
	int i;
	double ip = 0.0;
	for(i = 0; i < n; i++)
		ip += x[i] * y[i];
	return ip;
}

void hessmat(int n, double *par, optimgr grad, void *ex, double *gr1, double **hess){

	double eps = sqrt(DBL_EPSILON), parHold;
	double *gr_plus = vect(n), *gr_minus = vect(n);
	int i, j;

	grad(n, par, gr1, ex);
	if(toprint){
		Rprintvec("gr1 = ", "%g ", gr1, n);
	}
	
	for(i = 0; i < n; i++){
		parHold = par[i];
		par[i] += eps;
		grad(n, par, gr_plus, ex);
		par[i] -= eps;
		par[i] -= eps;
		grad(n, par, gr_minus, ex);
		for(j = 0; j < n; j++)
			hess[i][j] = 0.5 * (gr_plus[j] - gr_minus[j]) / eps;
		par[i] = parHold;
	}	
}

double norm_lp(double *x, int len, double exponent){
	
  int i;
  double xnorm = atof("-Inf"), absx;
	
  if(exponent == atof("Inf")){
		
    for(xnorm = fabs(x[0]), i = 1; i < len; i++)
      if((absx = fabs(x[i])) > xnorm)
				xnorm = absx;
		
  } else if(exponent == 0.0){
    
    for(xnorm = 0.0, i = 0; i < len; i++)
      xnorm += (x[i] != 0.0);
    
  } else if(exponent > 0.0){
		
    for(xnorm = 0.0, i = 0; i < len; i++)
      xnorm += pow(fabs(x[i]), exponent);
		
    xnorm = pow(xnorm, 1.0 / exponent);
  } else
    printf("Error in norm_lp: 'exponent' must be non-negative\n");
	
	
  return xnorm;
	
}

void rperm(int size, int *seq, int n){
  int i, j, temp;
  
  for(i = 0;i < size;i++){                                            /* At the i-th step of iteration               */
    j = i + floor((n - i) * runif(0.0,1.0));                          /* randomly select a number from <i,..,n - 1>  */
    temp = *(seq + i);                                                /* Switch seq[i] and seq[j]                    */
    *(seq + i) = *(seq + j);                                          /* The firs i-elements are the chosen sample   */
    *(seq + j) = temp;                                                /* The last ones form the potential candidates */
  }
}

void getf(int q, double *par, double *fest, double *piest){
	
	int i, j, l;
	double mu = par[0], sigma = exp(par[1]), pi0, tau = 1.0 + exp(par[3]), gam = Fp(par[4]);
	double mu0 = mu, sigma0 = sigma;
	double temp, L0, h, u, A0, B, v, wts, zstd;	

	piest[0] = 0.0;

	for(j = 0; j < ngrid; j++)
		fest[j] = 0.0;

	for(l = 0; l < nperm; l++){
		
		for(j = 0; j < ngrid; j++)
			fmix[j] = fmix0[j];
		pi0 = Fp(par[2]);
		
		for(i = 0; i < n; i++){
			
			//				zstd = (z[perm[l][i]] - mu) / sigma;
			
			L0 = dnorm(z[perm[l][i]], mu, sigma, 0);
			for(j = 0; j < ngrid; j++){
				L1[j] = dnorm(z[perm[l][i]], mu0 + tau * sigma0 * theta[j], sigma, 0);
				fnew[j] = fmix[j] * L1[j];
			}
			
			h = int_gq(fnew, theta, gqwt, ngrid);
			u = pi0 * L0 + (1.0 - pi0) * h;
			wts = exp(-gam * r[i]);
			
			A0 = 1.0 + wts * (L0 / u - 1.0);
			for(j = 0; j < ngrid; j++)
				A1[j] = 1.0 + wts * (L1[j] / u - 1.0);			
			B = (1.0 - pi0) / (1.0 - A0 * pi0);
			
			pi0 *= A0;
			for(j = 0; j < ngrid; j++)
					fmix[j] *= B * A1[j];			
		}
		
		piest[0] += pi0 / ((double)nperm);
		for(j = 0; j < ngrid; j++)
			fest[j] += (1.0 - pi0) * fmix[j] / ((double)nperm);
		
	}

	for(j = 0; j < ngrid; j++)
		fest[j] /= (1.0 - piest[0]);
	
}


double prml_f(int q, double *par, void *ex){

	int i, j, l;
	double mu = par[0], sigma = exp(par[1]), tau = 1.0 + exp(par[3]), pi0, gam = Fp(par[4]);
	double mu0 = mu, sigma0 = sigma;
	double temp, L0, h, u, A0, B, wts, zstd;	
	
	double v = 0.0, pi0est = 0.0;
	for(l = 0; l < nperm; l++){
		
		for(j = 0; j < ngrid; j++)
			fmix[j] = fmix0[j];
		pi0 = Fp(par[2]);
		
		temp = 0.0;		
		for(i = 0; i < n; i++){
			
			L0 = dnorm(z[perm[l][i]], mu, sigma, 0); 
			for(j = 0; j < ngrid; j++){
				L1[j] = dnorm(z[perm[l][i]], mu0 + tau * sigma0 * theta[j], sigma, 0); 
				fnew[j] = fmix[j] * L1[j];
			}
			
			h = int_gq(fnew, theta, gqwt, ngrid);
			u = pi0 * L0 + (1.0 - pi0) * h;
			temp -= log(u);
			wts = exp(-gam * r[i]);
				
			A0 = 1.0 + wts * (L0 / u - 1.0);
			for(j = 0; j < ngrid; j++)
				A1[j] = 1.0 + wts * (L1[j] / u - 1.0);			
			B = (1.0 - pi0) / (1.0 - A0 * pi0);
			
			pi0 *= A0;
			for(j = 0; j < ngrid; j++)
				fmix[j] *= B * A1[j];		
		}		
		v += temp / ((double)nperm);
		pi0est += pi0 / ((double)nperm);
	}
	
	double penalty = 0.0;
	penalty = (dnorm(mu, 0.0, 0.05 * sigma, 1) 
			   + dnorm(par[1], 0.0, 0.25, 1) 
			   + fp(par[2]) + dbeta(Fp(par[2]), 22.7, 1.0, 1) 
			   + 0.0
			   + dnorm(par[3], 0.0, 1.0, 1));
	return v + *(double *)ex - penalty;
}


void prml_g_num(int q, double *par, double *grad, void *ex){
	
	int i;
	double eps = sqrt(DBL_EPSILON);
	double f_atpar = prml_f(q, par, ex);
	for(i = 0; i < 5; i++)
		par_mv[i] = par[i];
	for(i = 0; i < q; i++){
		par_mv[i] += eps;
		grad[i] = (prml_f(q, par_mv, ex) - f_atpar) / eps;
		par_mv[i] = par[i];
	}
}


void prml_g_num2(int q, double *par, double *grad, void *ex){
	
	int i;
	double eps = sqrt(DBL_EPSILON);
	double f_minus, f_plus;
	for(i = 0; i < 5; i++)
		par_mv[i] = par[i];
	for(i = 0; i < q; i++){
		par_mv[i] += eps;
		f_plus = prml_f(q, par_mv, ex);
		par_mv[i] = par[i] - eps;
		f_minus = prml_f(q, par_mv, ex);
		grad[i] = 0.5 * (f_plus - f_minus) / eps;
		par_mv[i] = par[i];
	}
}

void prml_mtest_ad(int *dims, double *zscores, double *par, int *mask,
				   double *thetaPt, double *gqPt, double *f0Pt, double *logLik, 
				   double *gr1, double *hess, double *fest, double *piest,
				   int *indprint, int *fghcount, double *constFact, 
				   int *permDetails){
	
	n = dims[0];
	ngrid = dims[1];
	toprint = indprint[0];
		
	int i, j, l, q = 5;
	
	fmix = vect(ngrid);
	fnew = vect(ngrid);
	L1 = vect(ngrid);
	A1 = vect(ngrid);
	r = vect(n);
	z = zscores;

	
	// permutation matrix
	
	if(permDetails[0] < 0){
		nperm = permDetails[1];
		perm = imymatrix(nperm, n);
		GetRNGstate();
		for(l = 0; l < nperm; l++){
			for(i = 0; i < n; i++)
				perm[l][i] = i;
			rperm(n, perm[l], n);
		}
		PutRNGstate();
	} else {
		nperm = permDetails[0];
		perm = (int **)R_alloc(nperm, sizeof(int *));
		for(l = 0; l < nperm; l++)
			perm[l] = permDetails + 1 + n * l;
	}
	
		
	theta = thetaPt;
	fmix0 = f0Pt;	
//	Rprintvec("f0 = ", "%g ", fmix0, ngrid);
	gqwt = gqPt;

	for(i = 0; i < n; i++)
		r[i] = log(2.0 + (double)i);
	
	double Fmin;
	int fail;
	
	void *ex;
	ex = constFact;
	par_mv = vect(5);
	
	if(toprint){
		Rprintf("q = %d, ", q);
		Rprintveci("mask = ", "%d ", mask, q);
		Rprintvec("par.init = ", "%g ", par, q);
	}
	
	vmmin(q, par, &Fmin, prml_f, prml_g_num2,
		  1000, toprint, mask, 
		  1.0e-08, 1.0e-08, 10, ex,
		  fghcount, fghcount + 1, &fail);
	
	if(toprint){
		Rprintf("Fmin = %g, fcount = %d, gcount = %d\n", Fmin, fghcount[0], fghcount[1]);
		Rprintvec("par = ", "%g ", par, q);
	}
	
	// Begin Newton's iteration to finish off the optimization
	
	int maxit = 0, info, ipvt[q], job, newton_counter = 0;	
	double tol = 1.0e-08, grad_norm;
	double **hmat = (double **)R_alloc(q, sizeof(double *));

	
	for(i = 0; i < q; i++)
		hmat[i] = hess + q * i;
	
	// get current gradient and hessian

	if(maxit > 0){
		hessmat(q, par, prml_g_num2, ex, gr1, hmat);
		grad_norm = norm_lp(gr1, q, atof("Inf"));
	}
	
	while((grad_norm > tol) & (newton_counter < maxit)){
	
		newton_counter++;
		
		// Compute Newton direction
	
		info = dgefa ( hess, q, q, ipvt );	
		if ( info != 0 ){
			printf ( "  DGEFA returned an error flag INFO = %d\nNewton recursion stopped.\n", info );
			break;
		} else {
			job = 0;
			dgesl ( hess, q, q, ipvt, gr1, job );
		}	
		
		// update par
		for(i = 0; i < q; i++)
			par[i] -= gr1[i];

		// get current gradient and hessian
		
		hessmat(q, par, prml_g_num2, ex, gr1, hmat);
		grad_norm = norm_lp(gr1, q, atof("Inf"));
		
	}
	if(toprint){
		Rprintf("Done newton with iterations = %d\n", newton_counter);
	}
	// get hessian one last time
	// get current gradient and hessian

	fghcount[2] = newton_counter;
	*logLik = prml_f(q, par, ex) - *constFact;
	hessmat(q, par, prml_g_num2, ex, gr1, hmat);
	getf(q, par, fest, piest);
	
}


void prml_mtest_eval(int *dims, double *zscores, double *par,
					 double *thetaPt, double *gqPt, double *f0Pt, double *logLik, 
					 int *indprint, double *constFact, int *permDetails){
	
	n = dims[0];
	ngrid = dims[1];
	int ncase = dims[2];
	toprint = indprint[0];
	
	
	int i, j, l, q = 5;
	
	fmix = vect(ngrid);
	fnew = vect(ngrid);
	L1 = vect(ngrid);
	A1 = vect(ngrid);
	r = vect(n);
	z = zscores;
	
	
	// permutation matrix
	
	if(permDetails[0] < 0){
		nperm = permDetails[1];
		perm = imymatrix(nperm, n);
		GetRNGstate();
		for(l = 0; l < nperm; l++){
			for(i = 0; i < n; i++)
				perm[l][i] = i;
			rperm(n, perm[l], n);
		}
		PutRNGstate();
	} else {
		nperm = permDetails[0];
		perm = (int **)R_alloc(nperm, sizeof(int *));
		for(l = 0; l < nperm; l++)
			perm[l] = permDetails + 1 + n * l;
	}
	
	
	theta = thetaPt;
	fmix0 = f0Pt;	
	gqwt = gqPt;
	
	for(i = 0; i < n; i++)
		r[i] = log(2.0 + (double)i);
	
	double Fmin;
	int fail;
	
	int *mask = (int *)R_alloc(q, sizeof(int));
	for(i = 0; i < q; i++)
		mask[i] = 1;	
	void *ex;
	ex = constFact;
	
	for(i = 0; i < ncase; i++)
		logLik[i] = prml_f(q, par + i * q, ex) - *constFact;
	
	
}


void dmix(int *dims, double *x, double *par, double *theta, double *gqwt, 
					double *fest, double *dens){
	int n = dims[0];
	int ngrid = dims[1];
	
	int i, j;
	double mu = par[0], sigma = par[1], tau = par[2];
	double mu0 = mu, sigma0 = sigma;
	double *df = vect(ngrid);
	
	for(i = 0; i < n; i++){
		for(j = 0; j < ngrid; j++)
			df[j] = dnorm(x[i], mu0 + tau * sigma0 * theta[j], sigma, 0) * fest[j];
		dens[i] = int_gq(df, theta, gqwt, ngrid);
	}
	
	
}