/*
    * 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.math3.util.CombinatoricsUtils;
  import org.apache.commons.rng.UniformRandomProvider;
  import org.junit.jupiter.api.Assertions;
  import org.junit.jupiter.api.Test;
  
  /**
   * Tests for {@link CombinationSampler}.
   */
  class CombinationSamplerTest {
      @Test
      void testSampleIsInDomain() {
          final UniformRandomProvider rng = RandomAssert.seededRNG();
          final int n = 6;
          for (int k = 1; k <= n; k++) {
              final CombinationSampler sampler = new CombinationSampler(rng, n, k);
              final int[] random = sampler.sample();
              for (final int s : random) {
                  assertIsInDomain(n, s);
              }
          }
      }
  
      @Test
      void testUniformWithKlessThanHalfN() {
          final int n = 8;
          final int k = 2;
          assertUniformSamples(n, k);
      }
  
      @Test
      void testUniformWithKmoreThanHalfN() {
          final int n = 8;
          final int k = 6;
          assertUniformSamples(n, k);
      }
  
      @Test
      void testSampleWhenNequalsKIsNotShuffled() {
          final UniformRandomProvider rng = RandomAssert.seededRNG();
          // Check n == k boundary case.
          // This is allowed but the sample is not shuffled.
          for (int n = 1; n < 3; n++) {
              final int k = n;
              final CombinationSampler sampler = new CombinationSampler(rng, n, k);
              final int[] sample = sampler.sample();
              Assertions.assertEquals(n, sample.length, "Incorrect sample length");
              for (int i = 0; i < n; i++) {
                  Assertions.assertEquals(i, sample[i], "Sample was shuffled");
              }
          }
      }
  
      @Test
      void testKgreaterThanNThrows() {
          final UniformRandomProvider rng = RandomAssert.seededRNG();
          // Must fail for k > n.
          final int n = 2;
          final int k = 3;
          Assertions.assertThrows(IllegalArgumentException.class,
              () -> new CombinationSampler(rng, n, k));
      }
  
      @Test
      void testNequalsZeroThrows() {
          final UniformRandomProvider rng = RandomAssert.seededRNG();
          // Must fail for n = 0.
          final int n = 0;
          final int k = 3;
          Assertions.assertThrows(IllegalArgumentException.class,
              () -> new CombinationSampler(rng, n, k));
      }
  
      @Test
      void testKequalsZeroThrows() {
          final UniformRandomProvider rng = RandomAssert.seededRNG();
          // Must fail for k = 0.
          final int n = 2;
          final int k = 0;
          Assertions.assertThrows(IllegalArgumentException.class,
             () -> new CombinationSampler(rng, n, k));
     }
 
     @Test
     void testNisNegativeThrows() {
         final UniformRandomProvider rng = RandomAssert.seededRNG();
         // Must fail for n <= 0.
         final int n = -1;
         final int k = 3;
         Assertions.assertThrows(IllegalArgumentException.class,
             () -> new CombinationSampler(rng, n, k));
     }
 
     @Test
     void testKisNegativeThrows() {
         final UniformRandomProvider rng = RandomAssert.seededRNG();
         // Must fail for k <= 0.
         final int n = 0;
         final int k = -1;
         Assertions.assertThrows(IllegalArgumentException.class,
             () -> new CombinationSampler(rng, n, k));
     }
 
     /**
      * Test the SharedStateSampler implementation.
      */
     @Test
     void testSharedStateSampler() {
         final UniformRandomProvider rng1 = RandomAssert.seededRNG();
         final UniformRandomProvider rng2 = RandomAssert.seededRNG();
         final int n = 17;
         final int k = 3;
         final CombinationSampler sampler1 =
             new CombinationSampler(rng1, n, k);
         final CombinationSampler sampler2 = sampler1.withUniformRandomProvider(rng2);
         RandomAssert.assertProduceSameSequence(sampler1, sampler2);
     }
 
     //// Support methods.
 
     /**
      * Asserts the sample value is in the range 0 to n-1.
      *
      * @param n     the n
      * @param value the sample value
      */
     private static void assertIsInDomain(int n, int value) {
         if (value < 0 || value >= n) {
             Assertions.fail("sample " + value + " not in the domain " + n);
         }
     }
 
     private void assertUniformSamples(int n, int k) {
         // The C(n, k) should generate a sample of unspecified order.
         // To test this each combination is allocated a unique code
         // based on setting k of the first n-bits in an integer.
         // Codes are positive for all combinations of bits that use k-bits,
         // otherwise they are negative.
         final int totalBitCombinations = 1 << n;
         final int[] codeLookup = new int[totalBitCombinations];
         Arrays.fill(codeLookup, -1); // initialize as negative
         int codes = 0;
         for (int i = 0; i < totalBitCombinations; i++) {
             if (Integer.bitCount(i) == k) {
                 // This is a valid sample so allocate a code
                 codeLookup[i] = codes++;
             }
         }
 
         // The number of combinations C(n, k) is the binomial coefficient
         Assertions.assertEquals(CombinatoricsUtils.binomialCoefficient(n, k), codes,
             "Incorrect number of combination codes");
 
         final long[] observed = new long[codes];
         final int numSamples = 6000;
 
         final UniformRandomProvider rng = RandomAssert.createRNG();
         final CombinationSampler sampler = new CombinationSampler(rng, n, k);
         for (int i = 0; i < numSamples; i++) {
             observed[findCode(codeLookup, sampler.sample())]++;
         }
 
         // Chi squared test of uniformity
         final double numExpected = numSamples / (double) codes;
         final double[] expected = new double[codes];
         Arrays.fill(expected, numExpected);
         final ChiSquareTest chiSquareTest = new ChiSquareTest();
         // Pass if we cannot reject null hypothesis that distributions are the same.
         Assertions.assertFalse(chiSquareTest.chiSquareTest(expected, observed, 0.001));
     }
 
     private static int findCode(int[] codeLookup, int[] sample) {
         // Each sample index is used to set a bit in an integer.
         // The resulting bits should be a valid code.
         int bits = 0;
         for (final int s : sample) {
             // This shift will be from 0 to n-1 since it is from the
             // domain of size n.
             bits |= 1 << s;
         }
         if (bits >= codeLookup.length) {
             Assertions.fail("Bad bit combination: " + Arrays.toString(sample));
         }
         final int code = codeLookup[bits];
         if (code < 0) {
             Assertions.fail("Bad bit code: " + Arrays.toString(sample));
         }
         return code;
     }
 }