/*
 * 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 java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import org.apache.commons.math3.util.MathArrays;
import org.apache.commons.rng.UniformRandomProvider;
import org.apache.commons.rng.sampling.RandomAssert;

/**
 * List of samplers.
 */
public final class DiscreteSamplersList {
    /** List of all RNGs implemented in the library. */
    private static final List<DiscreteSamplerTestData> LIST = new ArrayList<>();

    static {
        try {
            // This test uses reference distributions from commons-math3 to compute the expected
            // PMF. These distributions have a dual functionality to compute the PMF and perform
            // sampling. When no sampling is needed for the created distribution, it is advised
            // to pass null as the random generator via the appropriate constructors to avoid the
            // additional initialisation overhead.
            org.apache.commons.math3.random.RandomGenerator unusedRng = null;

            // List of distributions to test.

            // Binomial ("inverse method").
            final int trialsBinomial = 20;
            final double probSuccessBinomial = 0.67;
            add(LIST, new org.apache.commons.math3.distribution.BinomialDistribution(unusedRng, trialsBinomial, probSuccessBinomial),
                MathArrays.sequence(8, 9, 1),
                RandomAssert.createRNG());
            add(LIST, new org.apache.commons.math3.distribution.BinomialDistribution(unusedRng, trialsBinomial, probSuccessBinomial),
                // range [9,16]
                MathArrays.sequence(8, 9, 1),
                MarsagliaTsangWangDiscreteSampler.Binomial.of(RandomAssert.createRNG(), trialsBinomial, probSuccessBinomial));
            // Inverted
            add(LIST, new org.apache.commons.math3.distribution.BinomialDistribution(unusedRng, trialsBinomial, 1 - probSuccessBinomial),
                // range [4,11] = [20-16, 20-9]
                MathArrays.sequence(8, 4, 1),
                MarsagliaTsangWangDiscreteSampler.Binomial.of(RandomAssert.createRNG(), trialsBinomial, 1 - probSuccessBinomial));

            // Geometric ("inverse method").
            final double probSuccessGeometric = 0.21;
            add(LIST, new org.apache.commons.math3.distribution.GeometricDistribution(unusedRng, probSuccessGeometric),
                MathArrays.sequence(10, 0, 1),
                RandomAssert.createRNG());
            // Geometric.
            add(LIST, new org.apache.commons.math3.distribution.GeometricDistribution(unusedRng, probSuccessGeometric),
                MathArrays.sequence(10, 0, 1),
                GeometricSampler.of(RandomAssert.createRNG(), probSuccessGeometric));

            // Hypergeometric ("inverse method").
            final int popSizeHyper = 34;
            final int numSuccessesHyper = 11;
            final int sampleSizeHyper = 12;
            add(LIST, new org.apache.commons.math3.distribution.HypergeometricDistribution(unusedRng, popSizeHyper, numSuccessesHyper, sampleSizeHyper),
                MathArrays.sequence(10, 0, 1),
                RandomAssert.createRNG());

            // Pascal ("inverse method").
            final int numSuccessesPascal = 6;
            final double probSuccessPascal = 0.2;
            add(LIST, new org.apache.commons.math3.distribution.PascalDistribution(unusedRng, numSuccessesPascal, probSuccessPascal),
                MathArrays.sequence(18, 1, 1),
                RandomAssert.createRNG());

            // Uniform ("inverse method").
            final int loUniform = -3;
            final int hiUniform = 4;
            add(LIST, new org.apache.commons.math3.distribution.UniformIntegerDistribution(unusedRng, loUniform, hiUniform),
                MathArrays.sequence(8, -3, 1),
                RandomAssert.createRNG());
            // Uniform (power of 2 range).
            add(LIST, new org.apache.commons.math3.distribution.UniformIntegerDistribution(unusedRng, loUniform, hiUniform),
                MathArrays.sequence(8, -3, 1),
                DiscreteUniformSampler.of(RandomAssert.createRNG(), loUniform, hiUniform));
            // Uniform (large range).
            final int halfMax = Integer.MAX_VALUE / 2;
            final int hiLargeUniform = halfMax + 10;
            final int loLargeUniform = -hiLargeUniform;
            add(LIST, new org.apache.commons.math3.distribution.UniformIntegerDistribution(unusedRng, loLargeUniform, hiLargeUniform),
                MathArrays.sequence(20, -halfMax, halfMax / 10),
                DiscreteUniformSampler.of(RandomAssert.createRNG(), loLargeUniform, hiLargeUniform));
            // Uniform (non-power of 2 range).
            final int rangeNonPowerOf2Uniform = 11;
            final int hiNonPowerOf2Uniform = loUniform + rangeNonPowerOf2Uniform;
            add(LIST, new org.apache.commons.math3.distribution.UniformIntegerDistribution(unusedRng, loUniform, hiNonPowerOf2Uniform),
                MathArrays.sequence(rangeNonPowerOf2Uniform, -3, 1),
                DiscreteUniformSampler.of(RandomAssert.createRNG(), loUniform, hiNonPowerOf2Uniform));

            // Zipf ("inverse method").
            final int numElementsZipf = 5;
            final double exponentZipf = 2.345;
            add(LIST, new org.apache.commons.math3.distribution.ZipfDistribution(unusedRng, numElementsZipf, exponentZipf),
                MathArrays.sequence(5, 1, 1),
                RandomAssert.createRNG());
            // Zipf.
            add(LIST, new org.apache.commons.math3.distribution.ZipfDistribution(unusedRng, numElementsZipf, exponentZipf),
                MathArrays.sequence(5, 1, 1),
                RejectionInversionZipfSampler.of(RandomAssert.createRNG(), numElementsZipf, exponentZipf));
            // Zipf (exponent close to 1).
            final double exponentCloseToOneZipf = 1 - 1e-10;
            add(LIST, new org.apache.commons.math3.distribution.ZipfDistribution(unusedRng, numElementsZipf, exponentCloseToOneZipf),
                MathArrays.sequence(5, 1, 1),
                RejectionInversionZipfSampler.of(RandomAssert.createRNG(), numElementsZipf, exponentCloseToOneZipf));
            // Zipf (exponent = 0).
            add(LIST, MathArrays.sequence(5, 1, 1), new double[] {0.2, 0.2, 0.2, 0.2, 0.2},
                RejectionInversionZipfSampler.of(RandomAssert.createRNG(), numElementsZipf, 0.0));

            // Poisson ("inverse method").
            final double epsilonPoisson = org.apache.commons.math3.distribution.PoissonDistribution.DEFAULT_EPSILON;
            final int maxIterationsPoisson = org.apache.commons.math3.distribution.PoissonDistribution.DEFAULT_MAX_ITERATIONS;
            final double meanPoisson = 3.21;
            add(LIST, new org.apache.commons.math3.distribution.PoissonDistribution(unusedRng, meanPoisson, epsilonPoisson, maxIterationsPoisson),
                MathArrays.sequence(10, 0, 1),
                RandomAssert.createRNG());
            // Poisson.
            add(LIST, new org.apache.commons.math3.distribution.PoissonDistribution(unusedRng, meanPoisson, epsilonPoisson, maxIterationsPoisson),
                MathArrays.sequence(10, 0, 1),
                PoissonSampler.of(RandomAssert.createRNG(), meanPoisson));
            // Dedicated small mean poisson samplers
            add(LIST, new org.apache.commons.math3.distribution.PoissonDistribution(unusedRng, meanPoisson, epsilonPoisson, maxIterationsPoisson),
                MathArrays.sequence(10, 0, 1),
                SmallMeanPoissonSampler.of(RandomAssert.createRNG(), meanPoisson));
            add(LIST, new org.apache.commons.math3.distribution.PoissonDistribution(unusedRng, meanPoisson, epsilonPoisson, maxIterationsPoisson),
                MathArrays.sequence(10, 0, 1),
                KempSmallMeanPoissonSampler.of(RandomAssert.createRNG(), meanPoisson));
            add(LIST, new org.apache.commons.math3.distribution.PoissonDistribution(unusedRng, meanPoisson, epsilonPoisson, maxIterationsPoisson),
                MathArrays.sequence(10, 0, 1),
                MarsagliaTsangWangDiscreteSampler.Poisson.of(RandomAssert.createRNG(), meanPoisson));
            // LargeMeanPoissonSampler should work at small mean.
            // Note: This hits a code path where the sample from the normal distribution is rejected.
            add(LIST, new org.apache.commons.math3.distribution.PoissonDistribution(unusedRng, meanPoisson, epsilonPoisson, maxIterationsPoisson),
                MathArrays.sequence(10, 0, 1),
                LargeMeanPoissonSampler.of(RandomAssert.createRNG(), meanPoisson));
            // Poisson (40 < mean < 80).
            final double largeMeanPoisson = 67.89;
            add(LIST, new org.apache.commons.math3.distribution.PoissonDistribution(unusedRng, largeMeanPoisson, epsilonPoisson, maxIterationsPoisson),
                MathArrays.sequence(50, (int) (largeMeanPoisson - 25), 1),
                PoissonSampler.of(RandomAssert.createRNG(), largeMeanPoisson));
            // Dedicated large mean poisson sampler
            add(LIST, new org.apache.commons.math3.distribution.PoissonDistribution(unusedRng, largeMeanPoisson, epsilonPoisson, maxIterationsPoisson),
                MathArrays.sequence(50, (int) (largeMeanPoisson - 25), 1),
                LargeMeanPoissonSampler.of(RandomAssert.createRNG(), largeMeanPoisson));
            add(LIST, new org.apache.commons.math3.distribution.PoissonDistribution(unusedRng, largeMeanPoisson, epsilonPoisson, maxIterationsPoisson),
                MathArrays.sequence(50, (int) (largeMeanPoisson - 25), 1),
                MarsagliaTsangWangDiscreteSampler.Poisson.of(RandomAssert.createRNG(), largeMeanPoisson));
            // Poisson (mean >> 40).
            final double veryLargeMeanPoisson = 543.21;
            add(LIST, new org.apache.commons.math3.distribution.PoissonDistribution(unusedRng, veryLargeMeanPoisson, epsilonPoisson, maxIterationsPoisson),
                MathArrays.sequence(100, (int) (veryLargeMeanPoisson - 50), 1),
                PoissonSampler.of(RandomAssert.createRNG(), veryLargeMeanPoisson));
            // Dedicated large mean poisson sampler
            add(LIST, new org.apache.commons.math3.distribution.PoissonDistribution(unusedRng, veryLargeMeanPoisson, epsilonPoisson, maxIterationsPoisson),
                MathArrays.sequence(100, (int) (veryLargeMeanPoisson - 50), 1),
                LargeMeanPoissonSampler.of(RandomAssert.createRNG(), veryLargeMeanPoisson));
            add(LIST, new org.apache.commons.math3.distribution.PoissonDistribution(unusedRng, veryLargeMeanPoisson, epsilonPoisson, maxIterationsPoisson),
                MathArrays.sequence(100, (int) (veryLargeMeanPoisson - 50), 1),
                MarsagliaTsangWangDiscreteSampler.Poisson.of(RandomAssert.createRNG(), veryLargeMeanPoisson));

            // Any discrete distribution
            final int[] discretePoints = {0, 1, 2, 3, 4};
            final double[] discreteProbabilities = {0.1, 0.2, 0.3, 0.4, 0.5};
            final long[] discreteFrequencies = {1, 2, 3, 4, 5};
            add(LIST, discretePoints, discreteProbabilities,
                MarsagliaTsangWangDiscreteSampler.Enumerated.of(RandomAssert.createRNG(), discreteProbabilities));
            add(LIST, discretePoints, discreteProbabilities,
                GuideTableDiscreteSampler.of(RandomAssert.createRNG(), discreteProbabilities));
            add(LIST, discretePoints, discreteProbabilities,
                AliasMethodDiscreteSampler.of(RandomAssert.createRNG(), discreteProbabilities));
            add(LIST, discretePoints, discreteProbabilities,
                FastLoadedDiceRollerDiscreteSampler.of(RandomAssert.createRNG(), discreteFrequencies));
            add(LIST, discretePoints, discreteProbabilities,
                FastLoadedDiceRollerDiscreteSampler.of(RandomAssert.createRNG(), discreteProbabilities));
        } catch (Exception e) {
            // CHECKSTYLE: stop Regexp
            System.err.println("Unexpected exception while creating the list of samplers: " + e);
            e.printStackTrace(System.err);
            // CHECKSTYLE: resume Regexp
            throw new RuntimeException(e);
        }
    }

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

    /**
     * @param list List of data (one the "parameters" tested by the Junit parametric test).
     * @param dist Distribution to which the samples are supposed to conform.
     * @param points Outcomes selection.
     * @param rng Generator of uniformly distributed sequences.
     */
    private static void add(List<DiscreteSamplerTestData> list,
                            final org.apache.commons.math3.distribution.IntegerDistribution dist,
                            int[] points,
                            UniformRandomProvider rng) {
        final DiscreteSampler inverseMethodSampler =
            InverseTransformDiscreteSampler.of(rng,
                new DiscreteInverseCumulativeProbabilityFunction() {
                    @Override
                    public int inverseCumulativeProbability(double p) {
                        return dist.inverseCumulativeProbability(p);
                    }
                    @Override
                    public String toString() {
                        return dist.toString();
                    }
                });
        list.add(new DiscreteSamplerTestData(inverseMethodSampler,
                                             points,
                                             getProbabilities(dist, points)));
    }

    /**
     * @param list List of data (one the "parameters" tested by the Junit parametric test).
     * @param dist Distribution to which the samples are supposed to conform.
     * @param points Outcomes selection.
     * @param sampler Sampler.
     */
    private static void add(List<DiscreteSamplerTestData> list,
                            final org.apache.commons.math3.distribution.IntegerDistribution dist,
                            int[] points,
                            final DiscreteSampler sampler) {
        list.add(new DiscreteSamplerTestData(sampler,
                                             points,
                                             getProbabilities(dist, points)));
    }

    /**
     * @param list List of data (one the "parameters" tested by the Junit parametric test).
     * @param points Outcomes selection.
     * @param probabilities Probability distribution to which the samples are supposed to conform.
     * @param sampler Sampler.
     */
    private static void add(List<DiscreteSamplerTestData> list,
                            int[] points,
                            final double[] probabilities,
                            final DiscreteSampler sampler) {
        list.add(new DiscreteSamplerTestData(sampler,
                                             points,
                                             probabilities));
    }

    /**
     * Subclasses that are "parametric" tests can forward the call to
     * the "@Parameters"-annotated method to this method.
     *
     * @return the list of all generators.
     */
    public static Iterable<DiscreteSamplerTestData> list() {
        return Collections.unmodifiableList(LIST);
    }

    /**
     * @param dist Distribution.
     * @param points Points.
     * @return the probabilities of the given points according to the distribution.
     */
    private static double[] getProbabilities(org.apache.commons.math3.distribution.IntegerDistribution dist,
                                             int[] points) {
        final int len = points.length;
        final double[] prob = new double[len];
        for (int i = 0; i < len; i++) {
            prob[i] = dist instanceof org.apache.commons.math3.distribution.UniformIntegerDistribution ? // XXX Workaround.
                getProbability((org.apache.commons.math3.distribution.UniformIntegerDistribution) dist) :
                dist.probability(points[i]);

            if (prob[i] < 0) {
                throw new IllegalStateException(dist + ": p < 0 (at " + points[i] + ", p=" + prob[i]);
            }
        }
        return prob;
    }

    /**
     * Workaround bugs in Commons Math's "UniformIntegerDistribution" (cf. MATH-1396).
     */
    private static double getProbability(org.apache.commons.math3.distribution.UniformIntegerDistribution dist) {
        return 1 / ((double) dist.getSupportUpperBound() - (double) dist.getSupportLowerBound() + 1);
    }
}
