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

import java.util.Arrays;
import java.util.function.DoubleUnaryOperator;
import org.apache.commons.math3.stat.inference.ChiSquareTest;
import org.apache.commons.rng.UniformRandomProvider;
import org.apache.commons.rng.core.source64.SplitMix64;
import org.apache.commons.rng.sampling.RandomAssert;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.Test;

/**
 * Test for {@link UnitBallSampler}.
 */
class UnitBallSamplerTest {
    /**
     * Test a non-positive dimension.
     */
    @Test
    void testInvalidDimensionThrows() {
        final UniformRandomProvider rng = RandomAssert.seededRNG();
        Assertions.assertThrows(IllegalArgumentException.class,
            () -> UnitBallSampler.of(rng, 0));
    }

    /**
     * Test the distribution of points in one dimension.
     */
    @Test
    void testDistribution1D() {
        testDistributionND(1);
    }

    /**
     * Test the distribution of points in two dimensions.
     */
    @Test
    void testDistribution2D() {
        testDistributionND(2);
    }

    /**
     * Test the distribution of points in three dimensions.
     */
    @Test
    void testDistribution3D() {
        testDistributionND(3);
    }

    /**
     * Test the distribution of points in four dimensions.
     */
    @Test
    void testDistribution4D() {
        testDistributionND(4);
    }

    /**
     * Test the distribution of points in five dimensions.
     */
    @Test
    void testDistribution5D() {
        testDistributionND(5);
    }

    /**
     * Test the distribution of points in six dimensions.
     */
    @Test
    void testDistribution6D() {
        testDistributionND(6);
    }

    /**
     * Test the distribution of points in n dimensions. The output coordinates
     * should be uniform in the unit n-ball. The unit n-ball is divided into inner
     * n-balls. The radii of the internal n-balls are varied to ensure that successive layers
     * have the same volume. This assigns each coordinate to an inner n-ball layer and an
     * orthant using the sign bits of the coordinates. The number of samples in each bin
     * should be the same.
     *
     * @see <a href="https://en.wikipedia.org/wiki/Volume_of_an_n-ball">Volume of an n-ball</a>
     * @see <a href="https://en.wikipedia.org/wiki/Orthant">Orthant</a>
     */
    private static void testDistributionND(int dimension) {
        // The number of inner layers and samples has been selected by trial and error using
        // random seeds and multiple runs to ensure correctness of the test (i.e. it fails with
        // approximately the fraction expected for the test p-value).
        // A fixed seed is used to make the test suite robust.
        final int layers = 10;
        final int samplesPerBin = 20;
        final int orthants = 1 << dimension;

        // Compute the radius for each layer to have the same volume.
        final double volume = createVolumeFunction(dimension).applyAsDouble(1);
        final DoubleUnaryOperator radius = createRadiusFunction(dimension);
        final double[] r = new double[layers];
        for (int i = 1; i < layers; i++) {
            r[i - 1] = radius.applyAsDouble(volume * ((double) i / layers));
        }
        // The final radius should be 1.0. Any coordinates with a radius above 1
        // should fail so explicitly set the value as 1.
        r[layers - 1] = 1.0;

        // Expect a uniform distribution
        final double[] expected = new double[layers * orthants];
        final int samples = samplesPerBin * expected.length;
        Arrays.fill(expected, (double) samples / layers);

        // Increase the loops to verify robustness
        final UniformRandomProvider rng = RandomAssert.createRNG();
        final UnitBallSampler sampler = UnitBallSampler.of(rng, dimension);
        for (int loop = 0; loop < 1; loop++) {
            // Assign each coordinate to a layer inside the ball and an orthant using the sign
            final long[] observed = new long[layers * orthants];
            NEXT:
            for (int i = 0; i < samples; i++) {
                final double[] v = sampler.sample();
                final double length = length(v);
                for (int layer = 0; layer < layers; layer++) {
                    if (length <= r[layer]) {
                        final int orthant = orthant(v);
                        observed[layer * orthants + orthant]++;
                        continue NEXT;
                    }
                }
                // Radius above 1
                Assertions.fail("Invalid sample length: " + length);
            }
            final double p = new ChiSquareTest().chiSquareTest(expected, observed);
            Assertions.assertFalse(p < 0.001, () -> "p-value too small: " + p);
        }
    }

    /**
     * Test the edge case where the normalisation sum to divide by is zero for 3D.
     */
    @Test
    void testInvalidInverseNormalisation3D() {
        testInvalidInverseNormalisationND(3);
    }

    /**
     * Test the edge case where the normalisation sum to divide by is zero for 4D.
     */
    @Test
    void testInvalidInverseNormalisation4D() {
        testInvalidInverseNormalisationND(4);
    }

    /**
     * Test the edge case where the normalisation sum to divide by is zero.
     * This test requires generation of Gaussian samples with the value 0.
     */
    private static void testInvalidInverseNormalisationND(final int dimension) {
        // Create a provider that will create a bad first sample but then recover.
        // This checks recursion will return a good value.
        final UniformRandomProvider bad = new SplitMix64(0x1a2b3cL) {
            private int count = -2 * dimension;

            @Override
            public long nextLong() {
                // Return enough zeros to create Gaussian samples of zero for all coordinates.
                return count++ < 0 ? 0 : super.nextLong();
            }
        };

        final double[] vector = UnitBallSampler.of(bad, dimension).sample();
        Assertions.assertEquals(dimension, vector.length);
        // A non-zero coordinate should occur with a SplitMix which returns 0 only once.
        Assertions.assertNotEquals(0.0, length(vector));
    }

    /**
     * Test the SharedStateSampler implementation for 1D.
     */
    @Test
    void testSharedStateSampler1D() {
        testSharedStateSampler(1);
    }

    /**
     * Test the SharedStateSampler implementation for 2D.
     */
    @Test
    void testSharedStateSampler2D() {
        testSharedStateSampler(2);
    }

    /**
     * Test the SharedStateSampler implementation for 3D.
     */
    @Test
    void testSharedStateSampler3D() {
        testSharedStateSampler(3);
    }

    /**
     * Test the SharedStateSampler implementation for 4D.
     */
    @Test
    void testSharedStateSampler4D() {
        testSharedStateSampler(4);
    }

    /**
     * Test the SharedStateSampler implementation for the given dimension.
     */
    private static void testSharedStateSampler(int dimension) {
        final UniformRandomProvider rng1 = RandomAssert.seededRNG();
        final UniformRandomProvider rng2 = RandomAssert.seededRNG();
        final UnitBallSampler sampler1 = UnitBallSampler.of(rng1, dimension);
        final UnitBallSampler sampler2 = sampler1.withUniformRandomProvider(rng2);
        RandomAssert.assertProduceSameSequence(sampler1, sampler2);
    }

    /**
     * @return the length (L2-norm) of given vector.
     */
    private static double length(double[] vector) {
        double total = 0;
        for (double d : vector) {
            total += d * d;
        }
        return Math.sqrt(total);
    }

    /**
     * Assign an orthant to the vector using the sign of each component.
     * The i<sup>th</sup> bit is set in the orthant for the i<sup>th</sup> component
     * if the component is negative.
     *
     * @return the orthant in the range [0, vector.length)
     * @see <a href="https://en.wikipedia.org/wiki/Orthant">Orthant</a>
     */
    private static int orthant(double[] vector) {
        int orthant = 0;
        for (int i = 0; i < vector.length; i++) {
            if (vector[i] < 0) {
                orthant |= 1 << i;
            }
        }
        return orthant;
    }

    /**
     * Check the n-ball volume functions can map the radius to the volume and back.
     * These functions are used to divide the n-ball into uniform volume bins to test sampling
     * within the n-ball.
     */
    @Test
    void checkVolumeFunctions() {
        final double[] radii = {0, 0.1, 0.25, 0.5, 0.75, 1.0};
        for (int n = 1; n <= 6; n++) {
            final DoubleUnaryOperator volume = createVolumeFunction(n);
            final DoubleUnaryOperator radius = createRadiusFunction(n);
            for (final double r : radii) {
                Assertions.assertEquals(r, radius.applyAsDouble(volume.applyAsDouble(r)), 1e-10);
            }
        }
    }

    /**
     * Creates a function to compute the volume of a ball of the given dimension
     * from the radius.
     *
     * @param dimension the dimension
     * @return the volume function
     * @see <a href="https://en.wikipedia.org/wiki/Volume_of_an_n-ball">Volume of an n-ball</a>
     */
    private static DoubleUnaryOperator createVolumeFunction(final int dimension) {
        if (dimension == 1) {
            return r -> r * 2;
        } else if (dimension == 2) {
            return r -> Math.PI * r * r;
        } else if (dimension == 3) {
            final double factor = 4 * Math.PI / 3;
            return r -> factor * Math.pow(r, 3);
        } else if (dimension == 4) {
            final double factor = Math.PI * Math.PI / 2;
            return r -> factor * Math.pow(r, 4);
        } else if (dimension == 5) {
            final double factor = 8 * Math.PI * Math.PI / 15;
            return r -> factor * Math.pow(r, 5);
        } else if (dimension == 6) {
            final double factor = Math.pow(Math.PI, 3) / 6;
            return r -> factor * Math.pow(r, 6);
        }
        throw new IllegalStateException("Unsupported dimension: " + dimension);
    }

    /**
     * Creates a function to compute the radius of a ball of the given dimension
     * from the volume.
     *
     * @param dimension the dimension
     * @return the radius function
     * @see <a href="https://en.wikipedia.org/wiki/Volume_of_an_n-ball">Volume of an n-ball</a>
     */
    private static DoubleUnaryOperator createRadiusFunction(final int dimension) {
        if (dimension == 1) {
            return v -> v * 0.5;
        } else if (dimension == 2) {
            return v -> Math.sqrt(v / Math.PI);
        } else if (dimension == 3) {
            final double factor = 3.0 / (4 * Math.PI);
            return v -> Math.cbrt(v * factor);
        } else if (dimension == 4) {
            final double factor = 2.0 / (Math.PI * Math.PI);
            return v -> Math.pow(v * factor, 0.25);
        } else if (dimension == 5) {
            final double factor = 15.0 / (8 * Math.PI * Math.PI);
            return v -> Math.pow(v * factor, 0.2);
        } else if (dimension == 6) {
            final double factor = 6.0 / Math.pow(Math.PI, 3);
            return v -> Math.pow(v * factor, 1.0 / 6);
        }
        throw new IllegalStateException("Unsupported dimension: " + dimension);
    }
}
