/*
    * Licensed to the Apache Software Foundation (ASF) under one or more
    * contributor license agreements.  See the NOTICE file distributed with
    * this work for additional information regarding copyright ownership.
    * The ASF licenses this file to You under the Apache License, Version 2.0
    * (the "License"); you may not use this file except in compliance with
    * the License.  You may obtain a copy of the License at
    *
    *      http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.commons.statistics.distribution;
  
  import org.apache.commons.numbers.gamma.RegularizedBeta;
  
  /**
   * Implementation of the binomial distribution.
   *
   * <p>The probability mass function of \( X \) is:
   *
   * <p>\[ f(k; n, p) = \binom{n}{k} p^k (1-p)^{n-k} \]
   *
   * <p>for \( n \in \{0, 1, 2, \dots\} \) the number of trials,
   * \( p \in [0, 1] \) the probability of success,
   * \( k \in \{0, 1, \dots, n\} \) the number of successes, and
   *
   * <p>\[ \binom{n}{k} = \frac{n!}{k! \, (n-k)!} \]
   *
   * <p>is the binomial coefficient.
   *
   * @see <a href="https://en.wikipedia.org/wiki/Binomial_distribution">Binomial distribution (Wikipedia)</a>
   * @see <a href="https://mathworld.wolfram.com/BinomialDistribution.html">Binomial distribution (MathWorld)</a>
   */
  public final class BinomialDistribution extends AbstractDiscreteDistribution {
      /** 1/2. */
      private static final float HALF = 0.5f;
  
      /** The number of trials. */
      private final int numberOfTrials;
      /** The probability of success. */
      private final double probabilityOfSuccess;
      /** Cached value for pmf(x=0). */
      private final double pmf0;
      /** Cached value for pmf(x=n). */
      private final double pmfn;
  
      /**
       * @param trials Number of trials.
       * @param p Probability of success.
       */
      private BinomialDistribution(int trials,
                                   double p) {
          probabilityOfSuccess = p;
          numberOfTrials = trials;
          // Special pmf cases where the power function is more accurate:
          //   (n choose k) == 1 for k=0, k=n
          //   pmf = p^k (1-p)^(n-k)
          // Note: This handles the edge case of n=0: pmf(k=0) = 1, else 0
          if (probabilityOfSuccess >= HALF) {
              pmf0 = Math.pow(1 - probabilityOfSuccess, numberOfTrials);
          } else {
              pmf0 = Math.exp(numberOfTrials * Math.log1p(-probabilityOfSuccess));
          }
          pmfn = Math.pow(probabilityOfSuccess, numberOfTrials);
      }
  
      /**
       * Creates a binomial distribution.
       *
       * @param trials Number of trials.
       * @param p Probability of success.
       * @return the distribution
       * @throws IllegalArgumentException if {@code trials < 0}, or if {@code p < 0}
       * or {@code p > 1}.
       */
      public static BinomialDistribution of(int trials,
                                            double p) {
          if (trials < 0) {
              throw new DistributionException(DistributionException.NEGATIVE,
                                              trials);
          }
          ArgumentUtils.checkProbability(p);
          // Avoid p = -0.0 to avoid returning -0.0 for some probability computations.
          return new BinomialDistribution(trials, Math.abs(p));
      }
  
      /**
       * Gets the number of trials parameter of this distribution.
       *
       * @return the number of trials.
       */
      public int getNumberOfTrials() {
          return numberOfTrials;
      }
 
     /**
      * Gets the probability of success parameter of this distribution.
      *
      * @return the probability of success.
      */
     public double getProbabilityOfSuccess() {
         return probabilityOfSuccess;
     }
 
     /** {@inheritDoc} */
     @Override
     public double probability(int x) {
         if (x < 0 || x > numberOfTrials) {
             return 0;
         } else if (x == 0) {
             return pmf0;
         } else if (x == numberOfTrials) {
             return pmfn;
         }
         return Math.exp(SaddlePointExpansionUtils.logBinomialProbability(x,
                         numberOfTrials, probabilityOfSuccess,
                         1.0 - probabilityOfSuccess));
     }
 
     /** {@inheritDoc} **/
     @Override
     public double logProbability(int x) {
         if (numberOfTrials == 0) {
             return (x == 0) ? 0.0 : Double.NEGATIVE_INFINITY;
         } else if (x < 0 || x > numberOfTrials) {
             return Double.NEGATIVE_INFINITY;
         }
         // Special cases for x=0, x=n
         // are handled in the saddle point expansion
         return SaddlePointExpansionUtils.logBinomialProbability(x,
                 numberOfTrials, probabilityOfSuccess,
                 1.0 - probabilityOfSuccess);
     }
 
     /** {@inheritDoc} */
     @Override
     public double cumulativeProbability(int x) {
         if (x < 0) {
             return 0.0;
         } else if (x >= numberOfTrials) {
             return 1.0;
         } else if (x == 0) {
             return pmf0;
         }
         return RegularizedBeta.complement(probabilityOfSuccess,
                                           x + 1.0, (double) numberOfTrials - x);
     }
 
     /** {@inheritDoc} */
     @Override
     public double survivalProbability(int x) {
         if (x < 0) {
             return 1.0;
         } else if (x >= numberOfTrials) {
             return 0.0;
         } else if (x == numberOfTrials - 1) {
             return pmfn;
         }
         return RegularizedBeta.value(probabilityOfSuccess,
                                      x + 1.0, (double) numberOfTrials - x);
     }
 
     /**
      * {@inheritDoc}
      *
      * <p>For number of trials \( n \) and probability of success \( p \), the mean is \( np \).
      */
     @Override
     public double getMean() {
         return numberOfTrials * probabilityOfSuccess;
     }
 
     /**
      * {@inheritDoc}
      *
      * <p>For number of trials \( n \) and probability of success \( p \), the variance is \( np (1 - p) \).
      */
     @Override
     public double getVariance() {
         final double p = probabilityOfSuccess;
         return numberOfTrials * p * (1 - p);
     }
 
     /**
      * {@inheritDoc}
      *
      * <p>The lower bound of the support is always 0 except for the probability
      * parameter {@code p = 1}.
      *
      * @return 0 or the number of trials.
      */
     @Override
     public int getSupportLowerBound() {
         return probabilityOfSuccess < 1.0 ? 0 : numberOfTrials;
     }
 
     /**
      * {@inheritDoc}
      *
      * <p>The upper bound of the support is the number of trials except for the
      * probability parameter {@code p = 0}.
      *
      * @return number of trials or 0.
      */
     @Override
     public int getSupportUpperBound() {
         return probabilityOfSuccess > 0.0 ? numberOfTrials : 0;
     }
 
     /** {@inheritDoc} */
     @Override
     int getMedian() {
         // Overridden for the probability(int, int) method.
         // This is intentionally not a public method.
         // Can be floor or ceiling of np. For the probability in a range use the floor
         // as this only used for values >= median+1.
         return (int) (numberOfTrials * probabilityOfSuccess);
     }
 }