/*
 * 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.rng.sampling.distribution;

import org.apache.commons.rng.UniformRandomProvider;

/**
 * Sampling from a <a href="https://en.wikipedia.org/wiki/Geometric_distribution">geometric
 * distribution</a>.
 *
 * <p>This distribution samples the number of failures before the first success taking values in the
 * set {@code [0, 1, 2, ...]}.</p>
 *
 * <p>The sample is computed using a related exponential distribution. If \( X \) is an
 * exponentially distributed random variable with parameter \( \lambda \), then
 * \( Y = \left \lfloor X \right \rfloor \) is a geometrically distributed random variable with
 * parameter \( p = 1 − e^\lambda \), with \( p \) the probability of success.</p>
 *
 * <p>This sampler outperforms using the {@link InverseTransformDiscreteSampler} with an appropriate
 * geometric inverse cumulative probability function.</p>
 *
 * <p>Usage note: As the probability of success (\( p \)) tends towards zero the mean of the
 * distribution (\( \frac{1-p}{p} \)) tends towards infinity and due to the use of {@code int}
 * for the sample this can result in truncation of the distribution.</p>
 *
 * <p>Sampling uses {@link UniformRandomProvider#nextDouble()}.</p>
 *
 * @see <a
 * href="https://en.wikipedia.org/wiki/Geometric_distribution#Related_distributions">Geometric
 * distribution - related distributions</a>
 * @since 1.3
 */
public final class GeometricSampler {
    /**
     * Sample from the geometric distribution when the probability of success is 1.
     */
    private static class GeometricP1Sampler
        implements SharedStateDiscreteSampler {
        /** The single instance. */
        static final GeometricP1Sampler INSTANCE = new GeometricP1Sampler();

        @Override
        public int sample() {
            // When probability of success is 1 the sample is always zero
            return 0;
        }

        @Override
        public String toString() {
            return "Geometric(p=1) deviate";
        }

        @Override
        public SharedStateDiscreteSampler withUniformRandomProvider(UniformRandomProvider rng) {
            // No requirement for a new instance
            return this;
        }
    }

    /**
     * Sample from the geometric distribution by using a related exponential distribution.
     */
    private static class GeometricExponentialSampler
        implements SharedStateDiscreteSampler {
        /** Underlying source of randomness. Used only for the {@link #toString()} method. */
        private final UniformRandomProvider rng;
        /** The related exponential sampler for the geometric distribution. */
        private final SharedStateContinuousSampler exponentialSampler;

        /**
         * @param rng Generator of uniformly distributed random numbers
         * @param probabilityOfSuccess The probability of success (must be in the range
         * {@code [0 < probabilityOfSuccess < 1]})
         */
        GeometricExponentialSampler(UniformRandomProvider rng, double probabilityOfSuccess) {
            this.rng = rng;
            // Use a related exponential distribution:
            // λ = −ln(1 − probabilityOfSuccess)
            // exponential mean = 1 / λ
            // --
            // Note on validation:
            // If probabilityOfSuccess == Math.nextDown(1.0) the exponential mean is >0 (valid).
            // If probabilityOfSuccess == Double.MIN_VALUE the exponential mean is +Infinity
            // and the sample will always be Integer.MAX_VALUE (the distribution is truncated). It
            // is noted in the class Javadoc that the use of a small p leads to truncation so
            // no checks are made for this case.
            final double exponentialMean = 1.0 / (-Math.log1p(-probabilityOfSuccess));
            exponentialSampler = ZigguratSampler.Exponential.of(rng, exponentialMean);
        }

        /**
         * @param rng Generator of uniformly distributed random numbers
         * @param source Source to copy.
         */
        GeometricExponentialSampler(UniformRandomProvider rng, GeometricExponentialSampler source) {
            this.rng = rng;
            exponentialSampler = source.exponentialSampler.withUniformRandomProvider(rng);
        }

        @Override
        public int sample() {
            // Return the floor of the exponential sample
            return (int) Math.floor(exponentialSampler.sample());
        }

        @Override
        public String toString() {
            return "Geometric deviate [" + rng.toString() + "]";
        }

        @Override
        public SharedStateDiscreteSampler withUniformRandomProvider(UniformRandomProvider rng) {
            return new GeometricExponentialSampler(rng, this);
        }
    }

    /** Class contains only static methods. */
    private GeometricSampler() {}

    /**
     * Creates a new geometric distribution sampler. The samples will be provided in the set
     * {@code k=[0, 1, 2, ...]} where {@code k} indicates the number of failures before the first
     * success.
     *
     * @param rng Generator of uniformly distributed random numbers.
     * @param probabilityOfSuccess The probability of success.
     * @return the sampler
     * @throws IllegalArgumentException if {@code probabilityOfSuccess} is not in the range
     * {@code [0 < probabilityOfSuccess <= 1]})
     */
    public static SharedStateDiscreteSampler of(UniformRandomProvider rng,
                                                double probabilityOfSuccess) {
        if (probabilityOfSuccess <= 0 || probabilityOfSuccess > 1) {
            throw new IllegalArgumentException(
                "Probability of success (p) must be in the range [0 < p <= 1]: " +
                    probabilityOfSuccess);
        }
        return probabilityOfSuccess == 1 ?
            GeometricP1Sampler.INSTANCE :
            new GeometricExponentialSampler(rng, probabilityOfSuccess);
    }
}