/*
    * 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 java.util.Arrays;
  import org.apache.commons.math3.linear.Array2DRowRealMatrix;
  import org.apache.commons.math3.stat.correlation.Covariance;
  import org.apache.commons.math3.stat.descriptive.moment.Mean;
  import org.apache.commons.rng.UniformRandomProvider;
  import org.apache.commons.rng.core.source64.SplitMix64;
  import org.apache.commons.rng.sampling.RandomAssert;
  import org.apache.commons.rng.simple.RandomSource;
  import org.junit.jupiter.api.Assertions;
  import org.junit.jupiter.api.Test;
  
  /**
   * Test for {@link DirichletSampler}.
   */
  class DirichletSamplerTest {
      @Test
      void testDistributionThrowsWithInvalidNumberOfCategories() {
          final UniformRandomProvider rng = RandomAssert.seededRNG();
          Assertions.assertThrows(IllegalArgumentException.class,
              () -> DirichletSampler.of(rng, 1.0));
      }
  
      @Test
      void testDistributionThrowsWithZeroConcentration() {
          final UniformRandomProvider rng = RandomAssert.seededRNG();
          Assertions.assertThrows(IllegalArgumentException.class,
              () -> DirichletSampler.of(rng, 1.0, 0.0));
      }
  
      @Test
      void testDistributionThrowsWithNaNConcentration() {
          final UniformRandomProvider rng = RandomAssert.seededRNG();
          Assertions.assertThrows(IllegalArgumentException.class,
              () -> DirichletSampler.of(rng, 1.0, Double.NaN));
      }
  
      @Test
      void testDistributionThrowsWithInfiniteConcentration() {
          final UniformRandomProvider rng = RandomAssert.seededRNG();
          Assertions.assertThrows(IllegalArgumentException.class,
              () -> DirichletSampler.of(rng, 1.0, Double.POSITIVE_INFINITY));
      }
  
      @Test
      void testSymmetricDistributionThrowsWithInvalidNumberOfCategories() {
          final UniformRandomProvider rng = RandomAssert.seededRNG();
          Assertions.assertThrows(IllegalArgumentException.class,
                  () -> DirichletSampler.symmetric(rng, 1, 1.0));
      }
  
      @Test
      void testSymmetricDistributionThrowsWithZeroConcentration() {
          final UniformRandomProvider rng = RandomAssert.seededRNG();
          Assertions.assertThrows(IllegalArgumentException.class,
              () -> DirichletSampler.symmetric(rng, 2, 0.0));
      }
  
      @Test
      void testSymmetricDistributionThrowsWithNaNConcentration() {
          final UniformRandomProvider rng = RandomAssert.seededRNG();
          Assertions.assertThrows(IllegalArgumentException.class,
              () -> DirichletSampler.symmetric(rng, 2, Double.NaN));
      }
  
      @Test
      void testSymmetricDistributionThrowsWithInfiniteConcentration() {
          final UniformRandomProvider rng = RandomAssert.seededRNG();
          Assertions.assertThrows(IllegalArgumentException.class,
              () -> DirichletSampler.symmetric(rng, 2, Double.POSITIVE_INFINITY));
      }
  
      /**
       * Create condition so that all samples are zero and it is impossible to normalise the
       * samples to sum to 1. These should be ignored and the sample is repeated until
       * normalisation is possible.
       */
      @Test
      void testInvalidSampleIsIgnored() {
          // An RNG implementation which should create zero samples from the underlying
          // exponential sampler for an initial sequence.
          final UniformRandomProvider rng = new SplitMix64(0L) {
             private int i;
 
             @Override
             public long next() {
                 return i++ < 10 ? 0L : super.next();
             }
         };
 
         // Alpha=1 will use an exponential sampler
         final DirichletSampler sampler = DirichletSampler.symmetric(rng, 2, 1.0);
         assertSample(2, sampler.sample());
     }
 
     @Test
     void testSharedStateSampler() {
         final RandomSource randomSource = RandomSource.XO_RO_SHI_RO_128_PP;
         final byte[] seed = randomSource.createSeed();
         final UniformRandomProvider rng1 = randomSource.create(seed);
         final UniformRandomProvider rng2 = randomSource.create(seed);
         final DirichletSampler sampler1 = DirichletSampler.of(rng1, 1, 2, 3);
         final DirichletSampler sampler2 = sampler1.withUniformRandomProvider(rng2);
         RandomAssert.assertProduceSameSequence(sampler1, sampler2);
     }
 
     @Test
     void testSharedStateSamplerForSymmetricCase() {
         final RandomSource randomSource = RandomSource.XO_RO_SHI_RO_128_PP;
         final byte[] seed = randomSource.createSeed();
         final UniformRandomProvider rng1 = randomSource.create(seed);
         final UniformRandomProvider rng2 = randomSource.create(seed);
         final DirichletSampler sampler1 = DirichletSampler.symmetric(rng1, 2, 1.5);
         final DirichletSampler sampler2 = sampler1.withUniformRandomProvider(rng2);
         RandomAssert.assertProduceSameSequence(sampler1, sampler2);
     }
 
     @Test
     void testSymmetricCaseMatchesGeneralCase() {
         final RandomSource randomSource = RandomSource.XO_RO_SHI_RO_128_PP;
         final byte[] seed = randomSource.createSeed();
         final UniformRandomProvider rng1 = randomSource.create(seed);
         final UniformRandomProvider rng2 = randomSource.create(seed);
         final int k = 3;
         final double[] alphas = new double[k];
         for (final double alpha : new double[] {0.5, 1.0, 1.5}) {
             Arrays.fill(alphas, alpha);
             final DirichletSampler sampler1 = DirichletSampler.symmetric(rng1, k, alpha);
             final DirichletSampler sampler2 = DirichletSampler.of(rng2, alphas);
             RandomAssert.assertProduceSameSequence(sampler1, sampler2);
         }
     }
 
     /**
      * Test the toString method. This is added to ensure coverage.
      */
     @Test
     void testToString() {
         final UniformRandomProvider rng = RandomAssert.seededRNG();
         final DirichletSampler sampler1 = DirichletSampler.symmetric(rng, 2, 1.0);
         final DirichletSampler sampler2 = DirichletSampler.of(rng, 0.5, 1, 1.5);
         Assertions.assertTrue(sampler1.toString().toLowerCase().contains("dirichlet"));
         Assertions.assertTrue(sampler2.toString().toLowerCase().contains("dirichlet"));
     }
 
     @Test
     void testSampling1() {
         assertSamples(1, 2, 3);
     }
 
     @Test
     void testSampling2() {
         assertSamples(1, 1, 1);
     }
 
     @Test
     void testSampling3() {
         assertSamples(0.5, 1, 1.5);
     }
 
     @Test
     void testSampling4() {
         assertSamples(1, 3);
     }
 
     @Test
     void testSampling5() {
         assertSamples(1, 2, 3, 4);
     }
 
     /**
      * Assert samples from the distribution. The variates are tested against the expected
      * mean and covariance for the given concentration parameters.
      *
      * @param alpha Concentration parameters.
      */
     private static void assertSamples(double... alpha) {
         // No fixed seed. Failed tests will be repeated by the JUnit test runner.
         final UniformRandomProvider rng = RandomAssert.createRNG();
         final DirichletSampler sampler = DirichletSampler.of(rng, alpha);
         final int k = alpha.length;
         final double[][] samples = new double[100000][];
         for (int i = 0; i < samples.length; i++) {
             final double[] x = sampler.sample();
             assertSample(k, x);
             samples[i] = x;
         }
 
         // Computation of moments:
         // https://en.wikipedia.org/wiki/Dirichlet_distribution#Moments
 
         // Compute the sum of the concentration parameters: alpha0
         double alpha0 = 0;
         for (int i = 0; i < k; i++) {
             alpha0 += alpha[i];
         }
 
         // Use a moderate tolerance.
         // Differences are usually observed in the 3rd significant figure.
         final double relativeTolerance = 5e-2;
 
         // Mean = alpha[i] / alpha0
         final double[] means = getColumnMeans(samples);
         for (int i = 0; i < k; i++) {
             final double mean = alpha[i] / alpha0;
             Assertions.assertEquals(mean, means[i], mean * relativeTolerance, "Mean");
         }
 
         // Variance = alpha_i (alpha_0 - alpha_i) / alpha_0^2 (alpha_0 + 1)
         // Covariance = -alpha_i * alpha_j / alpha_0^2 (alpha_0 + 1)
         final double[][] covars = getCovariance(samples);
         final double denom = alpha0 * alpha0 * (alpha0 + 1);
         for (int i = 0; i < k; i++) {
             final double var = alpha[i] * (alpha0 - alpha[i]) / denom;
             Assertions.assertEquals(var, covars[i][i], var * relativeTolerance, "Variance");
             for (int j = i + 1; j < k; j++) {
                 final double covar = -alpha[i] * alpha[j] / denom;
                 Assertions.assertEquals(covar, covars[i][j], Math.abs(covar) * relativeTolerance, "Covariance");
             }
         }
     }
 
     /**
      * Assert the sample has the correct length and sums to 1.
      *
      * @param k Expected number of categories.
      * @param x Sample.
      */
     private static void assertSample(int k, double[] x) {
         Assertions.assertEquals(k, x.length, "Number of categories");
         // There are always at least 2 categories
         double sum = x[0] + x[1];
         // Sum the rest
         for (int i = 2; i < x.length; i++) {
             sum += x[i];
         }
         Assertions.assertEquals(1.0, sum, 1e-10, "Invalid sum");
     }
 
     /**
      * Gets the column means. This is done using the same method as the means in the
      * Apache Commons Math Covariance class by using the Mean class.
      *
      * @param data the data
      * @return the column means
      */
     private static double[] getColumnMeans(double[][] data) {
         final Array2DRowRealMatrix m = new Array2DRowRealMatrix(data, false);
         final Mean mean = new Mean();
         final double[] means = new double[m.getColumnDimension()];
         for (int i = 0; i < means.length; i++) {
             means[i] = mean.evaluate(m.getColumn(i));
         }
         return means;
     }
 
     /**
      * Gets the covariance.
      *
      * @param data the data
      * @return the covariance
      */
     private static double[][] getCovariance(double[][] data) {
         final Array2DRowRealMatrix m = new Array2DRowRealMatrix(data, false);
         return new Covariance(m).getCovarianceMatrix().getData();
     }
 }