/*
 * 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.junit.jupiter.api.Assertions;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.Arguments;
import org.junit.jupiter.params.provider.MethodSource;
import java.util.Arrays;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Stream;
import org.apache.commons.math3.distribution.AbstractRealDistribution;
import org.apache.commons.math3.distribution.ExponentialDistribution;
import org.apache.commons.math3.distribution.NormalDistribution;
import org.apache.commons.math3.stat.inference.ChiSquareTest;
import org.apache.commons.rng.RestorableUniformRandomProvider;
import org.apache.commons.rng.UniformRandomProvider;
import org.apache.commons.rng.sampling.distribution.ContinuousSampler;
import org.apache.commons.rng.simple.RandomSource;

/**
 * Test for ziggurat samplers in the {@link ZigguratSamplerPerformance} class.
 *
 * <p>This test is copied from the {@code commons-rng-sampling} module to ensure all implementations
 * correctly sample from the distribution.
 */
class ZigguratSamplerTest {

    /**
     * The seed for the RNG used in the distribution sampling tests.
     *
     * <p>This has been chosen to allow the test to pass with all generators.
     * Set to null test with a random seed.
     *
     * <p>Note that the p-value of the chi-square test is 0.001. There are multiple assertions
     * per test and multiple samplers. The total number of chi-square tests is above 100
     * and failure of a chosen random seed on a few tests is common. When using a random
     * seed re-run the test multiple times. Systematic failure of the same sampler
     * should be investigated further.
     */
    private static final Long SEED = 0xd1342543de82ef95L;

    /**
     * Create arguments with the name of the factory.
     *
     * @param name Name of the factory
     * @param factory Factory to create the sampler
     * @return the arguments
     */
    private static Arguments args(String name) {
        // Create the factory.
        // Here we delegate to the static method used to create all the samplers for testing.
        final Function<UniformRandomProvider, ContinuousSampler> factory =
            rng -> ZigguratSamplerPerformance.Sources.createSampler(name, rng);
        return Arguments.of(name, factory);
    }

    /**
     * Create a stream of constructors of a Gaussian sampler.
     *
     * <p>Note: This method exists to allow this test to be duplicated in the examples JMH
     * module where many implementations are tested.
     *
     * @return the stream of constructors
     */
    private static Stream<Arguments> gaussianSamplers() {
        // Test all but MOD_GAUSSIAN (tested in the common-rng-sampling module)
        return Stream.of(
            args(ZigguratSamplerPerformance.GAUSSIAN_128),
            args(ZigguratSamplerPerformance.GAUSSIAN_256),
            args(ZigguratSamplerPerformance.MOD_GAUSSIAN2),
            args(ZigguratSamplerPerformance.MOD_GAUSSIAN_SIMPLE_OVERHANGS),
            args(ZigguratSamplerPerformance.MOD_GAUSSIAN_INLINING),
            args(ZigguratSamplerPerformance.MOD_GAUSSIAN_INLINING_SHIFT),
            args(ZigguratSamplerPerformance.MOD_GAUSSIAN_INLINING_SIMPLE_OVERHANGS),
            args(ZigguratSamplerPerformance.MOD_GAUSSIAN_INT_MAP),
            args(ZigguratSamplerPerformance.MOD_GAUSSIAN_E_MAX_TABLE),
            args(ZigguratSamplerPerformance.MOD_GAUSSIAN_E_MAX_2),
            args(ZigguratSamplerPerformance.MOD_GAUSSIAN_TERNARY),
            args(ZigguratSamplerPerformance.MOD_GAUSSIAN_512));
    }

    /**
     * Create a stream of constructors of an exponential sampler.
     *
     * <p>Note: This method exists to allow this test to be duplicated in the examples JMH
     * module where many implementations are tested.
     *
     * @return the stream of constructors
     */
    private static Stream<Arguments> exponentialSamplers() {
        // Test all but MOD_EXPONENTIAL (tested in the common-rng-sampling module)
        return Stream.of(
                args(ZigguratSamplerPerformance.EXPONENTIAL),
                args(ZigguratSamplerPerformance.MOD_EXPONENTIAL2),
                args(ZigguratSamplerPerformance.MOD_EXPONENTIAL_SIMPLE_OVERHANGS),
                args(ZigguratSamplerPerformance.MOD_EXPONENTIAL_INLINING),
                args(ZigguratSamplerPerformance.MOD_EXPONENTIAL_LOOP),
                args(ZigguratSamplerPerformance.MOD_EXPONENTIAL_LOOP2),
                args(ZigguratSamplerPerformance.MOD_EXPONENTIAL_RECURSION),
                args(ZigguratSamplerPerformance.MOD_EXPONENTIAL_INT_MAP),
                args(ZigguratSamplerPerformance.MOD_EXPONENTIAL_E_MAX_TABLE),
                args(ZigguratSamplerPerformance.MOD_EXPONENTIAL_E_MAX_2),
                args(ZigguratSamplerPerformance.MOD_EXPONENTIAL_TERNARY),
                args(ZigguratSamplerPerformance.MOD_EXPONENTIAL_TERNARY_SUBTRACT),
                args(ZigguratSamplerPerformance.MOD_EXPONENTIAL_512));
    }

    // -------------------------------------------------------------------------
    // All code below here is copied from:
    // commons-rng-sampling/src/test/java/org/apache/commons/rng/sampling/distribution/ZigguratSamplerTest.java
    // -------------------------------------------------------------------------

    /**
     * Creates the gaussian distribution.
     *
     * @return the distribution
     */
    private static AbstractRealDistribution createGaussianDistribution() {
        return new NormalDistribution(null, 0.0, 1.0);
    }

    /**
     * Creates the exponential distribution.
     *
     * @return the distribution
     */
    private static AbstractRealDistribution createExponentialDistribution() {
        return new ExponentialDistribution(null, 1.0);
    }

    /**
     * Test Gaussian samples using a large number of bins based on uniformly spaced
     * quantiles. Added for RNG-159.
     *
     * @param name Name of the sampler
     * @param factory Factory to create the sampler
     */
    @ParameterizedTest(name = "{index} => {0}")
    @MethodSource("gaussianSamplers")
    void testGaussianSamplesWithQuantiles(String name, Function<UniformRandomProvider, ContinuousSampler> factory) {
        final int bins = 2000;
        final AbstractRealDistribution dist = createGaussianDistribution();
        final double[] quantiles = new double[bins];
        for (int i = 0; i < bins; i++) {
            quantiles[i] = dist.inverseCumulativeProbability((i + 1.0) / bins);
        }
        testSamples(quantiles, factory, ZigguratSamplerTest::createGaussianDistribution,
            // Test positive and negative ranges to check symmetry.
            // Smallest layer is 0.292 (convex region)
            new double[] {0, 0.2},
            new double[] {-0.35, -0.1},
            // Around the mean
            new double[] {-0.1, 0.1},
            new double[] {-0.4, 0.6},
            // Inflection point at x=1
            new double[] {-1.1, -0.9},
            // A concave region
            new double[] {2.1, 2.5},
            // Tail = 3.64
            new double[] {2.5, 8});
    }

    /**
     * Test Gaussian samples using a large number of bins uniformly spaced in a range.
     * Added for RNG-159.
     *
     * @param name Name of the sampler
     * @param factory Factory to create the sampler
     */
    @ParameterizedTest(name = "{index} => {0}")
    @MethodSource("gaussianSamplers")
    void testGaussianSamplesWithUniformValues(String name, Function<UniformRandomProvider, ContinuousSampler> factory) {
        final int bins = 2000;
        final double[] values = new double[bins];
        final double minx = -8;
        final double maxx = 8;
        // Bin width = 16 / 2000 = 0.008
        for (int i = 0; i < bins; i++) {
            values[i] = minx + (maxx - minx) * (i + 1.0) / bins;
        }
        // Ensure upper bound is the support limit
        values[bins - 1] = Double.POSITIVE_INFINITY;
        testSamples(values, factory, ZigguratSamplerTest::createGaussianDistribution,
            // Test positive and negative ranges to check symmetry.
            // Smallest layer is 0.292 (convex region)
            new double[] {0, 0.2},
            new double[] {-0.35, -0.1},
            // Around the mean
            new double[] {-0.1, 0.1},
            new double[] {-0.4, 0.6},
            // Inflection point at x=1
            new double[] {-1.01, -0.99},
            new double[] {0.98, 1.03},
            // A concave region
            new double[] {1.03, 1.05},
            // Tail = 3.64
            new double[] {3.6, 3.8},
            new double[] {3.7, 8});
    }

    /**
     * Test exponential samples using a large number of bins based on uniformly spaced quantiles.
     *
     * @param name Name of the sampler
     * @param factory Factory to create the sampler
     */
    @ParameterizedTest(name = "{index} => {0}")
    @MethodSource("exponentialSamplers")
    void testExponentialSamplesWithQuantiles(String name, Function<UniformRandomProvider, ContinuousSampler> factory) {
        final int bins = 2000;
        final AbstractRealDistribution dist = createExponentialDistribution();
        final double[] quantiles = new double[bins];
        for (int i = 0; i < bins; i++) {
            quantiles[i] = dist.inverseCumulativeProbability((i + 1.0) / bins);
        }
        testSamples(quantiles, factory, ZigguratSamplerTest::createExponentialDistribution,
            // Smallest layer is 0.122
            new double[] {0, 0.1},
            new double[] {0.05, 0.15},
            // Around the mean
            new double[] {0.9, 1.1},
            // Tail = 7.57
            new double[] {1.5, 12});
    }

    /**
     * Test exponential samples using a large number of bins uniformly spaced in a range.
     *
     * @param name Name of the sampler
     * @param factory Factory to create the sampler
     */
    @ParameterizedTest(name = "{index} => {0}")
    @MethodSource("exponentialSamplers")
    void testExponentialSamplesWithUniformValues(String name, Function<UniformRandomProvider, ContinuousSampler> factory) {
        final int bins = 2000;
        final double[] values = new double[bins];
        final double minx = 0;
        // Enter the tail of the distribution
        final double maxx = 12;
        // Bin width = 12 / 2000 = 0.006
        for (int i = 0; i < bins; i++) {
            values[i] = minx + (maxx - minx) * (i + 1.0) / bins;
        }
        // Ensure upper bound is the support limit
        values[bins - 1] = Double.POSITIVE_INFINITY;

        testSamples(values, factory,  ZigguratSamplerTest::createExponentialDistribution,
            // Smallest layer is 0.122
            new double[] {0, 0.1},
            new double[] {0.05, 0.15},
            // Around the mean
            new double[] {0.9, 1.1},
            // Tail = 7.57
            new double[] {7.5, 7.7},
            new double[] {7.7, 12});
    }

    /**
     * Test samples using the provided bins. Values correspond to the bin upper
     * limit. It is assumed the values span most of the distribution. Additional
     * tests are performed using a region of the distribution sampled.
     *
     * @param values Bin upper limits
     * @param factory Factory to create the sampler
     * @param distribution The distribution under test
     * @param ranges Ranges of the distribution to test
     */
    private static void testSamples(double[] values,
                                    Function<UniformRandomProvider, ContinuousSampler> factory,
                                    Supplier<AbstractRealDistribution> distribution,
                                    double[]... ranges) {
        final int bins = values.length;

        final int samples = 10000000;
        final long[] observed = new long[bins];
        final RestorableUniformRandomProvider rng = RandomSource.XO_SHI_RO_128_PP.create(SEED);
        final ContinuousSampler sampler = factory.apply(rng);
        for (int i = 0; i < samples; i++) {
            final double x = sampler.sample();
            final int index = findIndex(values, x);
            observed[index]++;
        }

        // Compute expected
        final AbstractRealDistribution dist = distribution.get();
        final double[] expected = new double[bins];
        double x0 = Double.NEGATIVE_INFINITY;
        for (int i = 0; i < bins; i++) {
            final double x1 = values[i];
            expected[i] = dist.probability(x0, x1);
            x0 = x1;
        }

        final double significanceLevel = 0.001;

        final double lowerBound = dist.getSupportLowerBound();

        final ChiSquareTest chiSquareTest = new ChiSquareTest();
        // Pass if we cannot reject null hypothesis that the distributions are the same.
        final double pValue = chiSquareTest.chiSquareTest(expected, observed);
        Assertions.assertFalse(pValue < 0.001,
            () -> String.format("(%s <= x < %s) Chi-square p-value = %s",
                                lowerBound, values[bins - 1], pValue));

        // Test regions of the ziggurat.
        for (final double[] range : ranges) {
            final int min = findIndex(values, range[0]);
            final int max = findIndex(values, range[1]);
            // Must have a range of 2
            if (max - min + 1 < 2) {
                // This will probably occur if the quantiles test uses too small a range
                // for the tail. The tail is so far into the CDF that a single bin is
                // often used to represent it.
                Assertions.fail("Invalid range: " + Arrays.toString(range));
            }
            final long[] observed2 = Arrays.copyOfRange(observed, min, max + 1);
            final double[] expected2 = Arrays.copyOfRange(expected, min, max + 1);
            final double pValueB = chiSquareTest.chiSquareTest(expected2, observed2);
            Assertions.assertFalse(pValueB < significanceLevel,
                () -> String.format("(%s <= x < %s) Chi-square p-value = %s",
                                    min == 0 ? lowerBound : values[min - 1], values[max], pValueB));
        }
    }

    /**
     * Find the index of the value in the data such that:
     * <pre>
     * data[index - 1] <= x < data[index]
     * </pre>
     *
     * <p>This is a specialised binary search that assumes the bounds of the data are the
     * extremes of the support, and the upper support is infinite. Thus an index cannot
     * be returned as equal to the data length.
     *
     * @param data the data
     * @param x the value
     * @return the index
     */
    private static int findIndex(double[] data, double x) {
        int low = 0;
        int high = data.length - 1;

        // Bracket so that low is just above the value x
        while (low <= high) {
            final int mid = (low + high) >>> 1;
            final double midVal = data[mid];

            if (x < midVal) {
                // Reduce search range
                high = mid - 1;
            } else {
                // Set data[low] above the value
                low = mid + 1;
            }
        }
        // Verify the index is correct
        Assertions.assertTrue(x < data[low]);
        if (low != 0) {
            Assertions.assertTrue(x >= data[low - 1]);
        }
        return low;
    }
}
