
/*
   File:        hypergeo.c
   Authors:     Wong / Moore / Schneider
   Created:     Mon Oct 22 04:03:02 EDT 2001
   Description: Probability functions for a hypergeometric distribution

   Copyright 2001, The Auton Lab
*/

#include "amut.h"
#include "stats.h"
#include "choose.h"
#include "hypergeo.h"

extern choose_cache *Schoose_cache;

/*
  The bounds are:
  max(0,k-B) <= x <= min(k,A)
 */
static void calculate_bounds(int k, int A, int B, 
			     int* r_lower_bound, int* r_upper_bound)
{
  *r_lower_bound = (k-B) > 0 ? (k-B) : 0;
  *r_upper_bound = (k<A) ? k : A;
}

/*
double naive_hypergeometric_pdf(int x, int k, int A, int B, hg_choose_cache* cc)
{
  int lower_bound;
  int upper_bound;
  double result;
  double choose1;
  double choose2;
  double choose3;

  calculate_bounds(k,A,B,&lower_bound,&upper_bound);
  my_assert(x >= lower_bound);
  my_assert(x <= upper_bound);

  choose1 = log_choose(A,x,cc);
  choose2 = log_choose(B,k-x,cc);
  choose3 = log_choose(A+B,k,cc);
  result = choose1+choose2-choose3;
  result = exp(result);
  return result;
}
*/

double naive_hypergeometric_pval(int x, int k, int A, int B)
{
  int i;
  double* pf_array;
  int lower_bound;
  int upper_bound;
  double x_pf_val;
  double pval = 0;

  calculate_bounds(k,A,B,&lower_bound,&upper_bound);
  my_assert(x >= lower_bound);
  my_assert(x <= upper_bound);
  pf_array = AM_MALLOC_ARRAY(double,upper_bound-lower_bound+1);
  for( i = lower_bound; i <= upper_bound; i++ )
    {
      /*
      printf("Calling hypergeometric_pdf with %d, %d, %d, %d\n",
	     i,k,A,B);
      */
      pf_array[i-lower_bound] = hypergeometric_pdf(i,k,A,B);
      /* printf("pval[%d] = %.3f\n", i, pf_array[i-lower_bound]);*/
    }
  /*
    printf("Comparing against pf[%d] = %.20e\n",
	 x-lower_bound,
	 pf_array[x-lower_bound]);
  */
  x_pf_val = pf_array[x-lower_bound];
  for( i = lower_bound; i <= upper_bound; i++ )
    {
      /* printf("pf_array[%d] = %.20e\n", i, pf_array[i-lower_bound]); */
      if( pf_array[i-lower_bound] <= x_pf_val )
	{
	  pval += pf_array[i-lower_bound];
	}
      /*  
      printf("\n");
      */
    }
  AM_FREE_ARRAY(pf_array,double,upper_bound-lower_bound+1);
  /*
  printf("pval = %f\n", pval);
  */
  return pval;
}

/*
  The hypergeometric series is similar to having a box with A red balls
  and B blue balls.  Choose k balls at random without replacement.  Let X be 
  the number of red balls.  Then the probability of getting x red and k-x 
  blues is:
  f(x) = (A choose x)(B choose k-x)
         --------------------------
         (A+B choose k)
 */
double hypergeometric_pdf(int x, int k, int A, int B)
{
  int i;
  int lower_bound;
  int upper_bound;
  double result = 0;
  ivec* top1_exps;
  ivec* top1_factors;
  ivec* top2_exps;
  ivec* top2_factors;
  ivec* bot1_exps;
  ivec* bot1_factors;

  if (!Schoose_cache) 
    Schoose_cache = mk_choose_cache(CHOOSE_MAX_SIZE,CHOOSE_MAX_SIZE);

  calculate_bounds(k,A,B,&lower_bound,&upper_bound);
  my_assert(x >= lower_bound);
  my_assert(x <= upper_bound);

  mk_prime_factor_choose(A,x,&top1_factors,&top1_exps);
  mk_prime_factor_choose(B,k-x,&top2_factors,&top2_exps);
  mk_prime_factor_choose(A+B,k,&bot1_factors,&bot1_exps);

  /* We do it in this order to prevent overflow (although overflow is
   unlikely */
  prime_factor_ivec_subtract(top1_exps,bot1_exps,bot1_exps);
  prime_factor_ivec_plus(bot1_exps,top2_exps,bot1_exps);

  my_assert(ivec_size(bot1_factors) == ivec_size(bot1_exps));
  for( i = 0; i < ivec_size(bot1_factors); i++ )
    {
      result += (ivec_ref(bot1_exps,i) * log(ivec_ref(bot1_factors,i)));
    }
  free_ivec(top1_exps);
  free_ivec(top1_factors);
  free_ivec(top2_exps);
  free_ivec(top2_factors);
  free_ivec(bot1_exps);
  free_ivec(bot1_factors);
  return exp(result);
}

/*
   Let h(x; r, n, N) = (n choose x)(N-n choose r-x)
                       --------------------------
                         (N choose r)
   The function below will calculate h(0; r, n, N) ie.
   h(0; r, n, N) = (N-n)(N-n-1)...(N-n-r+1)
                   ------------------------
                   N(N-1)...(N-r+1)

   If you change this function, remember to change 
   hypergeometric_is_pval_lower() which does an early cutoff for the pval
   calculation if it is already known to be lower.
*/

double hypergeometric_pval(int x, int k, int A, int B, hg_pvaltype pvaltype)
{
  int i;
  int lower_bound;
  int upper_bound;
  double h_prev;
  double h;
  double pval = 0;
  double x_pf_val;

  calculate_bounds(k,A,B,&lower_bound,&upper_bound);
  my_assert(x >= lower_bound);
  my_assert(x <= upper_bound);
  x_pf_val = hypergeometric_pdf(x,k,A,B);
  h_prev = x_pf_val;
  pval = x_pf_val;

  /*printf("pf_array[%d] = %2.10e, pval = %2.10e\n", x, x_pf_val, pval);*/

  switch(pvaltype)
    {
    case hg_one_sided_decrease:
      /* 
         h(i-1) = i(B - k + i)           * h(i)
	          ----------------------
	          (A - i + 1)(k - i + 1)
      */
      for( i = x; i > lower_bound; i-- )
	{
	  h = h_prev * (double)(i * (B - k + i))/(double)((A - i + 1)*(k - i + 1));
	  /*printf("pf_array[%d] = %f\n", i-1, h);*/
	  h_prev = h;
	  pval += h;
	}
      break;
    case hg_one_sided_increase:
      for( i = x; i < upper_bound; i++ )
	{
	  /*
	    Note:
	    h(i+1) = h(i) * (A-i)(k-i)
	                    --------------
	                    (i+1)(B-k+i+1)
	  */
	  h = h_prev * (double)((A - i)*(k - i))/(double)((i+1)*(B - k + i + 1));
	  /*printf("pf_array[%d] = %f\n", i+1, h);*/
	  h_prev = h;
	  pval += h;
	}
      break;
    case hg_two_sided:
      /* 
         h(i-1) = i(B - k + i)           * h(i)
	          ----------------------
	          (A - i + 1)(k - i + 1)
      */
      for( i = x; i > lower_bound; i-- )
	{
	  h = h_prev * (double)(i * (B - k + i))/(double)((A - i + 1)*(k - i + 1));
	  h_prev = h;
	  if( h <= x_pf_val )
	    {
	      pval += h;
	      /*printf("*pf_array[%d] = %2.10e, pval = %2.10e\n", i-1, h, pval);*/
	    }
	  else
	    {
	      /*printf("pf_array[%d] = %2.10e, pval = %2.10e\n", i-1, h, pval);*/
	    }

	}
      h_prev = x_pf_val;
      for( i = x; i < upper_bound; i++ )
	{
	  /*
	    Note:
	    h(i+1) = h(i) * (A-i)(k-i)
	                    --------------
	                    (i+1)(B-k+i+1)
	  */
	  h = h_prev * (double)((A - i)*(k - i))/(double)((i+1)*(B - k + i + 1));
	  h_prev = h;
	  if( h <= x_pf_val )
	    {
	      pval += h;
	      /*printf("*pf_array[%d] = %2.10e, pval = %2.10e\n", i+1, h, pval);*/
	    }
	  else
	    {
	      /*printf("pf_array[%d] = %2.10e, pval = %2.10e\n", i+1, h, pval);*/
	    }
	}
      break;
    }
  return pval;
}

bool hypergeometric_is_pval_lower(int x, int k, int A, int B, 
				  hg_pvaltype pvaltype,
				  double comp_pval)
{
  int i;
  int lower_bound;
  int upper_bound;
  double h_prev;
  double h;
  double pval = 0.0;
  double x_pf_val;

  calculate_bounds(k,A,B,&lower_bound,&upper_bound);
  my_assert(x >= lower_bound);
  my_assert(x <= upper_bound);

  x_pf_val = hypergeometric_pdf(x,k,A,B);
  h_prev = x_pf_val;
  pval = x_pf_val;
  /*printf("comp_pval = %2.40e\n, pf_array[%d] = %2.40e\n", comp_pval, x, x_pf_val);*/
  if( pval > comp_pval )
    {
      return FALSE;
    }

  /* We compute from both ends because it is more stable numerically and
     because we can compute one-sided pvalues faster.
    h(i-1) = i(B - k + i)           * h(i)
             ----------------------
             (A - i + 1)(k - i + 1)
  */
  
  switch( pvaltype )
    {
    case hg_one_sided_decrease:
      for( i = x; i > lower_bound; i-- )
	{
	  h = h_prev * (double)(i * (B - k + i))/(double)((A - i + 1)*(k - i + 1));
	  h_prev = h;
	  pval += h;
	  if( pval > comp_pval )
	    return FALSE;
	}
      break;
    case hg_one_sided_increase:
      for( i = x; i < upper_bound; i++ )
	{
	  /*
	  Note:
          h(i+1) = h(i) * (A-i)(k-i)
                          --------------
                          (i+1)(B-k+i+1)
	  */
	  h = h_prev * (double)((A - i)*(k - i))/(double)((i+1)*(B - k + i + 1));
	  h_prev = h;
	  pval += h;
	  if( pval > comp_pval )
	    return FALSE;
	}
      break;
    case hg_two_sided:
      for( i = x; i > lower_bound; i-- )
	{
	  h = h_prev * (double)(i * (B - k + i))/(double)((A - i + 1)*(k - i + 1));
	  h_prev = h;
	  if( h <= x_pf_val )
	    {
	      pval += h;
	      /*printf("pf_array[%d] = %2.40e, pval = %2.40e\n", i-1, h, pval);*/
	    }
	  if( pval > comp_pval )
	    return FALSE;
	}
      h_prev = x_pf_val;
      for( i = x; i < upper_bound; i++ )
	{
	  /*
	  Note:
          h(i+1) = h(i) * (A-i)(k-i)
                          --------------
                          (i+1)(B-k+i+1)
	  */
	  h = h_prev * (double)((A - i)*(k - i))/(double)((i+1)*(B - k + i + 1));
	  h_prev = h;
	  if( h <= x_pf_val )
	    {
	      pval += h;
	      /*printf("pf_array[%d] = %2.40e, pval = %2.40e\n", i+1, h, pval);*/
	    }
	  if( pval > comp_pval )
	    return FALSE;
	}
      break;
    }
  return ( pval <= comp_pval);
}

double fisher_exact_test_pvalue(int num_match_today, int num_today,
				int num_match_other, int num_other,
				hg_pvaltype pvaltype)
{
  double p_value;
  double p_match_other;
  int num_nomatch_today;
  int num_nomatch_other;
  dyv* obs_dist;
  dyv* hyp_dist;
  /*
  printf("num_match_today = %d, num_today = %d, (%f)\n",num_match_today,
	 num_today, (double)num_match_today/(double)num_today);
  printf("num_match_other = %d, num_other = %d, (%f)\n",num_match_other,
	 num_other, (double)num_match_other/(double)num_other);
  */
  num_nomatch_today = num_today - num_match_today;
  num_nomatch_other = num_other - num_match_other;
  if( num_other > 0 )
    {
      p_match_other = num_match_other / (double) num_other;      
    }
  else
    {
      /* Just in case num_other is 0, say it is insignificant */
      /* printf("Not enough other cases\n");*/
      return 1.0;
    }

  obs_dist = mk_dyv_2((double)num_match_today,(double) num_nomatch_today);
  hyp_dist = mk_dyv_2(p_match_other * num_today,
		      (1-p_match_other) * num_today);
  if(( dyv_min(hyp_dist) > 5.0 ) && ( pvaltype == hg_two_sided ))
    {
      p_value = chi_squared_prob(obs_dist,hyp_dist, /*dof = */ 1);
    }
  else
    {
      p_value = hypergeometric_pval(num_match_today,num_today,
				    num_match_today+num_match_other,
				    num_nomatch_today+num_nomatch_other,
				    pvaltype);
    }
  free_dyv(obs_dist);
  free_dyv(hyp_dist);  
  return p_value;
}


bool fisher_exact_test_pvalue_is_lower(int num_match_today, int num_today,
				       int num_match_other, int num_other,
				       hg_pvaltype pvaltype,
				       double comp_pval)
{
  double p_value;
  double p_match_other;
  int num_nomatch_today;
  int num_nomatch_other;
  dyv* obs_dist;
  dyv* hyp_dist;
  bool result;

  num_nomatch_today = num_today - num_match_today;
  num_nomatch_other = num_other - num_match_other;
  if( num_other > 0 )
    {
      p_match_other = num_match_other / (double) num_other;      
    }
  else
    {
      /* Just in case num_other is 0, say it is insignificant */
      /*printf("Not enough other cases\n");*/
      return ( 1.0 <= comp_pval );
    }

  obs_dist = mk_dyv_2((double)num_match_today,(double) num_nomatch_today);
  hyp_dist = mk_dyv_2(p_match_other * num_today,
		      (1-p_match_other) * num_today);
  if(( dyv_min(hyp_dist) > 5.0 ) && ( pvaltype == hg_two_sided ))
    {
      p_value = chi_squared_prob(obs_dist,hyp_dist, /*dof = */ 1);
      result = ( p_value <= comp_pval );
    }
  else
    {
      result = hypergeometric_is_pval_lower(num_match_today,num_today,
					     num_match_today+num_match_other,
					     num_nomatch_today+num_nomatch_other,pvaltype,comp_pval);
    }
  free_dyv(obs_dist);
  free_dyv(hyp_dist);  
  return result;
}









