/*
 * 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;

/**
 * Discrete uniform distribution sampler.
 *
 * <p>Sampling uses {@link UniformRandomProvider#nextInt}.</p>
 *
 * <p>When the range is a power of two the number of calls is 1 per sample.
 * Otherwise a rejection algorithm is used to ensure uniformity. In the worst
 * case scenario where the range spans half the range of an {@code int}
 * (2<sup>31</sup> + 1) the expected number of calls is 2 per sample.</p>
 *
 * <p>This sampler can be used as a replacement for {@link UniformRandomProvider#nextInt}
 * with appropriate adjustment of the upper bound to be inclusive and will outperform that
 * method when the range is not a power of two. The advantage is gained by pre-computation
 * of the rejection threshold.</p>
 *
 * <p>The sampling algorithm is described in:</p>
 *
 * <blockquote>
 *  Lemire, D (2019). <i>Fast Random Integer Generation in an Interval.</i>
 *  <strong>ACM Transactions on Modeling and Computer Simulation</strong> 29 (1).
 * </blockquote>
 *
 * <p>The number of {@code int} values required per sample follows a geometric distribution with
 * a probability of success p of {@code 1 - ((2^32 % n) / 2^32)}. This requires on average 1/p random
 * {@code int} values per sample.</p>
 *
 * @see <a href="https://arxiv.org/abs/1805.10941">Fast Random Integer Generation in an Interval</a>
 *
 * @since 1.0
 */
public class DiscreteUniformSampler
    extends SamplerBase
    implements SharedStateDiscreteSampler {

    /** The appropriate uniform sampler for the parameters. */
    private final SharedStateDiscreteSampler delegate;

    /**
     * Base class for a sampler from a discrete uniform distribution. This contains the
     * source of randomness.
     */
    private abstract static class AbstractDiscreteUniformSampler
            implements SharedStateDiscreteSampler {
        /** Underlying source of randomness. */
        protected final UniformRandomProvider rng;

        /**
         * @param rng Generator of uniformly distributed random numbers.
         */
        AbstractDiscreteUniformSampler(UniformRandomProvider rng) {
            this.rng = rng;
        }

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

    /**
     * Discrete uniform distribution sampler when the sample value is fixed.
     */
    private static class FixedDiscreteUniformSampler
            extends AbstractDiscreteUniformSampler {
        /** The value. */
        private final int value;

        /**
         * @param value The value.
         */
        FixedDiscreteUniformSampler(int value) {
            // No requirement for the RNG
            super(null);
            this.value = value;
        }

        @Override
        public int sample() {
            return value;
        }

        @Override
        public String toString() {
            // No RNG to include in the string
            return "Uniform deviate [X=" + value + "]";
        }

        @Override
        public SharedStateDiscreteSampler withUniformRandomProvider(UniformRandomProvider rng) {
            // No requirement for the RNG
            return this;
        }
    }

    /**
     * Discrete uniform distribution sampler when the range is a power of 2 and greater than 1.
     * This sampler assumes the lower bound of the range is 0.
     *
     * <p>Note: This cannot be used when the range is 1 (2^0) as the shift would be 32-bits
     * which is ignored by the shift operator.</p>
     */
    private static class PowerOf2RangeDiscreteUniformSampler
            extends AbstractDiscreteUniformSampler {
        /** Bit shift to apply to the integer sample. */
        private final int shift;

        /**
         * @param rng Generator of uniformly distributed random numbers.
         * @param range Maximum range of the sample (exclusive).
         * Must be a power of 2 greater than 2^0.
         */
        PowerOf2RangeDiscreteUniformSampler(UniformRandomProvider rng,
                                            int range) {
            super(rng);
            this.shift = Integer.numberOfLeadingZeros(range) + 1;
        }

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

        @Override
        public int sample() {
            // Use a bit shift to favour the most significant bits.
            // Note: The result would be the same as the rejection method used in the
            // small range sampler when there is no rejection threshold.
            return rng.nextInt() >>> shift;
        }

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

    /**
     * Discrete uniform distribution sampler when the range is small
     * enough to fit in a positive integer.
     * This sampler assumes the lower bound of the range is 0.
     *
     * <p>Implements the algorithm of Lemire (2019).</p>
     *
     * @see <a href="https://arxiv.org/abs/1805.10941">Fast Random Integer Generation in an Interval</a>
     */
    private static class SmallRangeDiscreteUniformSampler
            extends AbstractDiscreteUniformSampler {
        /** Maximum range of the sample (exclusive). */
        private final long n;

        /**
         * The level below which samples are rejected based on the fraction remainder.
         *
         * <p>Any remainder below this denotes that there are still floor(2^32 / n) more
         * observations of this sample from the interval [0, 2^32), where n is the range.</p>
         */
        private final long threshold;

        /**
         * @param rng Generator of uniformly distributed random numbers.
         * @param range Maximum range of the sample (exclusive).
         */
        SmallRangeDiscreteUniformSampler(UniformRandomProvider rng,
                                         int range) {
            super(rng);
            // Handle range as an unsigned 32-bit integer
            this.n = range & 0xffffffffL;
            // Compute 2^32 % n
            threshold = (1L << 32) % n;
        }

        /**
         * @param rng Generator of uniformly distributed random numbers.
         * @param source Source to copy.
         */
        SmallRangeDiscreteUniformSampler(UniformRandomProvider rng,
                                         SmallRangeDiscreteUniformSampler source) {
            super(rng);
            this.n = source.n;
            this.threshold = source.threshold;
        }

        @Override
        public int sample() {
            // Rejection method using multiply by a fraction:
            // n * [0, 2^32 - 1)
            //     -------------
            //         2^32
            // The result is mapped back to an integer and will be in the range [0, n).
            // Note this is comparable to range * rng.nextDouble() but with compensation for
            // non-uniformity due to round-off.
            long result;
            do {
                // Compute 64-bit unsigned product of n * [0, 2^32 - 1).
                // The upper 32-bits contains the sample value in the range [0, n), i.e. result / 2^32.
                // The lower 32-bits contains the remainder, i.e. result % 2^32.
                result = n * (rng.nextInt() & 0xffffffffL);

                // Test the sample uniformity.
                // Samples are observed on average (2^32 / n) times at a frequency of either
                // floor(2^32 / n) or ceil(2^32 / n).
                // To ensure all samples have a frequency of floor(2^32 / n) reject any results with
                // a remainder < (2^32 % n), i.e. the level below which denotes that there are still
                // floor(2^32 / n) more observations of this sample.
            } while ((result & 0xffffffffL) < threshold);
            // Divide by 2^32 to get the sample.
            return (int)(result >>> 32);
        }

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

    /**
     * Discrete uniform distribution sampler when the range between lower and upper is too large
     * to fit in a positive integer.
     */
    private static class LargeRangeDiscreteUniformSampler
            extends AbstractDiscreteUniformSampler {
        /** Lower bound. */
        private final int lower;
        /** Upper bound. */
        private final int upper;

        /**
         * @param rng Generator of uniformly distributed random numbers.
         * @param lower Lower bound (inclusive) of the distribution.
         * @param upper Upper bound (inclusive) of the distribution.
         */
        LargeRangeDiscreteUniformSampler(UniformRandomProvider rng,
                                         int lower,
                                         int upper) {
            super(rng);
            this.lower = lower;
            this.upper = upper;
        }

        @Override
        public int sample() {
            // Use a simple rejection method.
            // This is used when (upper-lower) >= Integer.MAX_VALUE.
            // This will loop on average 2 times in the worst case scenario
            // when (upper-lower) == Integer.MAX_VALUE.
            while (true) {
                final int r = rng.nextInt();
                if (r >= lower &&
                    r <= upper) {
                    return r;
                }
            }
        }

        @Override
        public SharedStateDiscreteSampler withUniformRandomProvider(UniformRandomProvider rng) {
            return new LargeRangeDiscreteUniformSampler(rng, lower, upper);
        }
    }

    /**
     * Adds an offset to an underlying discrete sampler.
     */
    private static class OffsetDiscreteUniformSampler
            extends AbstractDiscreteUniformSampler {
        /** The offset. */
        private final int offset;
        /** The discrete sampler. */
        private final SharedStateDiscreteSampler sampler;

        /**
         * @param offset The offset for the sample.
         * @param sampler The discrete sampler.
         */
        OffsetDiscreteUniformSampler(int offset,
                                     SharedStateDiscreteSampler sampler) {
            super(null);
            this.offset = offset;
            this.sampler = sampler;
        }

        @Override
        public int sample() {
            return offset + sampler.sample();
        }

        @Override
        public String toString() {
            return sampler.toString();
        }

        @Override
        public SharedStateDiscreteSampler withUniformRandomProvider(UniformRandomProvider rng) {
            return new OffsetDiscreteUniformSampler(offset, sampler.withUniformRandomProvider(rng));
        }
    }

    /**
     * This instance delegates sampling. Use the factory method
     * {@link #of(UniformRandomProvider, int, int)} to create an optimal sampler.
     *
     * @param rng Generator of uniformly distributed random numbers.
     * @param lower Lower bound (inclusive) of the distribution.
     * @param upper Upper bound (inclusive) of the distribution.
     * @throws IllegalArgumentException if {@code lower > upper}.
     */
    public DiscreteUniformSampler(UniformRandomProvider rng,
                                  int lower,
                                  int upper) {
        super(null);
        delegate = of(rng, lower, upper);
    }

    /** {@inheritDoc} */
    @Override
    public int sample() {
        return delegate.sample();
    }

    /** {@inheritDoc} */
    @Override
    public String toString() {
        return delegate.toString();
    }

    /**
     * {@inheritDoc}
     *
     * @since 1.3
     */
    @Override
    public SharedStateDiscreteSampler withUniformRandomProvider(UniformRandomProvider rng) {
        // Direct return of the optimised sampler
        return delegate.withUniformRandomProvider(rng);
    }

    /**
     * Creates a new discrete uniform distribution sampler.
     *
     * @param rng Generator of uniformly distributed random numbers.
     * @param lower Lower bound (inclusive) of the distribution.
     * @param upper Upper bound (inclusive) of the distribution.
     * @return the sampler
     * @throws IllegalArgumentException if {@code lower > upper}.
     * @since 1.3
     */
    public static SharedStateDiscreteSampler of(UniformRandomProvider rng,
                                                int lower,
                                                int upper) {
        if (lower > upper) {
            throw new IllegalArgumentException(lower  + " > " + upper);
        }

        // Choose the algorithm depending on the range

        // Edge case for no range.
        // This must be done first as the methods to handle lower == 0
        // do not handle upper == 0.
        if (upper == lower) {
            return new FixedDiscreteUniformSampler(lower);
        }

        // Algorithms to ignore the lower bound if it is zero.
        if (lower == 0) {
            return createZeroBoundedSampler(rng, upper);
        }

        final int range = (upper - lower) + 1;
        // Check power of 2 first to handle range == 2^31.
        if (isPowerOf2(range)) {
            return new OffsetDiscreteUniformSampler(lower,
                                                    new PowerOf2RangeDiscreteUniformSampler(rng, range));
        }
        if (range <= 0) {
            // The range is too wide to fit in a positive int (larger
            // than 2^31); use a simple rejection method.
            // Note: if range == 0 then the input is [Integer.MIN_VALUE, Integer.MAX_VALUE].
            // No specialisation exists for this and it is handled as a large range.
            return new LargeRangeDiscreteUniformSampler(rng, lower, upper);
        }
        // Use a sample from the range added to the lower bound.
        return new OffsetDiscreteUniformSampler(lower,
                                                new SmallRangeDiscreteUniformSampler(rng, range));
    }

    /**
     * Create a new sampler for the range {@code 0} inclusive to {@code upper} inclusive.
     *
     * <p>This can handle any positive {@code upper}.
     *
     * @param rng Generator of uniformly distributed random numbers.
     * @param upper Upper bound (inclusive) of the distribution. Must be positive.
     * @return the sampler
     */
    private static AbstractDiscreteUniformSampler createZeroBoundedSampler(UniformRandomProvider rng,
                                                                           int upper) {
        // Note: Handle any range up to 2^31 (which is negative as a signed
        // 32-bit integer but handled as a power of 2)
        final int range = upper + 1;
        return isPowerOf2(range) ?
            new PowerOf2RangeDiscreteUniformSampler(rng, range) :
            new SmallRangeDiscreteUniformSampler(rng, range);
    }

    /**
     * Checks if the value is a power of 2.
     *
     * <p>This returns {@code true} for the value {@code Integer.MIN_VALUE} which can be
     * handled as an unsigned integer of 2^31.</p>
     *
     * @param value Value.
     * @return {@code true} if a power of 2
     */
    private static boolean isPowerOf2(final int value) {
        return value != 0 && (value & (value - 1)) == 0;
    }
}