/*
    * 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
    *
    *      https://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.statistics.distribution;
  
  import java.util.stream.IntStream;
  import org.junit.jupiter.api.Assertions;
  import org.junit.jupiter.api.Test;
  
  /**
   * Test cases for AbstractDiscreteDistribution default implementations.
   */
  class AbstractDiscreteDistributionTest {
      private final DiceDistribution diceDistribution = new DiceDistribution();
  
      @Test
      void testInverseCumulativeProbabilityMethod() {
          // This must be consistent with its own cumulative probability
          final double[] p = IntStream.rangeClosed(1, 6).mapToDouble(diceDistribution::cumulativeProbability).toArray();
          Assertions.assertEquals(1.0, p[5], "Incorrect cumulative probability at upper bound");
          Assertions.assertEquals(1, diceDistribution.inverseCumulativeProbability(0));
          for (int i = 0; i < 6; i++) {
              final int x = i + 1;
              Assertions.assertEquals(x, diceDistribution.inverseCumulativeProbability(Math.nextDown(p[i])));
              Assertions.assertEquals(x, diceDistribution.inverseCumulativeProbability(p[i]));
              if (x < 6) {
                  Assertions.assertEquals(x + 1, diceDistribution.inverseCumulativeProbability(Math.nextUp(p[i])));
              }
          }
      }
  
      @Test
      void testInverseSurvivalProbabilityMethod() {
          // This must be consistent with its own survival probability
          final double[] p = IntStream.rangeClosed(1, 6).mapToDouble(diceDistribution::survivalProbability).toArray();
          Assertions.assertEquals(0.0, p[5], "Incorrect survival probability at upper bound");
          Assertions.assertEquals(1, diceDistribution.inverseSurvivalProbability(1));
          for (int i = 0; i < 6; i++) {
              final int x = i + 1;
              Assertions.assertEquals(x, diceDistribution.inverseSurvivalProbability(Math.nextUp(p[i])));
              Assertions.assertEquals(x, diceDistribution.inverseSurvivalProbability(p[i]));
              if (x < 6) {
                  Assertions.assertEquals(x + 1, diceDistribution.inverseSurvivalProbability(Math.nextDown(p[i])));
              }
          }
      }
  
      @Test
      void testCumulativeProbabilitiesSingleArguments() {
          final double p = diceDistribution.probability(1);
          for (int i = 1; i <= 6; i++) {
              Assertions.assertEquals(p * i,
                      diceDistribution.cumulativeProbability(i), Math.ulp(p * i));
          }
          Assertions.assertEquals(0.0, diceDistribution.cumulativeProbability(0));
          Assertions.assertEquals(1.0, diceDistribution.cumulativeProbability(7));
      }
  
      @Test
      void testProbabilitiesRangeArguments() {
          final double p = diceDistribution.probability(1);
          int lower = 0;
          int upper = 6;
          for (int i = 0; i < 2; i++) {
              // cum(0,6) = p(0 < X <= 6) = 1, cum(1,5) = 4/6, cum(2,4) = 2/6
              Assertions.assertEquals(1 - p * 2 * i,
                      diceDistribution.probability(lower, upper), 1E-12);
              lower++;
              upper--;
          }
          for (int i = 0; i < 6; i++) {
              Assertions.assertEquals(p, diceDistribution.probability(i, i + 1), 1E-12);
          }
      }
  
      @Test
      void testInverseCumulativeProbabilityExtremes() {
          // Require a lower bound of MIN_VALUE and the cumulative probability
          // at that bound to be lower/higher than the argument cumulative probability.
          // Use a uniform distribution so that the default search is supported
          // for the inverse.
          final DiscreteDistribution dist = new AbstractDiscreteDistribution() {
              @Override
              public double probability(int x) {
                  throw new AssertionError();
              }
             @Override
             public double cumulativeProbability(int x) {
                 final int y = x - Integer.MIN_VALUE;
                 if (y == 0) {
                     return 0.25;
                 } else if (y == 1) {
                     return 0.5;
                 } else if (y == 2) {
                     return 0.75;
                 } else {
                     return 1.0;
                 }
             }
             @Override
             public double getMean() {
                 return 1.5 + Integer.MIN_VALUE;
             }
             @Override
             public double getVariance() {
                 // n = upper - lower + 1, the variance is (n^2 - 1) / 12
                 return 15.0 / 12;
             }
             @Override
             public int getSupportLowerBound() {
                 return Integer.MIN_VALUE;
             }
             @Override
             public int getSupportUpperBound() {
                 return Integer.MIN_VALUE + 3;
             }
         };
         Assertions.assertEquals(dist.getSupportLowerBound(), dist.inverseCumulativeProbability(0.0));
         Assertions.assertEquals(Integer.MIN_VALUE, dist.inverseCumulativeProbability(0.05));
         Assertions.assertEquals(Integer.MIN_VALUE + 1, dist.inverseCumulativeProbability(0.35));
         Assertions.assertEquals(Integer.MIN_VALUE + 2, dist.inverseCumulativeProbability(0.55));
         Assertions.assertEquals(dist.getSupportUpperBound(), dist.inverseCumulativeProbability(1.0));
 
         Assertions.assertEquals(dist.getSupportLowerBound(), dist.inverseSurvivalProbability(1.0));
         Assertions.assertEquals(Integer.MIN_VALUE, dist.inverseSurvivalProbability(0.95));
         Assertions.assertEquals(Integer.MIN_VALUE + 1, dist.inverseSurvivalProbability(0.65));
         Assertions.assertEquals(Integer.MIN_VALUE + 2, dist.inverseSurvivalProbability(0.45));
         Assertions.assertEquals(dist.getSupportUpperBound(), dist.inverseSurvivalProbability(0.0));
     }
 
     @Test
     void testInverseCumulativeProbabilityWithNoMean() {
         // A NaN mean will invalidate the Chebyshev inequality
         // to prevent bracketing
         final DiscreteDistribution dist = new AbstractDiscreteDistribution() {
             @Override
             public double probability(int x) {
                 throw new AssertionError();
             }
             @Override
             public double cumulativeProbability(int x) {
                 if (x == 0) {
                     return 0.25;
                 } else if (x == 1) {
                     return 0.5;
                 } else if (x == 2) {
                     return 0.75;
                 } else {
                     return 1.0;
                 }
             }
             @Override
             public double getMean() {
                 return Double.NaN;
             }
             @Override
             public double getVariance() {
                 return Double.NaN;
             }
             @Override
             public int getSupportLowerBound() {
                 return 0;
             }
             @Override
             public int getSupportUpperBound() {
                 return 3;
             }
         };
         Assertions.assertEquals(dist.getSupportLowerBound(), dist.inverseCumulativeProbability(0.0));
         Assertions.assertEquals(0, dist.inverseCumulativeProbability(0.05));
         Assertions.assertEquals(1, dist.inverseCumulativeProbability(0.35));
         Assertions.assertEquals(2, dist.inverseCumulativeProbability(0.55));
         Assertions.assertEquals(dist.getSupportUpperBound(), dist.inverseCumulativeProbability(1.0));
 
         Assertions.assertEquals(dist.getSupportLowerBound(), dist.inverseSurvivalProbability(1.0));
         Assertions.assertEquals(0, dist.inverseSurvivalProbability(0.95));
         Assertions.assertEquals(1, dist.inverseSurvivalProbability(0.65));
         Assertions.assertEquals(2, dist.inverseSurvivalProbability(0.45));
         Assertions.assertEquals(dist.getSupportUpperBound(), dist.inverseSurvivalProbability(0.0));
     }
 
     /**
      * Simple distribution modeling a 6-sided die
      */
     class DiceDistribution extends AbstractDiscreteDistribution {
         private final double p = 1d / 6d;
 
         @Override
         public double probability(int x) {
             if (x < 1 || x > 6) {
                 return 0;
             } else {
                 return p;
             }
         }
 
         @Override
         public double cumulativeProbability(int x) {
             if (x < 1) {
                 return 0;
             } else if (x >= 6) {
                 return 1;
             } else {
                 return x / 6d;
             }
         }
 
         @Override
         public double getMean() {
             return 3.5;
         }
 
         @Override
         public double getVariance() {
             return 70 / 24;  // E(X^2) - E(X)^2
         }
 
         @Override
         public int getSupportLowerBound() {
             return 1;
         }
 
         @Override
         public int getSupportUpperBound() {
             return 6;
         }
     }
 }