/*
    * 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.rng.UniformRandomProvider;
  import org.apache.commons.rng.sampling.distribution.DiscreteUniformSampler;
  
  /**
   * Implementation of the uniform discrete distribution.
   *
   * <p>The probability mass function of \( X \) is:
   *
   * <p>\[ f(k; a, b) = \frac{1}{b-a+1} \]
   *
   * <p>for integer \( a, b \) and \( a \le b \) and
   * \( k \in [a, b] \).
   *
   * @see <a href="https://en.wikipedia.org/wiki/Uniform_distribution_(discrete)">
   * Uniform distribution (discrete) (Wikipedia)</a>
   * @see <a href="https://mathworld.wolfram.com/DiscreteUniformDistribution.html">
   * Discrete uniform distribution (MathWorld)</a>
   */
  public final class UniformDiscreteDistribution extends AbstractDiscreteDistribution {
      /** Lower bound (inclusive) of this distribution. */
      private final int lower;
      /** Upper bound (inclusive) of this distribution. */
      private final int upper;
      /** "upper" - "lower" + 1 (as a double to avoid overflow). */
      private final double upperMinusLowerPlus1;
      /** Cache of the probability. */
      private final double pmf;
      /** Cache of the log probability. */
      private final double logPmf;
      /** Value of survival probability for x=0. Used in the inverse survival function. */
      private final double sf0;
  
      /**
       * @param lower Lower bound (inclusive) of this distribution.
       * @param upper Upper bound (inclusive) of this distribution.
       */
      private UniformDiscreteDistribution(int lower,
                                          int upper) {
          this.lower = lower;
          this.upper = upper;
          upperMinusLowerPlus1 = (double) upper - lower + 1;
          pmf = 1.0 / upperMinusLowerPlus1;
          logPmf = -Math.log(upperMinusLowerPlus1);
          sf0 = (upperMinusLowerPlus1 - 1) / upperMinusLowerPlus1;
      }
  
      /**
       * Creates a new uniform discrete distribution.
       *
       * @param lower Lower bound (inclusive) of this distribution.
       * @param upper Upper bound (inclusive) of this distribution.
       * @return the distribution
       * @throws IllegalArgumentException if {@code lower > upper}.
       */
      public static UniformDiscreteDistribution of(int lower,
                                                   int upper) {
          if (lower > upper) {
              throw new DistributionException(DistributionException.INVALID_RANGE_LOW_GT_HIGH,
                                              lower, upper);
          }
          return new UniformDiscreteDistribution(lower, upper);
      }
  
      /** {@inheritDoc} */
      @Override
      public double probability(int x) {
          if (x < lower || x > upper) {
              return 0;
          }
          return pmf;
      }
  
      /** {@inheritDoc} */
      @Override
      public double probability(int x0,
                                int x1) {
          if (x0 > x1) {
              throw new DistributionException(DistributionException.INVALID_RANGE_LOW_GT_HIGH, x0, x1);
          }
          if (x0 >= upper || x1 < lower) {
             // (x0, x1] does not overlap [lower, upper]
             return 0;
         }
 
         // x0 < upper
         // x1 >= lower
 
         // Find the range between x0 (exclusive) and x1 (inclusive) within [lower, upper].
         // In the case of x0 < lower set l so that u - l == (u - lower) + 1
         // long arithmetic prevents overflow
         final long l = Math.max(lower - 1L, x0);
         final long u = Math.min(upper, x1);
 
         return (u - l) / upperMinusLowerPlus1;
     }
 
     /** {@inheritDoc} */
     @Override
     public double logProbability(int x) {
         if (x < lower || x > upper) {
             return Double.NEGATIVE_INFINITY;
         }
         return logPmf;
     }
 
     /** {@inheritDoc} */
     @Override
     public double cumulativeProbability(int x) {
         if (x <= lower) {
             // Note: CDF(x=0) = PDF(x=0)
             return x == lower ? pmf : 0;
         }
         if (x >= upper) {
             return 1;
         }
         return ((double) x - lower + 1) / upperMinusLowerPlus1;
     }
 
     /** {@inheritDoc} */
     @Override
     public double survivalProbability(int x) {
         if (x <= lower) {
             // Note: SF(x=0) = 1 - PDF(x=0)
             // Use a pre-computed value to avoid cancellation when probabilityOfSuccess -> 0
             return x == lower ? sf0 : 1;
         }
         if (x >= upper) {
             return 0;
         }
         return ((double) upper - x) / upperMinusLowerPlus1;
     }
 
     /** {@inheritDoc} */
     @Override
     public int inverseCumulativeProbability(double p) {
         ArgumentUtils.checkProbability(p);
         if (p > sf0) {
             return upper;
         }
         if (p <= pmf) {
             return lower;
         }
         // p in ( pmf         , sf0             ]
         // p in ( 1 / {u-l+1} , {u-l} / {u-l+1} ]
         // x in ( l           , u-1             ]
         int x = (int) (lower + Math.ceil(p * upperMinusLowerPlus1) - 1);
 
         // Correct rounding errors.
         // This ensures x == icdf(cdf(x))
         // Note: Directly computing the CDF(x-1) avoids integer overflow if x=min_value
 
         if (((double) x - lower) / upperMinusLowerPlus1 >= p) {
             // No check for x > lower: cdf(x=lower) = 0 and thus is below p
             // cdf(x-1) >= p
             x--;
         } else if (((double) x - lower + 1) / upperMinusLowerPlus1 < p) {
             // No check for x < upper: cdf(x=upper) = 1 and thus is above p
             // cdf(x) < p
             x++;
         }
 
         return x;
     }
 
     /** {@inheritDoc} */
     @Override
     public int inverseSurvivalProbability(final double p) {
         ArgumentUtils.checkProbability(p);
         if (p < pmf) {
             return upper;
         }
         if (p >= sf0) {
             return lower;
         }
         // p in [ pmf         , sf0             )
         // p in [ 1 / {u-l+1} , {u-l} / {u-l+1} )
         // x in [ u-1         , l               )
         int x = (int) (upper - Math.floor(p * upperMinusLowerPlus1));
 
         // Correct rounding errors.
         // This ensures x == isf(sf(x))
         // Note: Directly computing the SF(x-1) avoids integer overflow if x=min_value
 
         if (((double) upper - x + 1) / upperMinusLowerPlus1 <= p) {
             // No check for x > lower: sf(x=lower) = 1 and thus is above p
             // sf(x-1) <= p
             x--;
         } else if (((double) upper - x) / upperMinusLowerPlus1 > p) {
             // No check for x < upper: sf(x=upper) = 0 and thus is below p
             // sf(x) > p
             x++;
         }
 
         return x;
     }
 
     /**
      * {@inheritDoc}
      *
      * <p>For lower bound \( a \) and upper bound \( b \), the mean is \( \frac{1}{2} (a + b) \).
      */
     @Override
     public double getMean() {
         // Avoid overflow
         return 0.5 * ((double) upper + lower);
     }
 
     /**
      * {@inheritDoc}
      *
      * <p>For lower bound \( a \) and upper bound \( b \), the variance is:
      *
      * <p>\[ \frac{1}{12} (n^2 - 1) \]
      *
      * <p>where \( n = b - a + 1 \).
      */
     @Override
     public double getVariance() {
         return (upperMinusLowerPlus1 * upperMinusLowerPlus1 - 1) / 12;
     }
 
     /**
      * {@inheritDoc}
      *
      * <p>The lower bound of the support is equal to the lower bound parameter
      * of the distribution.
      */
     @Override
     public int getSupportLowerBound() {
         return lower;
     }
 
     /**
      * {@inheritDoc}
      *
      * <p>The upper bound of the support is equal to the upper bound parameter
      * of the distribution.
      */
     @Override
     public int getSupportUpperBound() {
         return upper;
     }
 
     /** {@inheritDoc} */
     @Override
     public DiscreteDistribution.Sampler createSampler(final UniformRandomProvider rng) {
         // Discrete uniform distribution sampler.
         return DiscreteUniformSampler.of(rng, lower, upper)::sample;
     }
 }