/*
    * 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 java.util.function.IntToDoubleFunction;
  import org.apache.commons.rng.UniformRandomProvider;
  import org.apache.commons.rng.sampling.distribution.GeometricSampler;
  
  /**
   * Implementation of the geometric distribution.
   *
   * <p>The probability mass function of \( X \) is:
   *
   * <p>\[ f(k; p) = (1-p)^k \, p \]
   *
   * <p>for \( p \in (0, 1] \) the probability of success and
   * \( k \in \{0, 1, 2, \dots\} \) the number of failures.
   *
   * <p>This parameterization is used to model the number of failures until
   * the first success.
   *
   * @see <a href="https://en.wikipedia.org/wiki/Geometric_distribution">Geometric distribution (Wikipedia)</a>
   * @see <a href="https://mathworld.wolfram.com/GeometricDistribution.html">Geometric distribution (MathWorld)</a>
   */
  public final class GeometricDistribution extends AbstractDiscreteDistribution {
      /** 1/2. */
      private static final double HALF = 0.5;
  
      /** The probability of success. */
      private final double probabilityOfSuccess;
      /** {@code log(p)} where p is the probability of success. */
      private final double logProbabilityOfSuccess;
      /** {@code log(1 - p)} where p is the probability of success. */
      private final double log1mProbabilityOfSuccess;
      /** Value of survival probability for x=0.
       * Used in the survival functions. Equal to (1 - probability of success). */
      private final double sf0;
      /** Implementation of PMF(x). Assumes that {@code x > 0}. */
      private final IntToDoubleFunction pmf;
  
      /**
       * @param p Probability of success.
       */
      private GeometricDistribution(double p) {
          probabilityOfSuccess = p;
          logProbabilityOfSuccess = Math.log(p);
          log1mProbabilityOfSuccess = Math.log1p(-p);
          sf0 = 1 - p;
  
          // Choose the PMF implementation.
          // When p >= 0.5 then 1 - p is exact and using the power function
          // is consistently more accurate than the use of the exponential function.
          // When p -> 0 then the exponential function avoids large error propagation
          // of the power function used with an inexact 1 - p.
          // Also compute the survival probability for use when x=0.
          if (p >= HALF) {
              pmf = x -> Math.pow(sf0, x) * probabilityOfSuccess;
          } else {
              pmf = x -> Math.exp(log1mProbabilityOfSuccess * x) * probabilityOfSuccess;
          }
      }
  
      /**
       * Creates a geometric distribution.
       *
       * @param p Probability of success.
       * @return the geometric distribution
       * @throws IllegalArgumentException if {@code p <= 0} or {@code p > 1}.
       */
      public static GeometricDistribution of(double p) {
          if (p <= 0 || p > 1) {
              throw new DistributionException(DistributionException.INVALID_NON_ZERO_PROBABILITY, p);
          }
          return new GeometricDistribution(p);
      }
  
      /**
       * 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) {
             // Special case of x=0 exploiting cancellation.
             return x == 0 ? probabilityOfSuccess : 0;
         }
         return pmf.applyAsDouble(x);
     }
 
     /** {@inheritDoc} */
     @Override
     public double logProbability(int x) {
         if (x <= 0) {
             // Special case of x=0 exploiting cancellation.
             return x == 0 ? logProbabilityOfSuccess : Double.NEGATIVE_INFINITY;
         }
         return x * log1mProbabilityOfSuccess + logProbabilityOfSuccess;
     }
 
     /** {@inheritDoc} */
     @Override
     public double cumulativeProbability(int x) {
         if (x <= 0) {
             // Note: CDF(x=0) = PDF(x=0) = probabilityOfSuccess
             return x == 0 ? probabilityOfSuccess : 0;
         }
         // Note: Double addition avoids overflow. This may compute a value less than 1.0
         // for the max integer value when p is very small.
         return -Math.expm1(log1mProbabilityOfSuccess * (x + 1.0));
     }
 
     /** {@inheritDoc} */
     @Override
     public double survivalProbability(int x) {
         if (x <= 0) {
             // Note: SF(x=0) = 1 - PDF(x=0) = 1 - probabilityOfSuccess
             // Use a pre-computed value to avoid cancellation when probabilityOfSuccess -> 0
             return x == 0 ? sf0 : 1;
         }
         // Note: Double addition avoids overflow. This may compute a value greater than 0.0
         // for the max integer value when p is very small.
         return Math.exp(log1mProbabilityOfSuccess * (x + 1.0));
     }
 
     /** {@inheritDoc} */
     @Override
     public int inverseCumulativeProbability(double p) {
         ArgumentUtils.checkProbability(p);
         if (p == 1) {
             return getSupportUpperBound();
         }
         if (p <= probabilityOfSuccess) {
             return 0;
         }
         // p > probabilityOfSuccess
         // => log(1-p) < log(1-probabilityOfSuccess);
         // Both terms are negative as probabilityOfSuccess > 0.
         // This should be lower bounded to (2 - 1) = 1
         int x = (int) (Math.ceil(Math.log1p(-p) / log1mProbabilityOfSuccess) - 1);
 
         // Correct rounding errors.
         // This ensures x == icdf(cdf(x))
 
         if (cumulativeProbability(x - 1) >= p) {
             // No checks for x=0.
             // If x=0; cdf(-1) = 0 and the condition is false as p>0 at this point.
             x--;
         } else if (cumulativeProbability(x) < p && x < Integer.MAX_VALUE) {
             // The supported upper bound is max_value here as probabilityOfSuccess != 1
             x++;
         }
 
         return x;
     }
 
     /** {@inheritDoc} */
     @Override
     public int inverseSurvivalProbability(double p) {
         ArgumentUtils.checkProbability(p);
         if (p == 0) {
             return getSupportUpperBound();
         }
         if (p >= sf0) {
             return 0;
         }
 
         // p < 1 - probabilityOfSuccess
         // Inversion as for icdf using log(p) in place of log1p(-p)
         int x = (int) (Math.ceil(Math.log(p) / log1mProbabilityOfSuccess) - 1);
 
         // Correct rounding errors.
         // This ensures x == isf(sf(x))
 
         if (survivalProbability(x - 1) <= p) {
             // No checks for x=0
             // If x=0; sf(-1) = 1 and the condition is false as p<1 at this point.
             x--;
         } else if (survivalProbability(x) > p && x < Integer.MAX_VALUE) {
             // The supported upper bound is max_value here as probabilityOfSuccess != 1
             x++;
         }
 
         return x;
     }
 
     /**
      * {@inheritDoc}
      *
      * <p>For probability parameter \( p \), the mean is:
      *
      * <p>\[ \frac{1 - p}{p} \]
      */
     @Override
     public double getMean() {
         return (1 - probabilityOfSuccess) / probabilityOfSuccess;
     }
 
     /**
      * {@inheritDoc}
      *
      * <p>For probability parameter \( p \), the variance is:
      *
      * <p>\[ \frac{1 - p}{p^2} \]
      */
     @Override
     public double getVariance() {
         return (1 - probabilityOfSuccess) / (probabilityOfSuccess * probabilityOfSuccess);
     }
 
     /**
      * {@inheritDoc}
      *
      * <p>The lower bound of the support is always 0.
      *
      * @return 0.
      */
     @Override
     public int getSupportLowerBound() {
         return 0;
     }
 
     /**
      * {@inheritDoc}
      *
      * <p>The upper bound of the support is positive infinity except for the
      * probability parameter {@code p = 1.0}.
      *
      * @return {@link Integer#MAX_VALUE} or 0.
      */
     @Override
     public int getSupportUpperBound() {
         return probabilityOfSuccess < 1 ? Integer.MAX_VALUE : 0;
     }
 
     /** {@inheritDoc} */
     @Override
     public Sampler createSampler(UniformRandomProvider rng) {
         return GeometricSampler.of(rng, probabilityOfSuccess)::sample;
     }
 }