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

import org.apache.commons.math3.distribution.BinomialDistribution;
import org.apache.commons.math3.distribution.IntegerDistribution;
import org.apache.commons.math3.distribution.PoissonDistribution;
import org.apache.commons.rng.UniformRandomProvider;
import org.apache.commons.rng.sampling.distribution.AliasMethodDiscreteSampler;
import org.apache.commons.rng.sampling.distribution.DirichletSampler;
import org.apache.commons.rng.sampling.distribution.DiscreteSampler;
import org.apache.commons.rng.sampling.distribution.FastLoadedDiceRollerDiscreteSampler;
import org.apache.commons.rng.sampling.distribution.GuideTableDiscreteSampler;
import org.apache.commons.rng.sampling.distribution.MarsagliaTsangWangDiscreteSampler;
import org.apache.commons.rng.simple.RandomSource;

import org.openjdk.jmh.annotations.Benchmark;
import org.openjdk.jmh.annotations.BenchmarkMode;
import org.openjdk.jmh.annotations.Fork;
import org.openjdk.jmh.annotations.Level;
import org.openjdk.jmh.annotations.Measurement;
import org.openjdk.jmh.annotations.Mode;
import org.openjdk.jmh.annotations.OutputTimeUnit;
import org.openjdk.jmh.annotations.Param;
import org.openjdk.jmh.annotations.Scope;
import org.openjdk.jmh.annotations.Setup;
import org.openjdk.jmh.annotations.State;
import org.openjdk.jmh.annotations.Warmup;

import java.util.Arrays;
import java.util.concurrent.TimeUnit;
import java.util.function.Supplier;

/**
 * Executes benchmark to compare the speed of generation of random numbers from an enumerated
 * discrete probability distribution.
 */
@BenchmarkMode(Mode.AverageTime)
@OutputTimeUnit(TimeUnit.NANOSECONDS)
@Warmup(iterations = 5, time = 1, timeUnit = TimeUnit.SECONDS)
@Measurement(iterations = 5, time = 1, timeUnit = TimeUnit.SECONDS)
@State(Scope.Benchmark)
@Fork(value = 1, jvmArgs = {"-server", "-Xms128M", "-Xmx128M"})
public class EnumeratedDistributionSamplersPerformance {
    /**
     * The value for the baseline generation of an {@code int} value.
     *
     * <p>This must NOT be final!</p>
     */
    private int value;

    /**
     * The random sources to use for testing. This is a smaller list than all the possible
     * random sources; the list is composed of generators of different speeds.
     */
    @State(Scope.Benchmark)
    public static class LocalRandomSources {
        /**
         * RNG providers.
         *
         * <p>Use different speeds.</p>
         *
         * @see <a href="https://commons.apache.org/proper/commons-rng/userguide/rng.html">
         *      Commons RNG user guide</a>
         */
        @Param({"WELL_44497_B",
                "ISAAC",
                "XO_RO_SHI_RO_128_PLUS"})
        private String randomSourceName;

        /** RNG. */
        private UniformRandomProvider generator;

        /**
         * @return the RNG.
         */
        public UniformRandomProvider getGenerator() {
            return generator;
        }

        /** Create the random source. */
        @Setup
        public void setup() {
            final RandomSource randomSource = RandomSource.valueOf(randomSourceName);
            generator = randomSource.create();
        }
    }

    /**
     * The {@link DiscreteSampler} samplers to use for testing. Creates the sampler for each
     * random source.
     *
     * <p>This class is abstract. The probability distribution is created by implementations.</p>
     */
    @State(Scope.Benchmark)
    public abstract static class SamplerSources extends LocalRandomSources {
        /**
         * The sampler type.
         */
        @Param({"BinarySearchDiscreteSampler",
                "AliasMethodDiscreteSampler",
                "GuideTableDiscreteSampler",
                "MarsagliaTsangWangDiscreteSampler",
                "FastLoadedDiceRollerDiscreteSampler",
                "FastLoadedDiceRollerDiscreteSamplerLong",
                "FastLoadedDiceRollerDiscreteSampler53",

                // Uncomment to test non-default parameters
                //"AliasMethodDiscreteSamplerNoPad", // Not optimal for sampling
                //"AliasMethodDiscreteSamplerAlpha1",
                //"AliasMethodDiscreteSamplerAlpha2",

                // The AliasMethod memory requirement doubles for each alpha increment.
                // A fair comparison is to use 2^alpha for the equivalent guide table method.
                //"GuideTableDiscreteSamplerAlpha2",
                //"GuideTableDiscreteSamplerAlpha4",
        })
        private String samplerType;

        /** The factory. */
        private Supplier<DiscreteSampler> factory;

        /** The sampler. */
        private DiscreteSampler sampler;

        /**
         * Gets the sampler.
         *
         * @return the sampler.
         */
        public DiscreteSampler getSampler() {
            return sampler;
        }

        /** Create the distribution (per iteration as it may vary) and instantiates sampler. */
        @Override
        @Setup(Level.Iteration)
        public void setup() {
            super.setup();

            final double[] probabilities = createProbabilities();
            createSamplerFactory(getGenerator(), probabilities);
            sampler = factory.get();
        }

        /**
         * Creates the probabilities for the distribution.
         *
         * @return The probabilities.
         */
        protected abstract double[] createProbabilities();

        /**
         * Creates the sampler factory.
         *
         * @param rng The random generator.
         * @param probabilities The probabilities.
         */
        private void createSamplerFactory(final UniformRandomProvider rng,
            final double[] probabilities) {
            if ("BinarySearchDiscreteSampler".equals(samplerType)) {
                factory = () -> new BinarySearchDiscreteSampler(rng, probabilities);
            } else if ("AliasMethodDiscreteSampler".equals(samplerType)) {
                factory = () -> AliasMethodDiscreteSampler.of(rng, probabilities);
            } else if ("AliasMethodDiscreteSamplerNoPad".equals(samplerType)) {
                factory = () -> AliasMethodDiscreteSampler.of(rng, probabilities, -1);
            } else if ("AliasMethodDiscreteSamplerAlpha1".equals(samplerType)) {
                factory = () -> AliasMethodDiscreteSampler.of(rng, probabilities, 1);
            } else if ("AliasMethodDiscreteSamplerAlpha2".equals(samplerType)) {
                factory = () -> AliasMethodDiscreteSampler.of(rng, probabilities, 2);
            } else if ("GuideTableDiscreteSampler".equals(samplerType)) {
                factory = () -> GuideTableDiscreteSampler.of(rng, probabilities);
            } else if ("GuideTableDiscreteSamplerAlpha2".equals(samplerType)) {
                factory = () -> GuideTableDiscreteSampler.of(rng, probabilities, 2);
            } else if ("GuideTableDiscreteSamplerAlpha8".equals(samplerType)) {
                factory = () -> GuideTableDiscreteSampler.of(rng, probabilities, 8);
            } else if ("MarsagliaTsangWangDiscreteSampler".equals(samplerType)) {
                factory = () -> MarsagliaTsangWangDiscreteSampler.Enumerated.of(rng, probabilities);
            } else if ("FastLoadedDiceRollerDiscreteSampler".equals(samplerType)) {
                factory = () -> FastLoadedDiceRollerDiscreteSampler.of(rng, probabilities);
            } else if ("FastLoadedDiceRollerDiscreteSamplerLong".equals(samplerType)) {
                // Avoid exact floating-point arithmetic in construction.
                // Frequencies must sum to less than 2^63; here the sum is ~2^62.
                // This conversion may omit very small probabilities.
                final double sum = Arrays.stream(probabilities).sum();
                final long[] frequencies = Arrays.stream(probabilities)
                                                 .mapToLong(x -> Math.round(0x1.0p62 * x / sum))
                                                 .toArray();
                factory = () -> FastLoadedDiceRollerDiscreteSampler.of(rng, frequencies);
            } else if ("FastLoadedDiceRollerDiscreteSampler53".equals(samplerType)) {
                factory = () -> FastLoadedDiceRollerDiscreteSampler.of(rng, probabilities, 53);
            } else {
                throw new IllegalStateException();
            }
        }

        /**
         * Creates a new instance of the sampler.
         *
         * @return The sampler.
         */
        public DiscreteSampler createSampler() {
            return factory.get();
        }
    }

    /**
     * Define known probability distributions for testing. These are expected to have well
     * behaved cumulative probability functions.
     */
    @State(Scope.Benchmark)
    public static class KnownDistributionSources extends SamplerSources {
        /** The cumulative probability limit for unbounded distributions. */
        private static final double CUMULATIVE_PROBABILITY_LIMIT = 1 - 1e-9;
        /** Binomial distribution number of trials. */
        private static final int BINOM_N = 67;
        /** Binomial distribution probability of success. */
        private static final double BINOM_P = 0.7;
        /** Geometric distribution probability of success. */
        private static final double GEO_P = 0.2;
        /** Poisson distribution mean. */
        private static final double POISS_MEAN = 3.22;
        /** Bimodal distribution mean 1. */
        private static final double BIMOD_MEAN1 = 10;
        /** Bimodal distribution mean 1. */
        private static final double BIMOD_MEAN2 = 20;

        /**
         * The distribution.
         */
        @Param({"Binomial_N67_P0.7",
                "Geometric_P0.2",
                "4SidedLoadedDie",
                "Poisson_Mean3.22",
                "Poisson_Mean10_Mean20"})
        private String distribution;

        /** {@inheritDoc} */
        @Override
        protected double[] createProbabilities() {
            if ("Binomial_N67_P0.7".equals(distribution)) {
                final BinomialDistribution dist = new BinomialDistribution(null, BINOM_N, BINOM_P);
                return createProbabilities(dist, 0, BINOM_N);
            } else if ("Geometric_P0.2".equals(distribution)) {
                final double probabilityOfFailure = 1 - GEO_P;
                // https://en.wikipedia.org/wiki/Geometric_distribution
                // PMF = (1-p)^k * p
                // k is number of failures before a success
                double p = 1.0; // (1-p)^0
                // Build until the cumulative function is big
                double[] probabilities = new double[100];
                double sum = 0;
                int k = 0;
                while (k < probabilities.length) {
                    probabilities[k] = p * GEO_P;
                    sum += probabilities[k++];
                    if (sum > CUMULATIVE_PROBABILITY_LIMIT) {
                        break;
                    }
                    // For the next PMF
                    p *= probabilityOfFailure;
                }
                return Arrays.copyOf(probabilities, k);
            } else if ("4SidedLoadedDie".equals(distribution)) {
                return new double[] {1.0 / 2, 1.0 / 3, 1.0 / 12, 1.0 / 12};
            } else if ("Poisson_Mean3.22".equals(distribution)) {
                final IntegerDistribution dist = createPoissonDistribution(POISS_MEAN);
                final int max = dist.inverseCumulativeProbability(CUMULATIVE_PROBABILITY_LIMIT);
                return createProbabilities(dist, 0, max);
            } else if ("Poisson_Mean10_Mean20".equals(distribution)) {
                // Create a Bimodel using two Poisson distributions
                final IntegerDistribution dist1 = createPoissonDistribution(BIMOD_MEAN2);
                final int max = dist1.inverseCumulativeProbability(CUMULATIVE_PROBABILITY_LIMIT);
                final double[] p1 = createProbabilities(dist1, 0, max);
                final double[] p2 = createProbabilities(createPoissonDistribution(BIMOD_MEAN1), 0, max);
                for (int i = 0; i < p1.length; i++) {
                    p1[i] += p2[i];
                }
                // Leave to the distribution to normalise the sum
                return p1;
            }
            throw new IllegalStateException();
        }

        /**
         * Creates the poisson distribution.
         *
         * @param mean the mean
         * @return the distribution
         */
        private static IntegerDistribution createPoissonDistribution(double mean) {
            return new PoissonDistribution(null, mean,
                PoissonDistribution.DEFAULT_EPSILON, PoissonDistribution.DEFAULT_MAX_ITERATIONS);
        }

        /**
         * Creates the probabilities from the distribution.
         *
         * @param dist the distribution
         * @param lower the lower bounds (inclusive)
         * @param upper the upper bounds (inclusive)
         * @return the probabilities
         */
        private static double[] createProbabilities(IntegerDistribution dist, int lower, int upper) {
            double[] probabilities = new double[upper - lower + 1];
            int index = 0;
            for (int x = lower; x <= upper; x++) {
                probabilities[index++] = dist.probability(x);
            }
            return probabilities;
        }
    }

    /**
     * Define random probability distributions of known size for testing. These are random but
     * the average cumulative probability function will be a straight line given the increment
     * average is 0.5.
     */
    @State(Scope.Benchmark)
    public static class RandomDistributionSources extends SamplerSources {
        /**
         * The distribution size.
         * These are spaced half-way between powers-of-2 to minimise the advantage of
         * padding by the Alias method sampler.
         */
        @Param({"6",
                //"12",
                //"24",
                //"48",
                "96",
                //"192",
                //"384",
                // Above 2048 forces the Alias method to use more than 64-bits for sampling
                "3072"})
        private int randomNonUniformSize;

        /** {@inheritDoc} */
        @Override
        protected double[] createProbabilities() {
            return RandomSource.XO_RO_SHI_RO_128_PP.create()
                .doubles(randomNonUniformSize).toArray();
        }
    }

    /**
     * Sample random probability arrays from a Dirichlet distribution.
     *
     * <p>The distribution ensures the probabilities sum to 1.
     * The <a href="https://en.wikipedia.org/wiki/Entropy_(information_theory)">entropy</a>
     * of the probabilities increases with parameters k and alpha.
     * The following shows the mean and sd of the entropy from 100 samples
     * for a range of parameters.
     * <pre>
     * k   alpha   mean    sd
     * 4   0.500   1.299   0.374
     * 4   1.000   1.531   0.294
     * 4   2.000   1.754   0.172
     * 8   0.500   2.087   0.348
     * 8   1.000   2.490   0.266
     * 8   2.000   2.707   0.142
     * 16  0.500   3.023   0.287
     * 16  1.000   3.454   0.166
     * 16  2.000   3.693   0.095
     * 32  0.500   4.008   0.182
     * 32  1.000   4.406   0.125
     * 32  2.000   4.692   0.075
     * 64  0.500   4.986   0.151
     * 64  1.000   5.392   0.115
     * 64  2.000   5.680   0.048
     * </pre>
     */
    @State(Scope.Benchmark)
    public static class DirichletDistributionSources extends SamplerSources {
        /** Number of categories. */
        @Param({"4", "8", "16"})
        private int k;

        /** Concentration parameter. */
        @Param({"0.5", "1", "2"})
        private double alpha;

        /** {@inheritDoc} */
        @Override
        protected double[] createProbabilities() {
            return DirichletSampler.symmetric(RandomSource.XO_RO_SHI_RO_128_PP.create(),
                                              k, alpha).sample();
        }
    }

    /**
     * The {@link FastLoadedDiceRollerDiscreteSampler} samplers to use for testing.
     * Creates the sampler for each random source and the probabilities using
     * a Dirichlet distribution.
     *
     * <p>This class is a specialized source to allow examination of the effect of the
     * {@link FastLoadedDiceRollerDiscreteSampler} {@code alpha} parameter.
     */
    @State(Scope.Benchmark)
    public static class FastLoadedDiceRollerDiscreteSamplerSources extends LocalRandomSources {
        /** Number of categories. */
        @Param({"4", "8", "16"})
        private int k;

        /** Concentration parameter. */
        @Param({"0.5", "1", "2"})
        private double concentration;

        /** The constructor {@code alpha} parameter. */
        @Param({"0", "30", "53"})
        private int alpha;

        /** The factory. */
        private Supplier<DiscreteSampler> factory;

        /** The sampler. */
        private DiscreteSampler sampler;

        /**
         * Gets the sampler.
         *
         * @return the sampler.
         */
        public DiscreteSampler getSampler() {
            return sampler;
        }

        /** Create the distribution probabilities (per iteration as it may vary), the sampler
         * factory and instantiates sampler. */
        @Override
        @Setup(Level.Iteration)
        public void setup() {
            super.setup();

            final double[] probabilities =
                DirichletSampler.symmetric(RandomSource.XO_RO_SHI_RO_128_PP.create(),
                                           k, concentration).sample();
            final UniformRandomProvider rng = getGenerator();
            factory = () -> FastLoadedDiceRollerDiscreteSampler.of(rng, probabilities, alpha);
            sampler = factory.get();
        }

        /**
         * Creates a new instance of the sampler.
         *
         * @return The sampler.
         */
        public DiscreteSampler createSampler() {
            return factory.get();
        }
    }

    /**
     * Compute a sample by binary search of the cumulative probability distribution.
     */
    static final class BinarySearchDiscreteSampler
        implements DiscreteSampler {
        /** Underlying source of randomness. */
        private final UniformRandomProvider rng;
        /**
         * The cumulative probability table.
         */
        private final double[] cumulativeProbabilities;

        /**
         * @param rng Generator of uniformly distributed random numbers.
         * @param probabilities The probabilities.
         * @throws IllegalArgumentException if {@code probabilities} is null or empty, a
         * probability is negative, infinite or {@code NaN}, or the sum of all
         * probabilities is not strictly positive.
         */
        BinarySearchDiscreteSampler(UniformRandomProvider rng,
                                    double[] probabilities) {
            // Minimal set-up validation
            if (probabilities == null || probabilities.length == 0) {
                throw new IllegalArgumentException("Probabilities must not be empty.");
            }

            final int size = probabilities.length;
            cumulativeProbabilities = new double[size];

            double sumProb = 0;
            int count = 0;
            for (final double prob : probabilities) {
                if (prob < 0 ||
                    Double.isInfinite(prob) ||
                    Double.isNaN(prob)) {
                    throw new IllegalArgumentException("Invalid probability: " +
                                                       prob);
                }

                // Compute and store cumulative probability.
                sumProb += prob;
                cumulativeProbabilities[count++] = sumProb;
            }

            if (Double.isInfinite(sumProb) || sumProb <= 0) {
                throw new IllegalArgumentException("Invalid sum of probabilities: " + sumProb);
            }

            this.rng = rng;

            // Normalise cumulative probability.
            for (int i = 0; i < size; i++) {
                final double norm = cumulativeProbabilities[i] / sumProb;
                cumulativeProbabilities[i] = (norm < 1) ? norm : 1.0;
            }
        }

        /** {@inheritDoc} */
        @Override
        public int sample() {
            final double u = rng.nextDouble();

            // Java binary search
            //int index = Arrays.binarySearch(cumulativeProbabilities, u);
            //if (index < 0) {
            //    index = -index - 1;
            //}
            //
            //return index < cumulativeProbabilities.length ?
            //    index :
            //    cumulativeProbabilities.length - 1;

            // Binary search within known cumulative probability table.
            // Find x so that u > f[x-1] and u <= f[x].
            // This is a looser search than Arrays.binarySearch:
            // - The output is x = upper.
            // - The table stores probabilities where f[0] is >= 0 and the max == 1.0.
            // - u should be >= 0 and <= 1 (or the random generator is broken).
            // - It avoids comparisons using Double.doubleToLongBits.
            // - It avoids the low likelihood of equality between two doubles for fast exit
            //   so uses only 1 compare per loop.
            int lower = 0;
            int upper = cumulativeProbabilities.length - 1;
            while (lower < upper) {
                final int mid = (lower + upper) >>> 1;
                final double midVal = cumulativeProbabilities[mid];
                if (u > midVal) {
                    // Change lower such that
                    // u > f[lower - 1]
                    lower = mid + 1;
                } else {
                    // Change upper such that
                    // u <= f[upper]
                    upper = mid;
                }
            }
            return upper;
        }
    }

    // Benchmarks methods below.

    /**
     * Baseline for the JMH timing overhead for production of an {@code int} value.
     *
     * @return the {@code int} value
     */
    @Benchmark
    public int baselineInt() {
        return value;
    }

    /**
     * Baseline for the production of a {@code double} value.
     * This is used to assess the performance of the underlying random source.
     *
     * @param sources Source of randomness.
     * @return the {@code int} value
     */
    @Benchmark
    public int baselineNextDouble(LocalRandomSources sources) {
        return sources.getGenerator().nextDouble() < 0.5 ? 1 : 0;
    }

    /**
     * Run the sampler.
     *
     * @param sources Source of randomness.
     * @return the sample value
     */
    @Benchmark
    public int sampleKnown(KnownDistributionSources sources) {
        return sources.getSampler().sample();
    }

    /**
     * Create and run the sampler.
     *
     * @param sources Source of randomness.
     * @return the sample value
     */
    @Benchmark
    public int singleSampleKnown(KnownDistributionSources sources) {
        return sources.createSampler().sample();
    }

    /**
     * Run the sampler.
     *
     * @param sources Source of randomness.
     * @return the sample value
     */
    @Benchmark
    public int sampleRandom(RandomDistributionSources sources) {
        return sources.getSampler().sample();
    }

    /**
     * Create and run the sampler.
     *
     * @param sources Source of randomness.
     * @return the sample value
     */
    @Benchmark
    public int singleSampleRandom(RandomDistributionSources sources) {
        return sources.createSampler().sample();
    }

    /**
     * Run the sampler.
     *
     * @param sources Source of randomness.
     * @return the sample value
     */
    @Benchmark
    public int sampleDirichlet(DirichletDistributionSources sources) {
        return sources.getSampler().sample();
    }

    /**
     * Create and run the sampler.
     *
     * @param sources Source of randomness.
     * @return the sample value
     */
    @Benchmark
    public int singleSampleDirichlet(DirichletDistributionSources sources) {
        return sources.createSampler().sample();
    }

    /**
     * Run the sampler.
     *
     * @param sources Source of randomness.
     * @return the sample value
     */
    @Benchmark
    public int sampleFast(FastLoadedDiceRollerDiscreteSamplerSources sources) {
        return sources.getSampler().sample();
    }

    /**
     * Create and run the sampler.
     *
     * @param sources Source of randomness.
     * @return the sample value
     */
    @Benchmark
    public int singleSampleFast(FastLoadedDiceRollerDiscreteSamplerSources sources) {
        return sources.createSampler().sample();
    }
}
