/*
    * 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 org.apache.commons.math3.distribution.BinomialDistribution;
  import org.apache.commons.math3.distribution.PoissonDistribution;
  import org.apache.commons.math3.stat.inference.ChiSquareTest;
  import org.apache.commons.rng.UniformRandomProvider;
  import org.apache.commons.rng.sampling.RandomAssert;
  import org.junit.jupiter.api.Assertions;
  import org.junit.jupiter.api.Test;
  import org.junit.jupiter.params.ParameterizedTest;
  import org.junit.jupiter.params.provider.ValueSource;
  
  /**
   * Test for the {@link GuideTableDiscreteSampler}.
   */
  class GuideTableDiscreteSamplerTest {
      @Test
      void testConstructorThrowsWithNullProbabilites() {
          assertConstructorThrows(null, 1.0);
      }
  
      @Test
      void testConstructorThrowsWithZeroLengthProbabilites() {
          assertConstructorThrows(new double[0], 1.0);
      }
  
      @Test
      void testConstructorThrowsWithNegativeProbabilites() {
          assertConstructorThrows(new double[] {-1, 0.1, 0.2}, 1.0);
      }
  
      @Test
      void testConstructorThrowsWithNaNProbabilites() {
          assertConstructorThrows(new double[] {0.1, Double.NaN, 0.2}, 1.0);
      }
  
      @Test
      void testConstructorThrowsWithInfiniteProbabilites() {
          assertConstructorThrows(new double[] {0.1, Double.POSITIVE_INFINITY, 0.2}, 1.0);
      }
  
      @Test
      void testConstructorThrowsWithInfiniteSumProbabilites() {
          assertConstructorThrows(new double[] {Double.MAX_VALUE, Double.MAX_VALUE}, 1.0);
      }
  
      @Test
      void testConstructorThrowsWithZeroSumProbabilites() {
          assertConstructorThrows(new double[4], 1.0);
      }
  
      @Test
      void testConstructorThrowsWithZeroAlpha() {
          assertConstructorThrows(new double[] {0.5, 0.5}, 0.0);
      }
  
      @Test
      void testConstructorThrowsWithNegativeAlpha() {
          assertConstructorThrows(new double[] {0.5, 0.5}, -1.0);
      }
  
      /**
       * Assert the factory constructor throws an {@link IllegalArgumentException}.
       *
       * @param probabilities the probabilities
       * @param alpha the alpha
       */
      private static void assertConstructorThrows(double[] probabilities, double alpha) {
          final UniformRandomProvider rng = RandomAssert.seededRNG();
          Assertions.assertThrows(IllegalArgumentException.class,
              () -> GuideTableDiscreteSampler.of(rng, probabilities, alpha));
      }
  
      @Test
      void testToString() {
          final UniformRandomProvider rng = RandomAssert.createRNG();
          final SharedStateDiscreteSampler sampler = GuideTableDiscreteSampler.of(rng, new double[] {0.5, 0.5}, 1.0);
          Assertions.assertTrue(sampler.toString().toLowerCase().contains("guide table"));
      }
  
      /**
       * Test sampling from a binomial distribution.
       */
     @Test
     void testBinomialSamples() {
         final int trials = 67;
         final double probabilityOfSuccess = 0.345;
         final BinomialDistribution dist = new BinomialDistribution(null, trials, probabilityOfSuccess);
         final double[] expected = new double[trials + 1];
         for (int i = 0; i < expected.length; i++) {
             expected[i] = dist.probability(i);
         }
         checkSamples(expected, 1.0);
     }
 
     /**
      * Test sampling from a Poisson distribution.
      */
     @Test
     void testPoissonSamples() {
         final double mean = 3.14;
         final PoissonDistribution dist = new PoissonDistribution(null, mean,
             PoissonDistribution.DEFAULT_EPSILON, PoissonDistribution.DEFAULT_MAX_ITERATIONS);
         final int maxN = dist.inverseCumulativeProbability(1 - 1e-6);
         final double[] expected = new double[maxN];
         for (int i = 0; i < expected.length; i++) {
             expected[i] = dist.probability(i);
         }
         checkSamples(expected, 1.0);
     }
 
     /**
      * Test sampling from a non-uniform distribution of probabilities (these sum to 1).
      * The alpha used in the default (1.0) or a smaller or larger value than the default.
      */
     @ParameterizedTest
     @ValueSource(doubles = {1.0, 0.1, 10.0})
     void testNonUniformSamplesWithProbabilities(double alpha) {
         final double[] expected = {0.1, 0.2, 0.3, 0.1, 0.3};
         checkSamples(expected, alpha);
     }
 
     /**
      * Test sampling from a non-uniform distribution of observations (i.e. the sum is not 1 as per
      * probabilities).
      */
     @Test
     void testNonUniformSamplesWithObservations() {
         final double[] expected = {1, 2, 3, 1, 3};
         checkSamples(expected, 1.0);
     }
 
     /**
      * Test sampling from a non-uniform distribution of probabilities (these sum to 1).
      * Extra zero-values are added.
      */
     @Test
     void testNonUniformSamplesWithZeroProbabilities() {
         final double[] expected = {0.1, 0, 0.2, 0.3, 0.1, 0.3, 0};
         checkSamples(expected, 1.0);
     }
 
     /**
      * Test sampling from a non-uniform distribution of observations (i.e. the sum is not 1 as per
      * probabilities). Extra zero-values are added.
      */
     @Test
     void testNonUniformSamplesWithZeroObservations() {
         final double[] expected = {1, 2, 3, 0, 1, 3, 0};
         checkSamples(expected, 1.0);
     }
 
     /**
      * Test sampling from a uniform distribution. This is an edge case where there
      * are no probabilities less than the mean.
      */
     @Test
     void testUniformSamplesWithNoObservationLessThanTheMean() {
         final double[] expected = {2, 2, 2, 2, 2, 2};
         checkSamples(expected, 1.0);
     }
 
     /**
      * Check the distribution of samples match the expected probabilities.
      *
      * <p>If the expected probability is zero then this should never be sampled. The non-zero
      * probabilities are compared to the sample distribution using a Chi-square test.</p>
      *
      * @param probabilies the probabilities
      * @param alpha the alpha
      */
     private static void checkSamples(double[] probabilies, double alpha) {
         final UniformRandomProvider rng = RandomAssert.createRNG();
         final SharedStateDiscreteSampler sampler = GuideTableDiscreteSampler.of(rng, probabilies, alpha);
 
         final int numberOfSamples = 10000;
         final long[] samples = new long[probabilies.length];
         for (int i = 0; i < numberOfSamples; i++) {
             samples[sampler.sample()]++;
         }
 
         // Handle a test with some zero-probability observations by mapping them out.
         // The results is the Chi-square test is performed using only the non-zero probabilities.
         int mapSize = 0;
         for (int i = 0; i < probabilies.length; i++) {
             if (probabilies[i] != 0) {
                 mapSize++;
             }
         }
 
         final double[] expected = new double[mapSize];
         final long[] observed = new long[mapSize];
         for (int i = 0; i < probabilies.length; i++) {
             if (probabilies[i] == 0) {
                 Assertions.assertEquals(0, samples[i], "No samples expected from zero probability");
             } else {
                 // This can be added for the Chi-square test
                 --mapSize;
                 expected[mapSize] = probabilies[i];
                 observed[mapSize] = samples[i];
             }
         }
 
         final ChiSquareTest chiSquareTest = new ChiSquareTest();
         // Pass if we cannot reject null hypothesis that the distributions are the same.
         Assertions.assertFalse(chiSquareTest.chiSquareTest(expected, observed, 0.001));
     }
 
     /**
      * Test the SharedStateSampler implementation.
      */
     @Test
     void testSharedStateSampler() {
         final UniformRandomProvider rng1 = RandomAssert.seededRNG();
         final UniformRandomProvider rng2 = RandomAssert.seededRNG();
         final double[] probabilities = {0.1, 0, 0.2, 0.3, 0.1, 0.3, 0};
         final SharedStateDiscreteSampler sampler1 =
             GuideTableDiscreteSampler.of(rng1, probabilities);
         final SharedStateDiscreteSampler sampler2 = sampler1.withUniformRandomProvider(rng2);
         RandomAssert.assertProduceSameSequence(sampler1, sampler2);
     }
 }