/*
 *  pp_est.c
 *
 *  Created by Wei Sun on Thu Sep 20 2007.
 *
 */

#include "ttest.h"
#include "utility.h"
#include "pp_est.h"
#include <R.h>
#include <Rmath.h>


/*********************************************************************
 * 
 * p_r_est
 *
 * given one marker genotype profile mm, and one trait y, randomly
 * generate K genotype profile whose distance to mm is r1:r2,
 * calculate the K t-test p-values comparing the generated genotype
 * profiles with y
 * 
 *********************************************************************/
void p_r_est(double *y, int *mm, double *Rpd, int *dims){
	
	double **pd, *e0, *e1, t, pval;
	int r1 = dims[0];
	int r2 = dims[1];
	int K  = dims[2];
	int p  = dims[3];
	int i, j, u, v, le0, le1;
	int r = r2 - r1 + 1;
	int *md, *tc;
	long idum = (long)dims[4];
	
	md = (int *)calloc(p, sizeof(int));
	tc = (int *)calloc(p, sizeof(int));

  e0 = (double *)calloc(p, sizeof(double));
  e1 = (double *)calloc(p, sizeof(double));
  
	reorg(Rpd, &pd, r, K);

	for(i=0,j=r1; j<=r2; i++,j++){
    for(u=0; u<K; u++){
    	for(v=0; v<p; v++){
    		md[v] = mm[v];
    	}
      
      getPermute(tc, p, &idum);
      
      for(v=0; v<j; v++){
        md[tc[v]] = 1-mm[tc[v]];
      }
      
      le0 = 0;
      le1 = 0;
      for(v=0; v<p; v++){
        if(md[v]==0){
          e0[le0]=y[v];
          le0++;
        }else{
          e1[le1]=y[v];
          le1++;
        }
      }
      
      tutest(e0, &le0, e1, &le1, &t, &pval);
      pd[i][u] = pval;
    }
  }
  
  free(md);
  free(tc);
  free(e0);
  free(e1);
}

/*********************************************************************
 * 
 * p_p_est
 *
 * estimate permutation test p-value
 * 
 *********************************************************************/
 
void p_p_est(int *RmD, double *pcuts, double *Rpp, int *dims, int* trace,
  double *pers){
  
	int i, j, D0, D1;
	int n = dims[0];
  int p = dims[1];
  int rmax = dims[2];
  int npcuts = dims[3];
  int p2 = p/2;
  int **mD;
  double **pp, **cts, *ctsSum, *ctsT, tmp, total;
  
  if(*trace){
    Rprintf("n=%d, p=%d, rmax=%d, npcuts=%d\n", n, p, rmax, npcuts);
  }
  
  reorg_int(RmD, &mD, n, p);
  reorg(Rpp, &pp, npcuts, rmax+1);
  
  if(*trace){
    Rprintf("allocating memory...\n");
  }
  
  ctsT  = (double*) calloc(npcuts, sizeof(double));
  ctsSum  = (double*) calloc(rmax+1, sizeof(double));

	cts = (double**) malloc(n*sizeof(double*));
  cts[0] = (double*) calloc(n*(rmax+1), sizeof(double));
  for(i=1; i<n; i++){
    cts[i] = cts[i-1] + (rmax+1);
  }

  if(*trace){
    Rprintf("counts1...\n");
  }
  
  counts1(mD, n, p, rmax, cts);
  
  if(*trace){
    Rprintf("counts2...\n");
  }
  
  if(n==1){
  	for(j=0; j<=rmax; j++){
  		ctsSum[j] = cts[0][j];
  	}
  }else{
    counts2(mD, cts, n, p, rmax, ctsSum);
  }
  
  if(*trace > 1){
    Rprintf("\n ctsSum:\n");
    for(i=0; i<=rmax; i++){
      Rprintf("%.4f ", ctsSum[i]);
      if((i+1)%10==0){ Rprintf("\n"); }
    }
    Rprintf("\n");
  }
  
  if(*trace){
    Rprintf("counting...\n");
  }
  
  for(i=0; i<npcuts; i++){
  	
  	for(D0 = rmax; pp[i][D0]==0; D0--){}
  	for(D1 = 0; pp[i][D1]==1; D1++) {}

    ctsT[i] = ctsSum[D1-1];
    for(j = D1; j<=D0; j++){
      tmp = log_exp_sub(ctsSum[j], ctsSum[j-1]) + log(pp[i][j]);
      ctsT[i] = log_exp_add(ctsT[i], tmp);
    }
  }
  
  if(*trace){
    Rprintf("finishing up ...\n");
  }
  
  total = lchoose((double)p, (double)p2);
  
  for(i=0; i<npcuts; i++){
    pers[i] = 2*exp(ctsT[i] - total);
  }
  
  free(cts);
  free(ctsT);
  free(ctsSum);
}

/**********************************************************************
 * 
 * counts1
 * 
 * given a matrix of genotype profiles, with each row for one marker 
 * and each column for one individual, count genotype configurations 
 * for each marker with ditance from 0 to r, obtain the cumsum
 * 
 **********************************************************************/
void counts1(int** mD, int n, int p, int r, double **cts){
	int i, j, t, k1, *k1_used;
	double **cts_k1;
	int p2 = p/2;
	
	k1_used  = (int*) calloc(p2+1, sizeof(int));
	for(i=0; i<=p2; i++){
	  k1_used[i] = 0;
	}
	
	cts_k1 = (double**) malloc((p2+1)*sizeof(double*));
  cts_k1[0] = (double*) calloc((p2+1)*(r+1), sizeof(double));
  for(i=1; i<=p2; i++){
    cts_k1[i] = cts_k1[i-1] + (r+1);
  }
  
	for(i=0; i<n; i++){
		k1 = 0;
		for(j=0; j<p; j++){ k1 += mD[i][j]; }
		if(k1 > p2){ k1 = p-k1; }
		if(!k1_used[k1]){
		  count_1m((double)k1, (double)p, (double)r, cts_k1[k1]);
		}
		cts[i][0] = cts_k1[k1][0];
		for(t=1; t<=r; t++){
			cts[i][t] = log_exp_add(cts[i][t-1], cts_k1[k1][t]);
		}
	}
	free(k1_used);
	free(cts_k1);
}

/**********************************************************************
 * 
 * count_1m
 *
 * given one genotype profile m1
 * count genotype configurations with distance to m1 from 0 to r
 *
 * assume k1 <= p2
 * 
 **********************************************************************/
int count_1m(double k1, double p, double r, double *ct){
	int i, t, tmax;
	double d;
	double p2 = p/2.0;
	double k0 = p - k1;
	
	if(k1 > p2){ error("expect k1 <= p2\n in count_1m"); }
		
	for(i=0; i<=r; i++){
		ct[i] = -INFINITY;
	}
  
  d = k0 - p2;
  tmax = floor((r - d)/2);
  if(tmax < 0){ return(1); }
  	
  for(t=0; t<=tmax; t++){
    ct[(int)d+2*t] = lchoose(k0, (d+t)) + lchoose(k1, (double)t);
  }

	return(1);
}

/*********************************************************************
 * 
 * counts2
 * 
 * find those desired genotype profiles belong to the union of two 
 * hyper-spheres both with radius r
 * 
 *********************************************************************/
void counts2(int **mD, double **cts, int n, int p, int r, 
double *ctsSum){
  int i, j, k, *ds, k1, db0, da0, da, db;
  int p2 = p/2;
  double **cts_intersect, **cts_ab, tmp, tmp1;
  
  ds = (int*) calloc(n, sizeof(int));
  
  cts_intersect = (double**) malloc(n*sizeof(double*));
  cts_intersect[0] = (double*) calloc(n*(r+1), sizeof(double));
  for(i=1; i<n; i++){
    cts_intersect[i] = cts_intersect[i-1] + (r+1);
  }
  
  cts_ab = (double**) malloc((r+1)*sizeof(double*));
  cts_ab[0] = (double*) calloc((r+1)*(r+1), sizeof(double));
  for(i=1; i<=r; i++){
    cts_ab[i] = cts_ab[i-1] + (r+1);
  }
  
  for(i=0; i<n; i++){
  	k1 = 0;
  	for(j=0; j<p; j++){	k1 += mD[i][j]; }
    if(k1 >= p2){ ds[i] = k1 - p2; }else{ ds[i] = p2 - k1; }
  }
  
  /* calculate the intersection between mD[i,] and mD[i-1,] */
  for(i=1; i<n; i++){
  	db0 = ds[i];
  	if(db0 > r){ 
      for(k=0; k<=r; k++){
  		  cts_intersect[i][k] = -INFINITY;
      }
  	  continue; 
  	}
  	
  	j = i-1;
  	da0 = ds[j];
    while(da0 > r && j>0){
      j--;
      da0 = ds[j];
    }
    if(j==0 && da0 > r){
    	for(k=0; k<=r; k++){
  		  cts_intersect[i][k] = -INFINITY;
      }
  	  continue; 
    }

    /* initialize cts_ab first */
    for(da=0; da<=r; da++){
      for(db=0; db<=r; db++){
      	cts_ab[da][db] = -INFINITY;
      }
    }
  
    for(da=da0; da<=r; da+=2){
      for(db=db0; db<=r; db+=2){            
        cts_ab[da][db] = count_2m(mD[j], mD[i], p, da, db);
      }
    }
    
    cts_intersect[i][0] = cts_ab[0][0];
    
    for(k=1; k<=r; k++){
    	tmp = cts_ab[k][k];
    	for(j=0; j<k; j++){
        tmp = log_exp_add(tmp, cts_ab[k][j]);
        tmp = log_exp_add(tmp, cts_ab[j][k]);
      }
      cts_intersect[i][k] = log_exp_add(cts_intersect[i][k-1], tmp);
    }
  }
  
  for(k=0; k<=r; k++){
  	tmp  = cts[0][k];
  	tmp1 = cts_intersect[0][k];

  	for(i=1; i<n; i++){
  		tmp  = log_exp_add(tmp, cts[i][k]);
  		tmp1 = log_exp_add(tmp1, cts_intersect[i][k]);
  	}
  	/*
  	Rprintf("k=%d, tmp=%f, tmp1=%f\n", k, tmp, tmp1);
  	*/
    ctsSum[k] = log_exp_sub(tmp, tmp1);
  }
  
  free(ds);
  free(cts_intersect[0]);
  free(cts_intersect);
  free(cts_ab[0]);
  free(cts_ab);
}

/**********************************************************************
 * 
 * count_2m
 * 
 * given two genotype profiles ma and mb, count genotype configurations
 * m with d(m, ma) = da, d(m, mb) = db, and #{m==0} = p/2
 * 
 **********************************************************************/
double count_2m(int *ma, int *mb, int p, int da, int db){
	int i, s000, s010, s100, s110, c1, c2, c3, lB, uB;
	int p2 = p/2;
	double lowB[4];
	double uppB[5];
	double st, st1;
	int s00 = 0;
	int s01 = 0;
	int s10 = 0;
	int s11 = 0;
	
	if(da==0 && db==0){ return -INFINITY; }
				
	for(i=0; i<p; i++){
		if(ma[i]==0){
		  if(mb[i]==0){ s00++; } else{ s01++; }
		}else{
		  if(mb[i]==0){ s10++; } else{ s11++; }
		}
	}
  
  lowB[0] = 0;
  lowB[1] = 0.5*(p2 - da + s00 - s01);
  lowB[2] = 0.5*(p2 - db + s00 - s10);
  lowB[3] = s00 + 0.5*(s01 + s10 - da - db);
  uppB[0] = s00;
  uppB[1] = 0.5*(p2 - da + s00 + s01);
  uppB[2] = 0.5*(p2 - db + s00 + s10);
  uppB[3] = s11 + lowB[3];
  uppB[4] = p2;
  
  lB = ceil(d_max(lowB, 4));
  uB = floor(d_min(uppB, 5));
  
  /*
  for(i = 0; i<4; i++){
  	Rprintf("%.2f ", lowB[i]);
  }
  Rprintf("\n");
  
  for(i = 0; i<5; i++){
  	Rprintf("%.2f ", uppB[i]);
  }
  Rprintf("\n");
  
  Rprintf("lB=%d, uB=%d, s00=%d, s01=%d, s10=%d, s11=%d\n", lB, uB, s00, s01, s10, s11);
  */
  
  st = -INFINITY;
  if(lB > uB){ return st; }
  
  for(s000=lB; s000<=uB; s000++){
    c1 = p2 - s000;
    c2 = da - s00 - s01 + s000;
    c3 = db - s00 - s10 + s000;
    
    if((c1 - c2)%2 == 1){ continue; }
    if((c1 - c3)%2 == 1){ continue; }
    if((c2 + c3)%2 == 1){ continue; }

    s010 = (c1 - c2)/2;
    s100 = (c1 - c3)/2;
    s110 = (c2 + c3)/2;
                
    st1 = lchoose((double)s00, (double)s000) + 
          lchoose((double)s01, (double)s010) + 
          lchoose((double)s10, (double)s100) + 
          lchoose((double)s11, (double)s110);
            
    st = log_exp_add(st, st1);
  }
  return(st);
}
