/* *
 * Title:        

 * Description:  

 * Copyright:    Copyright (c) 

 * Company:      

 * @author
 * @version 1.0
 */


package umontreal.iro.lecuyer.probdist;

import umontreal.iro.lecuyer.util.RootFinder;
import umontreal.iro.lecuyer.functions.MathFunction;


/**
 * The Pascal distribution is a special case of the
 * negative binomial distribution with
 * parameters n and p, where n is a positive
 * integer and 
 * 0 <= p <= 1.
 * Its mass function is
 * 

 * 
 * 

 * p(x) = nCr(n + x - 1, x)pn(1 - p)x,        for x = 0, 1, 2,…
 * 

 * where nCr is defined in {@link BinomialDist}.
 * This p(x) can be interpreted as the probability of having x failures
 * before the nth success in a sequence of independent Bernoulli trials
 * with  probability of success p.
 * For n = 1, this gives the geometric distribution.
 * 
 */
public class PascalDist extends NegativeBinomialDist {
   private int n;
   private static final double EPSI = 1.0E-10;

   private static class Function implements MathFunction {
      protected int m;
      protected int max;
      protected double mean;
      protected int[] Fj;

      public Function (int m, int max, double mean, int[] Fj)
      {
         this.m = m;
         this.max = max;
         this.mean = mean;
         this.Fj = new int[Fj.length];
         System.arraycopy(Fj, 0, this.Fj, 0, Fj.length);
      }

      public double evaluate (double p)
      {
         double sum = 0.0;
         double s = (p * mean) / (1.0 - p);

         for (int j = 0; j < max; j++)
            sum += Fj[j] / (s + (double) j);

         return sum + m * Math.log (p);
      }

      public double evaluateN (int n, double p)
      {
         double sum = 0.0;

         for (int j = 0; j < max; j++)
            sum += Fj[j] / (n + j);

         return sum + m * Math.log (p);
      }
   }



   /**
    * Creates an object that contains the probability
    *    terms and the distribution function for
    *    the Pascal distribution with parameter n and p.
    * 
    */
   public PascalDist (int n, double p) {
      setParams (n, p);
   }


   /**
    * Estimates the parameter (n, p) of the Pascal distribution
    *    using the maximum likelihood method, from the m observations 
    *    x[i], 
    * i = 0, 1,…, m - 1. The estimates are returned in a two-element
    *     array, in regular order: [n, p].
    * 
    * @param x the list of observations used to evaluate parameters
    * 
    *    @param m the number of observations used to evaluate parameters
    * 
    *    @return returns the parameters [hat(n), hat(p)]
    * 
    */
   public static double[] getMLE (int[] x, int m) {
      if (m <= 0)
         throw new IllegalArgumentException ("m <= 0");

      double sum = 0.0;
      int max = Integer.MIN_VALUE;
      for (int i = 0; i < m; i++)
      {
         sum += x[i];
         if (x[i] > max)
            max = x[i];
      }
      double mean = (double) sum / (double) m;

      double var = 0.0;
      for (int i = 0; i < m; i++)
         var += (x[i] - mean) * (x[i] - mean);
      var /= (double) m;

      if (mean >= var)
           throw new UnsupportedOperationException("mean >= variance");

      int[] Fj = new int[max];
      for (int j = 0; j < max; j++)
      {
         int prop = 0;
         for (int i = 0; i < m; i++)
            if (x[i] > j)
               prop++;

         Fj[j] = prop;
      }

      double[] parameters = new double[2];
      Function f = new Function (m, max, mean, Fj);

      parameters[1] = RootFinder.brentDekker (EPSI, 1 - EPSI, f, 1e-5);
      if (parameters[1] >= 1.0)
         parameters[1] = 1.0 - 1e-15;

      parameters[0] = Math.round ((parameters[1] * mean) / (1.0 - parameters[1]));
      if(parameters[0] == 0)
          parameters[0] = 1;

      return parameters;
   }


   /**
    * Creates a new instance of a Pascal distribution with parameters n and
    *    p estimated using the maximum likelihood method based on the m
    *    observations x[i], 
    * i = 0, 1,…, m - 1.
    * 
    * @param x the list of observations to use to evaluate parameters
    * 
    *    @param m the number of observations to use to evaluate parameters
    * 
    * 
    */
   public static PascalDist getInstanceFromMLE (int[] x, int m) {
      double parameters[] = getMLE (x, m);
      return new PascalDist ((int) parameters[0], parameters[1]);
   }


   /**
    * Returns the parameter n of this object.
    * 
    */
   public int getN() {
      return n;
   }
   


   /**
    * Sets the parameter n and p of this object.
    * 
    */
   public void setParams (int n, double p) {
      setParams ((double) n, p);
      this.n = n;
   }


   /**
    * Return a table containing the parameters of the current distribution.
    *    This table is put in regular order: [n, p].
    * 
    * 
    */
   public double[] getParams () {
      double[] retour = {n, p};
      return retour;
   }


   public String toString () {
      return super.toString();
   }

}