/*
    * 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.core;
  
  import org.junit.jupiter.api.Assertions;
  import org.junit.jupiter.api.Assumptions;
  import org.junit.jupiter.params.ParameterizedTest;
  import org.junit.jupiter.params.provider.MethodSource;
  
  import java.util.Arrays;
  
  import org.apache.commons.rng.JumpableUniformRandomProvider;
  import org.apache.commons.rng.LongJumpableUniformRandomProvider;
  import org.apache.commons.rng.RandomProviderState;
  import org.apache.commons.rng.RestorableUniformRandomProvider;
  import org.apache.commons.rng.UniformRandomProvider;
  import org.apache.commons.rng.core.source32.IntProvider;
  import org.apache.commons.rng.core.source64.LongProvider;
  
  /**
   * Tests which all {@link JumpableUniformRandomProvider} generators must pass.
   */
  class JumpableProvidersParametricTest {
      /** The default state for the IntProvider. */
      private static final byte[] INT_PROVIDER_STATE;
      /** The default state for the LongProvider. */
      private static final byte[] LONG_PROVIDER_STATE;
  
      static {
          INT_PROVIDER_STATE = new State32Generator().getState();
          LONG_PROVIDER_STATE = new State64Generator().getState();
      }
  
      /**
       * Gets the list of Jumpable generators.
       *
       * @return the list
       */
      private static Iterable<JumpableUniformRandomProvider> getJumpableProviders() {
          return ProvidersList.listJumpable();
      }
  
      /**
       * Gets the function using the {@link LongJumpableUniformRandomProvider#longJump()} method.
       * If the RNG is not long jumpable then this will raise an exception to skip the test.
       *
       * @param generator RNG under test.
       * @return the jump function
       */
      private static TestJumpFunction getLongJumpFunction(JumpableUniformRandomProvider generator) {
          Assumptions.assumeTrue(generator instanceof LongJumpableUniformRandomProvider, "No long jump function");
          final LongJumpableUniformRandomProvider rng2 = (LongJumpableUniformRandomProvider) generator;
          return rng2::jump;
      }
  
      /**
       * Test that the random generator returned from the jump is a new instance of the same class.
       */
      @ParameterizedTest
      @MethodSource("getJumpableProviders")
      void testJumpReturnsACopy(JumpableUniformRandomProvider generator) {
          assertJumpReturnsACopy(generator::jump, generator);
      }
  
      /**
       * Test that the random generator returned from the long jump is a new instance of the same class.
       */
      @ParameterizedTest
      @MethodSource("getJumpableProviders")
      void testLongJumpReturnsACopy(JumpableUniformRandomProvider generator) {
          assertJumpReturnsACopy(getLongJumpFunction(generator), generator);
      }
  
      /**
       * Assert that the random generator returned from the jump function is a new instance of the same class.
       *
       * @param jumpFunction Jump function to test.
       * @param generator RNG under test.
       */
      private static void assertJumpReturnsACopy(TestJumpFunction jumpFunction,
                                                 JumpableUniformRandomProvider generator) {
          final UniformRandomProvider copy = jumpFunction.jump();
          Assertions.assertNotSame(generator, copy, "The copy instance should be a different object");
          Assertions.assertEquals(generator.getClass(), copy.getClass(), "The copy instance should be the same class");
      }
 
     /**
      * Test that the random generator state of the copy instance returned from the jump
      * matches the input state.
      */
     @ParameterizedTest
     @MethodSource("getJumpableProviders")
     void testJumpCopyMatchesPreJumpState(JumpableUniformRandomProvider generator) {
         assertCopyMatchesPreJumpState(generator::jump, generator);
     }
 
     /**
      * Test that the random generator state of the copy instance returned from the long jump
      * matches the input state.
      */
     @ParameterizedTest
     @MethodSource("getJumpableProviders")
     void testLongJumpCopyMatchesPreJumpState(JumpableUniformRandomProvider generator) {
         assertCopyMatchesPreJumpState(getLongJumpFunction(generator), generator);
     }
 
     /**
      * Assert that the random generator state of the copy instance returned from the jump
      * function matches the input state.
      *
      * <p>The generator must be a {@link RestorableUniformRandomProvider} and return an
      * instance of {@link RandomProviderDefaultState}.</p>
      *
      * <p>The input generator is sampled using methods in the
      * {@link UniformRandomProvider} interface, the state is saved and a jump is
      * performed. The states from the pre-jump generator and the returned copy instance
      * must match.</p>
      *
      * <p>This test targets any cached state of the default implementation of a generator
      * in {@link IntProvider} and {@link LongProvider} such as the state cached for the
      * nextBoolean() and nextInt() functions.</p>
      *
      * @param jumpFunction Jump function to test.
      * @param generator RNG under test.
      */
     private static void assertCopyMatchesPreJumpState(TestJumpFunction jumpFunction,
                                                       JumpableUniformRandomProvider generator) {
         Assumptions.assumeTrue(generator instanceof RestorableUniformRandomProvider, "Not a restorable RNG");
 
         for (int repeats = 0; repeats < 2; repeats++) {
             // Exercise the generator.
             // This calls nextInt() once so the default implementation of LongProvider
             // should have cached a state for nextInt() in one of the two repeats.
             // Calls nextBoolean() to ensure a cached state in one of the two repeats.
             generator.nextInt();
             generator.nextBoolean();
 
             final RandomProviderState preJumpState = ((RestorableUniformRandomProvider) generator).saveState();
             Assumptions.assumeTrue(preJumpState instanceof RandomProviderDefaultState, "Not a recognised state");
 
             final UniformRandomProvider copy = jumpFunction.jump();
 
             final RandomProviderState copyState = ((RestorableUniformRandomProvider) copy).saveState();
             final RandomProviderDefaultState expected = (RandomProviderDefaultState) preJumpState;
             final RandomProviderDefaultState actual = (RandomProviderDefaultState) copyState;
             Assertions.assertArrayEquals(expected.getState(), actual.getState(),
                 "The copy instance state should match the state of the original");
         }
     }
 
     /**
      * Test that a jump resets the state of the default implementation of a generator in
      * {@link IntProvider} and {@link LongProvider}.
      */
     @ParameterizedTest
     @MethodSource("getJumpableProviders")
     void testJumpResetsDefaultState(JumpableUniformRandomProvider generator) {
         assertJumpResetsDefaultState(generator::jump, generator);
     }
 
     /**
      * Test that a long jump resets the state of the default implementation of a generator in
      * {@link IntProvider} and {@link LongProvider}.
      */
     @ParameterizedTest
     @MethodSource("getJumpableProviders")
     void testLongJumpResetsDefaultState(JumpableUniformRandomProvider generator) {
         assertJumpResetsDefaultState(getLongJumpFunction(generator), generator);
     }
 
     /**
      * Assert the jump resets the specified number of bytes of the state. The bytes are
      * checked from the end of the saved state.
      *
      * <p>This is intended to check the default state of the base implementation of
      * {@link IntProvider} and {@link LongProvider} is reset.</p>
      *
      * @param jumpFunction Jump function to test.
      * @param generator RNG under test.
      */
     private static void assertJumpResetsDefaultState(TestJumpFunction jumpFunction,
                                                      JumpableUniformRandomProvider generator) {
         byte[] expected;
         if (generator instanceof IntProvider) {
             expected = INT_PROVIDER_STATE;
         } else if (generator instanceof LongProvider) {
             expected = LONG_PROVIDER_STATE;
         } else {
             throw new AssertionError("Unsupported RNG");
         }
         final int stateSize = expected.length;
         for (int repeats = 0; repeats < 2; repeats++) {
             // Exercise the generator.
             // This calls nextInt() once so the default implementation of LongProvider
             // should have cached a state for nextInt() in one of the two repeats.
             // Calls nextBoolean() to ensure a cached state in one of the two repeats.
             generator.nextInt();
             generator.nextBoolean();
 
             jumpFunction.jump();
 
             // An Int/LongProvider so must be a RestorableUniformRandomProvider
             final RandomProviderState postJumpState = ((RestorableUniformRandomProvider) generator).saveState();
             final byte[] actual = ((RandomProviderDefaultState) postJumpState).getState();
 
             Assumptions.assumeTrue(actual.length >= stateSize, "Implementation has removed default state");
 
             // The implementation requires that any sub-class state is prepended to the
             // state thus the default state is at the end.
             final byte[] defaultState = Arrays.copyOfRange(actual, actual.length - stateSize, actual.length);
             Assertions.assertArrayEquals(expected, defaultState,
                 "The jump should reset the default state to zero");
         }
     }
 
     /**
      * Dummy class for checking the state size of the IntProvider.
      */
     static class State32Generator extends IntProvider {
         /** {@inheritDoc} */
         @Override
         public int next() {
             return 0;
         }
 
         /**
          * Gets the default state. This captures the initial state of the IntProvider
          * including the values of the cache variables.
          *
          * @return the state
          */
         byte[] getState() {
             return getStateInternal();
         }
     }
 
     /**
      * Dummy class for checking the state size of the LongProvider.
      */
     static class State64Generator extends LongProvider {
         /** {@inheritDoc} */
         @Override
         public long next() {
             return 0;
         }
 
         /**
          * Gets the default state. This captures the initial state of the LongProvider
          * including the values of the cache variables.
          *
          * @return the state size
          */
         byte[] getState() {
             return getStateInternal();
         }
     }
 
     /**
      * Specify the jump operation to test.
      *
      * <p>This allows testing {@link JumpableUniformRandomProvider} or
      * {@link LongJumpableUniformRandomProvider}.</p>
      */
     interface TestJumpFunction {
         /**
          * Perform the jump and return a pre-jump copy.
          *
          * @return the pre-jump copy.
          */
         UniformRandomProvider jump();
     }
 }