/*
    * 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;
  
  import java.util.Arrays;
  import org.apache.commons.math3.stat.inference.ChiSquareTest;
  import org.apache.commons.rng.UniformRandomProvider;
  import org.apache.commons.rng.core.source64.SplitMix64;
  import org.junit.jupiter.api.Assertions;
  import org.junit.jupiter.api.Test;
  
  /**
   * Test for {@link UnitSphereSampler}.
   */
  class UnitSphereSamplerTest {
      /** 2 pi */
      private static final double TWO_PI = 2 * Math.PI;
  
      /**
       * Test a non-positive dimension.
       */
      @Test
      void testInvalidDimensionThrows() {
          // Use instance constructor not factory constructor to exercise 1.X public API
          Assertions.assertThrows(IllegalArgumentException.class,
              () -> new UnitSphereSampler(0, null));
      }
  
      /**
       * Test a non-positive dimension.
       */
      @Test
      void testInvalidDimensionThrowsWithFactoryConstructor() {
          Assertions.assertThrows(IllegalArgumentException.class,
              () -> UnitSphereSampler.of(null, 0));
      }
  
      /**
       * Test the distribution of points in one dimension.
       */
      @Test
      void testDistribution1D() {
          testDistribution1D(false);
      }
  
      /**
       * Test the distribution of points in one dimension with the factory constructor.
       */
      @Test
      void testDistribution1DWithFactoryConstructor() {
          testDistribution1D(true);
      }
  
      /**
       * Test the distribution of points in one dimension.
       * RNG-130: All samples should be 1 or -1.
       */
      private static void testDistribution1D(boolean factoryConstructor) {
          final UniformRandomProvider rng = RandomAssert.createRNG();
          final UnitSphereSampler generator = createUnitSphereSampler(1, rng, factoryConstructor);
          final int samples = 10000;
          // Count the negatives.
          int count = 0;
          for (int i = 0; i < samples; i++) {
              // Test the deprecated method once in the test suite.
              @SuppressWarnings("deprecation")
              final double[] v = generator.nextVector();
              Assertions.assertEquals(1, v.length);
              final double d = v[0];
              if (d == -1.0) {
                  count++;
              } else if (d != 1.0) {
                  // RNG-130: All samples should be 1 or -1.
                  Assertions.fail("Invalid unit length: " + d);
              }
          }
          // Test the number of negatives is approximately 50%
          assertMonobit(count, samples);
      }
  
      /**
       * Assert that the number of 1 bits is approximately 50%. This is based upon a fixed-step "random
       * walk" of +1/-1 from zero.
       *
       * <p>The test is equivalent to the NIST Monobit test with a fixed p-value of 0.01. The number of
      * bits is recommended to be above 100.</p>
      *
      * @see <a href="https://csrc.nist.gov/publications/detail/sp/800-22/rev-1a/final">Bassham, et al
      *      (2010) NIST SP 800-22: A Statistical Test Suite for Random and Pseudorandom Number
      *      Generators for Cryptographic Applications. Section 2.1.</a>
      *
      * @param bitCount The bit count.
      * @param numberOfBits Number of bits.
      */
     private static void assertMonobit(int bitCount, int numberOfBits) {
         // Convert the bit count into a number of +1/-1 steps.
         final double sum = 2.0 * bitCount - numberOfBits;
         // The reference distribution is Normal with a standard deviation of sqrt(n).
         // Check the absolute position is not too far from the mean of 0 with a fixed
         // p-value of 0.01 taken from a 2-tailed Normal distribution. Computation of
         // the p-value requires the complimentary error function.
         final double absSum = Math.abs(sum);
         final double max = Math.sqrt(numberOfBits) * 2.576;
         Assertions.assertTrue(absSum <= max,
             () -> "Walked too far astray: " + absSum + " > " + max +
                   " (test will fail randomly about 1 in 100 times)");
     }
 
     /**
      * Test the distribution of points in two dimensions.
      */
     @Test
     void testDistribution2D() {
         testDistribution2D(false);
     }
 
     /**
      * Test the distribution of points in two dimensions with the factory constructor.
      */
     @Test
     void testDistribution2DWithFactoryConstructor() {
         testDistribution2D(true);
     }
 
     /**
      * Test the distribution of points in two dimensions.
      * Obtains polar coordinates and checks the angle distribution is uniform.
      */
     private static void testDistribution2D(boolean factoryConstructor) {
         final UniformRandomProvider rng = RandomAssert.createRNG();
         final UnitSphereSampler generator = createUnitSphereSampler(2, rng, factoryConstructor);
 
         // In 2D, angles with a given vector should be uniformly distributed.
         final int angleBins = 200;
         final long[] observed = new long[angleBins];
         final int steps = 100000;
         for (int i = 0; i < steps; ++i) {
             final double[] v = generator.sample();
             Assertions.assertEquals(2, v.length);
             Assertions.assertEquals(1.0, length(v), 1e-10);
             // Get the polar angle bin from xy
             final int angleBin = angleBin(angleBins, v[0], v[1]);
             observed[angleBin]++;
         }
 
         final double[] expected = new double[observed.length];
         Arrays.fill(expected, (double) steps / observed.length);
         final double p = new ChiSquareTest().chiSquareTest(expected, observed);
         Assertions.assertFalse(p < 0.01, () -> "p-value too small: " + p);
     }
 
     /**
      * Test the distribution of points in three dimensions.
      */
     @Test
     void testDistribution3D() {
         testDistribution3D(false);
     }
 
     /**
      * Test the distribution of points in three dimensions with the factory constructor.
      */
     @Test
     void testDistribution3DWithFactoryConstructor() {
         testDistribution3D(true);
     }
 
     /**
      * Test the distribution of points in three dimensions.
      * Obtains spherical coordinates and checks the distribution is uniform.
      */
     private static void testDistribution3D(boolean factoryConstructor) {
         final UniformRandomProvider rng = RandomAssert.createRNG();
         final UnitSphereSampler generator = createUnitSphereSampler(3, rng, factoryConstructor);
 
         // Get 3D spherical coordinates. Assign to a bin.
         //
         // polar angle (theta) in [0, 2pi)
         // azimuthal angle (phi) in [0, pi)
         //
         // theta = arctan(y/x) and is uniformly distributed
         // phi = arccos(z); and cos(phi) is uniformly distributed
         final int angleBins = 20;
         final int depthBins = 10;
         final long[] observed = new long[angleBins * depthBins];
         final int steps = 1000000;
         for (int i = 0; i < steps; ++i) {
             final double[] v = generator.sample();
             Assertions.assertEquals(3, v.length);
             Assertions.assertEquals(1.0, length(v), 1e-10);
             // Get the polar angle bin from xy
             final int angleBin = angleBin(angleBins, v[0], v[1]);
             // Map cos(phi) = z from [-1, 1) to [0, 1) then assign a bin
             final int depthBin = (int) (depthBins * (v[2] + 1) / 2);
             observed[depthBin * angleBins + angleBin]++;
         }
 
         final double[] expected = new double[observed.length];
         Arrays.fill(expected, (double) steps / observed.length);
         final double p = new ChiSquareTest().chiSquareTest(expected, observed);
         Assertions.assertFalse(p < 0.01, () -> "p-value too small: " + p);
     }
 
     /**
      * Test the distribution of points in four dimensions.
      */
     @Test
     void testDistribution4D() {
         testDistribution4D(false);
     }
 
     /**
      * Test the distribution of points in four dimensions with the factory constructor.
      */
     @Test
     void testDistribution4DWithFactoryConstructor() {
         testDistribution4D(true);
     }
 
     /**
      * Test the distribution of points in four dimensions.
      * Checks the surface of the 3-sphere can be used to generate uniform samples within a circle.
      */
     private static void testDistribution4D(boolean factoryConstructor) {
         final UniformRandomProvider rng = RandomAssert.createRNG();
         final UnitSphereSampler generator = createUnitSphereSampler(4, rng, factoryConstructor);
 
         // No uniform distribution of spherical coordinates for a 3-sphere.
         // https://en.wikipedia.org/wiki/N-sphere#Spherical_coordinates
         // Here we exploit the fact that the uniform distribution of a (n+1)-sphere
         // when discarding two coordinates is uniform within a n-ball.
         // Thus any two coordinates of the 4 are uniform within a circle.
         // Here we separately test pairs (0, 1) and (2, 3).
         // Note: We cannot create a single bin from the two bins in each circle as they are
         // not independent. A point close to the edge in one circle requires a point close to
         // the centre in the other circle to create a unit radius from all 4 coordinates.
         // This test exercises the N-dimension sampler and demonstrates the vectors obey
         // properties of the (n+1)-sphere.
 
         // Divide the circle into layers of concentric rings and an angle.
         final int layers = 10;
         final int angleBins = 20;
 
         // Compute the radius for each layer to have the same area
         // (i.e. incrementally larger concentric circles must increase area by a constant).
         // r = sqrt(fraction * maxArea / pi)
         // Unit circle has area pi so we just use sqrt(fraction).
         final double[] r = new double[layers];
         for (int i = 1; i < layers; i++) {
             r[i - 1] = Math.sqrt((double) i / layers);
         }
         // The final radius should be 1.0.
         r[layers - 1] = 1.0;
 
         final long[] observed1 = new long[layers * angleBins];
         final long[] observed2 = new long[observed1.length];
         final int steps = 1000000;
         for (int i = 0; i < steps; ++i) {
             final double[] v = generator.sample();
             Assertions.assertEquals(4, v.length);
             Assertions.assertEquals(1.0, length(v), 1e-10);
             // Circle 1
             observed1[circleBin(angleBins, r, v[0], v[1])]++;
             // Circle 2
             observed2[circleBin(angleBins, r, v[2], v[3])]++;
         }
 
         final double[] expected = new double[observed1.length];
         Arrays.fill(expected, (double) steps / observed1.length);
         final ChiSquareTest chi = new ChiSquareTest();
         final double p1 = chi.chiSquareTest(expected, observed1);
         Assertions.assertFalse(p1 < 0.01, () -> "Circle 1 p-value too small: " + p1);
         final double p2 = chi.chiSquareTest(expected, observed2);
         Assertions.assertFalse(p2 < 0.01, () -> "Circle 2 p-value too small: " + p2);
     }
 
     /**
      * Compute a bin inside the circle using the polar angle theta and the radius thresholds.
      *
      * @param angleBins the number of angle bins
      * @param r the radius bin thresholds (ascending order)
      * @param x the x coordinate
      * @param y the y coordinate
      * @return the bin
      */
     private static int circleBin(int angleBins, double[] r, double x, double y) {
         final int angleBin = angleBin(angleBins, x, y);
         final int radiusBin = radiusBin(r, x, y);
         return radiusBin * angleBins + angleBin;
     }
 
     /**
      * Compute an angle bin from the xy vector. The bin will represent the range [0, 2pi).
      *
      * @param angleBins the number of angle bins
      * @param x the x coordinate
      * @param y the y coordinate
      * @return the bin
      */
     private static int angleBin(int angleBins, double x, double y) {
         final double angle = Math.atan2(y, x);
         // Map [-pi, pi) to [0, 1) then assign a bin
         return (int) (angleBins * (angle + Math.PI) / TWO_PI);
     }
 
     /**
      * Compute a radius bin from the xy vector. The bin is assigned if the length of the vector
      * is above the threshold of the bin.
      *
      * @param r the radius bin thresholds (ascending order)
      * @param x the x coordinate
      * @param y the y coordinate
      * @return the bin
      */
     private static int radiusBin(double[] r, double x, double y) {
         final double length = Math.sqrt(x * x + y * y);
 
         // Note: The bin should be uniformly distributed.
         // A test using a custom binary search (avoiding double NaN checks)
         // shows the simple loop over a small number of bins is comparable in speed.
         // The loop is preferred for simplicity. A binary search may be better
         // if the number of bins is increased.
         for (int layer = 0; layer < r.length; layer++) {
             if (length <= r[layer]) {
                 return layer;
             }
         }
         // Unreachable if the xy component is from a vector of length <= 1
         throw new AssertionError("Invalid sample length: " + length);
     }
 
     /**
      * Test infinite recursion occurs with a bad provider in 2D.
      */
     @Test
     void testBadProvider2D() {
         Assertions.assertThrows(StackOverflowError.class,
             () -> testBadProvider(2));
     }
 
     /**
      * Test infinite recursion occurs with a bad provider in 3D.
      */
     @Test
     void testBadProvider3D() {
         Assertions.assertThrows(StackOverflowError.class,
             () -> testBadProvider(3));
     }
 
     /**
      * Test infinite recursion occurs with a bad provider in 4D.
      */
     @Test
     void testBadProvider4D() {
         Assertions.assertThrows(StackOverflowError.class,
             () -> testBadProvider(4));
     }
 
     /**
      * Test the edge case where the normalisation sum to divide by is always zero.
      * This test requires generation of Gaussian samples with the value 0.
      * The sample eventually fails due to infinite recursion.
      * See RNG-55.
      *
      * @param dimension the dimension
      */
     private static void testBadProvider(final int dimension) {
         // A provider that will create zero valued Gaussian samples
         // from the ZigguratNormalizedGaussianSampler.
         final UniformRandomProvider bad = new SplitMix64(0L) {
             @Override
             public long nextLong() {
                 return 0;
             }
         };
 
         UnitSphereSampler.of(bad, dimension).sample();
     }
 
     /**
      * Test the edge case where the normalisation sum to divide by is zero for 2D.
      */
     @Test
     void testInvalidInverseNormalisation2D() {
         testInvalidInverseNormalisationND(2);
     }
 
     /**
      * 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.
      * See RNG-55.
      */
     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 = UnitSphereSampler.of(bad, dimension).sample();
         Assertions.assertEquals(dimension, vector.length);
         Assertions.assertEquals(1.0, length(vector), 1e-10);
     }
 
     /**
      * Test to demonstrate that using floating-point equality of the norm squared with
      * zero is valid. Any norm squared after zero should produce a valid scaling factor.
      */
     @Test
     void testNextNormSquaredAfterZeroIsValid() {
         // The sampler explicitly handles length == 0 using recursion.
         // Anything above zero should be valid.
         final double normSq = Math.nextUp(0.0);
         // Map to the scaling factor
         final double f = 1 / Math.sqrt(normSq);
         // As long as this is finite positive then the sampler is valid
         Assertions.assertTrue(f > 0 && f <= Double.MAX_VALUE);
     }
 
     /**
      * Test the SharedStateSampler implementation for 1D.
      */
     @Test
     void testSharedStateSampler1D() {
         testSharedStateSampler(1, false);
     }
 
     /**
      * Test the SharedStateSampler implementation for 2D.
      */
     @Test
     void testSharedStateSampler2D() {
         testSharedStateSampler(2, false);
     }
 
     /**
      * Test the SharedStateSampler implementation for 3D.
      */
     @Test
     void testSharedStateSampler3D() {
         testSharedStateSampler(3, false);
     }
 
     /**
      * Test the SharedStateSampler implementation for 4D.
      */
     @Test
     void testSharedStateSampler4D() {
         testSharedStateSampler(4, false);
     }
 
     /**
      * Test the SharedStateSampler implementation for 1D using the factory constructor.
      */
     @Test
     void testSharedStateSampler1DWithFactoryConstructor() {
         testSharedStateSampler(1, true);
     }
 
     /**
      * Test the SharedStateSampler implementation for 2D using the factory constructor.
      */
     @Test
     void testSharedStateSampler2DWithFactoryConstructor() {
         testSharedStateSampler(2, true);
     }
 
     /**
      * Test the SharedStateSampler implementation for 3D using the factory constructor.
      */
     @Test
     void testSharedStateSampler3DWithFactoryConstructor() {
         testSharedStateSampler(3, true);
     }
 
     /**
      * Test the SharedStateSampler implementation for 4D using the factory constructor.
      */
     @Test
     void testSharedStateSampler4DWithFactoryConstructor() {
         testSharedStateSampler(4, true);
     }
 
     /**
      * Test the SharedStateSampler implementation for the given dimension.
      *
      * @param dimension the dimension
      * @param factoryConstructor true to use the factory constructor
      */
     private static void testSharedStateSampler(int dimension, boolean factoryConstructor) {
         final UniformRandomProvider rng1 = RandomAssert.seededRNG();
         final UniformRandomProvider rng2 = RandomAssert.seededRNG();
         final UnitSphereSampler sampler1 = createUnitSphereSampler(dimension, rng1, factoryConstructor);
         final UnitSphereSampler sampler2 = sampler1.withUniformRandomProvider(rng2);
         RandomAssert.assertProduceSameSequence(sampler1, sampler2);
     }
 
     /**
      * Creates a UnitSphereSampler.
      *
      * @param dimension the dimension
      * @param rng the source of randomness
      * @param factoryConstructor true to use the factory constructor
      * @return the sampler
      */
     private static UnitSphereSampler createUnitSphereSampler(int dimension, UniformRandomProvider rng,
             boolean factoryConstructor) {
         return factoryConstructor ?
                 UnitSphereSampler.of(rng, dimension) : new UnitSphereSampler(dimension, rng);
     }
 
     /**
      * @return the length (L2-norm) of given vector.
      */
     private static double length(double[] vector) {
         double total = 0;
         for (final double d : vector) {
             total += d * d;
         }
         return Math.sqrt(total);
     }
 }