#include <RanGen.h>

using namespace std;

ClassImp(RanGen);

RanGen::RanGen(Int_t seed) : TRandom3(seed){
}

RanGen::~RanGen(){};

Double_t RanGen::Gamma(Int_t n, Double_t theta){

  Double_t value = 0;

  for(Int_t i=0; i<n; i++) value -= TMath::Log( this->Uniform() );

  return value;

}


vector<Double_t> RanGen::Dirichlet(Int_t n, vector<Int_t> x){

  Double_t sum = 0;
  vector<Double_t> gammas(n);
  vector<Double_t> value(n);

  for(Int_t i=0; i<n; i++){
    gammas[i] = Gamma(x[i], 1);
    sum += gammas[i];
  }

  for(Int_t i=0; i<n; i++){
    if(sum)
      value[i] = gammas[i]/sum;
    else
      value[i] = 0;
  }

  return value;

}


Int_t RanGen::Binomial(Int_t n, Double_t p){
  unsigned int i, a, b, k = 0;
  while (n > 10)        /* This parameter is tunable */
    {
      double X;
      a = 1 + (n / 2);
      b = 1 + n - a;
      
      X = this->Beta((Int_t) a, (Int_t) b);
      
      if (X >= p)
	{
	  n = a - 1;
	  p /= X;
	}
      else
	{
	  k += a;
	  n = b - 1;
	  p = (p - X) / (1 - X);
	}
    }
  
  for (i = 0; i < n; i++)
    {
      double u = this->Uniform();
      if (u < p)
	k++;
    }
  
  return k;
  
}

Double_t RanGen::Beta(const Int_t a, const Int_t b){
  double x1 = this->Gamma(a, 1.0);
  double x2 = this->Gamma(b, 1.0);
  
  return x1 / (x1 + x2);
}

vector<Int_t> RanGen::Multinomial(Int_t N, vector<Double_t> p){

/*  From the GSL Library:
    The multinomial distribution has the form
   
                                       N!           n_1  n_2      n_K
    prob(n_1, n_2, ... n_K) = -------------------- p_1  p_2  ... p_K
                              (n_1! n_2! ... n_K!) 
 
    where n_1, n_2, ... n_K are nonnegative integers, sum_{k=1,K} n_k = N,
    and p = (p_1, p_2, ..., p_K) is a probability distribution. 
 
    Random variates are generated using the conditional binomial method.
    This scales well with N and does not require a setup step.
 
    Ref: 
    C.S. David, The computer generation of multinomial random variates,
    Comp. Stat. Data Anal. 16 (1993) 205-217
 */

  size_t K = p.size();
  vector<Int_t> n(K);

  size_t k;
  double norm = 0.0;
  double sum_p = 0.0;
  
  unsigned int sum_n = 0;
  
  /* p[k] may contain non-negative weights that do not sum to 1.0.
   * Even a probability distribution will not exactly sum to 1.0
   * due to rounding errors. 
   */
  
  for (k = 0; k < K; k++)
    {
      norm += p[k];
    }
  
  for (k = 0; k < K; k++)
    {
      if (p[k] > 0.0)
	{
	  n[k] = this->Binomial(N - sum_n, p[k] / (norm - sum_p));
	}
      else
	{
	  n[k] = 0;
	}
      
      sum_p += p[k];
      sum_n += n[k];
    }
  return n;
  
}
