/***************************************************************************
 * Copyright (C) 1995, 1996, Jun Sun, jsun@junsun.net.                     *
 *                                                                         *
 ***************************************************************************
 */
/* implementation of negtive binomial distribution */

#include <math.h>
#include "random_variates.h"
#include "macro.h"

Negbin::Negbin(int s, double p)
{ _s = s;
  _p = p;
  gp = new Geometric(p);
}

double h(int s,int *f,int n,int M,double mean)
{ int k;
  double sum=0.0;
  double p = (double)s / (mean + (double)s);

  for(k=1;k<=M;k++)
    sum += f[k]*log((double)(s+k-1));
  sum += n*s*log(p) + n*mean*log(1-p);
  return sum;
}

Negbin::Negbin(int *data, int num_data)
{ double mean;
  double v;
  double sum,sum2;
  double        max;
  int   *f;
  int   M;
  int  i;

  M = -1;
  sum = 0.0;
  sum2 = 0.0;
  for(i=0;i<num_data;i++)
  { sum += data[i];
    M = MAX(data[i],M);
    sum2 += SQR(data[i]);
  }

  mean = sum / (double) num_data;
  v = sum2/(double)num_data - SQR(mean);

  /* check if these data forms a valid negbin distribution or not */

  if ( v <= mean )
    mythrow("sample data don't form a valid negtive binomial distribution in Negbin::Negbin(int *data, int num_data)");

  /* estimate p and t */
  /* init f */
  f = new int[M+1];
  for(i=0;i<=M;i++) f[i] = 0;
  for(i=0;i<num_data;i++)
    f[data[i]]++;
  for(i=M-1;i>=0;i--)
    f[i] += f[i+1];

  /* find _s */
  max = h(1,f,num_data,M,mean);
  for(_s=2;;_s++)
    if (h(_s,f,num_data,M,mean) > max)
      max = h(_s,f,num_data,M,mean);
    else
      break;
  _s--;
  _p = (double)_s / (mean + (double)_s);
  gp = new Geometric(_p);


  delete[] f;
}

Negbin::~Negbin() { delete gp; }


double Negbin::mass(int x)
{ if (x<0)
    return 0.0;
  else
    return combination(_s+x-1,x)*pow(_p,(double)_s)*pow(1-_p,(double)(x));
}


double Negbin::distribution(double x)
{ double f;

  if (x<0.0) return 0;

  f= 0.0;
  for(int i=0;i<=FLOOR(x);i++)
    f += combination(_s+i-1,i)*pow(_p,(double)_s)*pow(1-_p,(double)i);
  return f;
}


int Negbin::sample()
{  
  int i=0;
  int sum = 0;
  for(; i<_s; i++)
    sum += gp->sample();

  return sum;
}