/*
   File:        choose.h
   Authors:     Moore / Schneider / Kubica / Liu / Wong
   Created:     Sat Sep 22 09:17:30 EDT 2001
   Description: Header for a cached n-choose k implementation

   Copyright 2001, The Auton Lab
   */

#include "choose.h"
#include "amma.h"

/*****************************************************************************
 Prime factorization algorithms		
 *****************************************************************************/
/* Eratosthenes sieve algorithm */
ivec* mk_primes_less_than_n_ivec(int n)
{
  int i;
  int multiple;

  /* Ignore elements in position 0 and 1 of this ivec */
  ivec* is_prime = mk_constant_ivec(n+1,1);
  ivec* result = mk_ivec(0);

  for( i = 2; (i*i) <= n; i++ )
    {
      if( ivec_ref(is_prime,i) )
	{
	  multiple = i+i;
	  while (multiple <= n)
	    {
	      ivec_set(is_prime,multiple,0);
	      multiple += i;
	    }
	}
    }
 for( i = 2; i <= n; i++ )
   {
     if( ivec_ref(is_prime,i) )
       {
	 add_to_ivec(result,i);
       }
   }
 free_ivec(is_prime);
 return result;
}

/* P must be a prime number.  This function returns the exponent of p that
   divides n! */
int find_prime_exponent_that_divides_fac(int n, int p)
{
  int result = 0;
  int div = 0;
  double pexp;
  
  pexp = (double)p;
  for( pexp = (double)p, div = 1; div > 0; pexp *= p, result += div )
    {
      div = (int)floor((double)n/pexp);
    }
  return result;
}

void mk_prime_factorization_for_fac(int n, ivec** r_factors, 
					   ivec** r_factor_exponents)
{
  int i;
  int exp;

  *r_factors = mk_primes_less_than_n_ivec(n);
  *r_factor_exponents = mk_ivec(0);
  for( i = 0; i < ivec_size(*r_factors); i++ )
    {
      exp = find_prime_exponent_that_divides_fac(n,ivec_ref(*r_factors,i));
      add_to_ivec(*r_factor_exponents,exp);
    }
}

/*
  A factor ivec is an ivec containing the exponents for the prime
  factorizations of a factorial.  For instance, if the number was 20!,
  the prime factorization would be:
    2^18 * 3^8 * 5^4 * 7^2 * 11 * 13 * 17 * 19
  And the factor ivec would look like:
    18, 8, 4, 2, 1, 1, 1, 1
  Note that the number in the ith position of the factor ivec is the ith
  prime in the sequence.  A prime factorization of a factorial does not
  skip any primes.  Thus, the order of the primes in a factor ivec is:
  Position     Prime
  0            2
  1            3
  2            5
  3            7
  4            11
  .            .
  .            .
  .            .
  The prime_factor_ivec_plus adds two prime factor ivecs together
  and places the result in r_result.  pf1 and pf2 do not have to be of 
  the same size.
*/

void prime_factor_ivec_grow(ivec* pf, int grow_by)
{
  int i;
  for( i = 0; i < grow_by; i++ )
    {
      add_to_ivec(pf,0);
    }
}

void prime_factor_ivec_make_same_shape(ivec* pf1, ivec* pf2)
{
  if( ivec_size(pf1) > ivec_size(pf2) )
    {
      prime_factor_ivec_grow(pf2,ivec_size(pf1) - ivec_size(pf2));
    }
  else if( ivec_size(pf2) > ivec_size(pf1) )
    {
      prime_factor_ivec_grow(pf1,ivec_size(pf2) - ivec_size(pf1));
    }
}

void prime_factor_ivec_result_check(ivec* pf1, ivec* pf2, ivec* r_result)
{
  int bigger;
  
  bigger = (ivec_size(pf1) > ivec_size(pf2) ? ivec_size(pf1) : ivec_size(pf2));
  if( ivec_size( r_result ) < bigger )
    {
      prime_factor_ivec_grow(r_result,bigger-ivec_size(r_result));
    }
  else if( ivec_size( r_result ) > bigger )
    {
      my_error("result ivec for prime_factor_ivec_plus must be the same size as the bigger of pf1 and pf2");
    }
}

void prime_factor_ivec_plus(ivec* pf1, ivec* pf2, ivec* r_result)
{
  prime_factor_ivec_result_check(pf1,pf2,r_result);
  prime_factor_ivec_make_same_shape(pf1,pf2);
  ivec_plus(pf1,pf2,r_result);
}

void prime_factor_ivec_subtract(ivec* pf1, ivec* pf2, ivec* r_result)
{
  prime_factor_ivec_result_check(pf1,pf2,r_result);
  prime_factor_ivec_make_same_shape(pf1,pf2);
  ivec_subtract(pf1,pf2,r_result);
}

/* Calculates (n choose k) and expresses the result as a prime factorization
   ie. returns the result in two ivec:
   1. r_factors contains the prime factors
   2. r_factor_exponents contains the exponent of the prime factors ie. 
      the ith prime factor in r_factors has the exponent in the ith position
      of r_factor_exponents
*/
void mk_prime_factor_choose(int n, int k, ivec** r_factors, 
			    ivec** r_factor_exponents)
{
  
  ivec* n_factors;
  ivec* n_exps;
  ivec* k_factors;
  ivec* k_exps;
  ivec* n_k_factors;
  ivec* n_k_exps;

  my_assert( n >= 0 );
  if( n == 0 )
    {
      if( k == 0 )
	{
	  *r_factors = mk_ivec(0);
	  *r_factor_exponents = mk_ivec(0);
	  return;
	}
      else
	{
	  my_error("n is 0 and k > 0 for mk_prime_factor_choose");
	}
    } 
  my_assert( k >= 0 );

  if(( k == 0 ) || (k == n ))
    {
      *r_factors = mk_ivec(0);
      *r_factor_exponents = mk_ivec(0);
      return;
    }
  mk_prime_factorization_for_fac(n,&n_factors,&n_exps);
  mk_prime_factorization_for_fac(k,&k_factors,&k_exps);
  mk_prime_factorization_for_fac(n-k,&n_k_factors,&n_k_exps);

  prime_factor_ivec_plus(k_exps,n_k_exps,k_exps);
  prime_factor_ivec_subtract(n_exps,k_exps,n_exps);
  *r_factors = mk_copy_ivec(n_factors);
  *r_factor_exponents = mk_copy_ivec(n_exps);

  free_ivec(n_factors);
  free_ivec(n_exps);
  free_ivec(k_factors);
  free_ivec(n_k_factors);
  free_ivec(k_exps);
  free_ivec(n_k_exps);
}

/******************** End prime factorization ********************************/

choose_cache* mk_choose_cache(int max_n, int max_k)
{
  int i, j;
  choose_cache *choose;

  /* check that the parameters are valid */
  my_assert(max_n > 0);
  my_assert(max_k > 0);
  my_assert(max_n >= max_k);

  /* allocate memory */
  choose = AM_MALLOC(choose_cache);
  choose->cache = mk_dym(max_n+1,max_k+1);

  /* set internal variables */
  choose->max_k = max_k;  
  choose->max_n = max_n;

  /* initial choose settings to nothing */
  /* (i.e. don't do the work if it is not needed) */
  for(i=0;i<=max_k;i++)
    for(j=0;j<=max_n;j++)
      dym_set(choose->cache, j, i, -1.0);

  return choose;
}


void free_choose_cache(choose_cache* c)
{
  free_dym(c->cache);
  AM_FREE(c,choose_cache);
}

/* The straightforward choose algorithm */
double simple_choose(int n, int k)
{
  int i;
  double numerator = 1.0;
  double denominator = 1.0;
  
  if((n==0)||(k==0))
    {
      return 1.0;
    }

  if ( k > n/2 ) k = n-k;

  for ( i = 0 ; i < k ; i++ )
  {
    numerator *= (double) (n - i);
    denominator *= (double) (i+1);
  }
  return (numerator/denominator);
}

double optimized_choose(int n, int k)
{
  int i;
  ivec* factors;
  ivec* exps;
  double result = 0.0;
  mk_prime_factor_choose(n,k,&factors,&exps);
  for( i = 0; i < ivec_size(factors); i++ )
    {
      result += (ivec_ref(exps,i) * log(ivec_ref(factors,i)));
    }  
  free_ivec(factors);
  free_ivec(exps);
  return exp(result);
}

/* returns a cached value of n choose k if its in the cache.  otherwise it
   computes the value.  if n and k are within the limits of what is stored
   in the cache the result is stored there for future use.  if not, you will
   end up recomputing those values of n choose k every time.
*/
double choose_cache_ref(choose_cache* c, int n, int k)
{
  double result = 0.0;

  /* First check the cache */
  if ((n < c->max_n) && (k < c->max_k) && (dym_ref(c->cache,n,k) >= 0.0))
  {
    return dym_ref(c->cache,n,k);
  }

  /* check that n ad k are greater than zero */
  if((n==0)||(k==0))
  {
    if ((n < c->max_n) && (k < c->max_k)) dym_set(c->cache,n,k,1.0);
    return 1.0;
  }

  result = simple_choose(n,k);

  /* If possible, save the result into the cache and return it. */
  if ((n < c->max_n) && (k < c->max_k))
    dym_set(c->cache,n,k,result);

  return result;
}









