/*
    * 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.examples.stress;
  
  import org.apache.commons.rng.UniformRandomProvider;
  import org.apache.commons.rng.core.source32.IntProvider;
  import org.apache.commons.rng.core.source32.RandomIntSource;
  import org.apache.commons.rng.core.source64.RandomLongSource;
  import org.apache.commons.rng.simple.RandomSource;
  import org.junit.jupiter.api.Assertions;
  import org.junit.jupiter.api.Test;
  
  import java.io.ByteArrayOutputStream;
  import java.io.DataOutputStream;
  import java.io.IOException;
  import java.io.OutputStream;
  import java.nio.ByteOrder;
  import java.util.function.BiConsumer;
  import java.util.function.UnaryOperator;
  
  /**
   * Tests for {@link RngDataOutput}.
   */
  class RngDataOutputTest {
      /**
       * A factory for creating RngDataOutput objects.
       */
      interface RngDataOutputFactory {
          /**
           * Create a new instance.
           *
           * @param out Output stream.
           * @param size Number of values to write.
           * @param byteOrder Byte order.
           * @return the data output
           */
          RngDataOutput create(OutputStream out, int size, ByteOrder byteOrder);
      }
  
      @Test
      void testIntBigEndian() throws IOException {
          assertRngOutput(RandomSource.PCG_MCG_XSH_RS_32,
              UnaryOperator.identity(),
              RngDataOutputTest::writeInt,
              RngDataOutput::ofInt, ByteOrder.BIG_ENDIAN);
      }
  
      @Test
      void testIntLittleEndian() throws IOException {
          assertRngOutput(RandomSource.PCG_MCG_XSH_RS_32,
              RNGUtils::createReverseBytesProvider,
              RngDataOutputTest::writeInt,
              RngDataOutput::ofInt, ByteOrder.LITTLE_ENDIAN);
      }
  
      @Test
      void testLongBigEndian() throws IOException {
          assertRngOutput(RandomSource.SPLIT_MIX_64,
              UnaryOperator.identity(),
              RngDataOutputTest::writeLong,
              RngDataOutput::ofLong, ByteOrder.BIG_ENDIAN);
      }
  
      @Test
      void testLongLittleEndian() throws IOException {
          assertRngOutput(RandomSource.SPLIT_MIX_64,
              RNGUtils::createReverseBytesProvider,
              RngDataOutputTest::writeLong,
              RngDataOutput::ofLong, ByteOrder.LITTLE_ENDIAN);
      }
  
      @Test
      void testLongAsHLIntBigEndian() throws IOException {
          assertRngOutput(RandomSource.SPLIT_MIX_64,
              // Convert to an int provider so it is detected as requiring double the
              // length output.
              rng -> new HiLoCachingIntProvider(rng),
              RngDataOutputTest::writeInt,
              RngDataOutput::ofLongAsHLInt, ByteOrder.BIG_ENDIAN);
      }
  
      @Test
      void testLongAsHLIntLittleEndian() throws IOException {
          assertRngOutput(RandomSource.SPLIT_MIX_64,
              // Convert SplitMix64 to an int provider so it is detected as requiring double the
             // length output. Then reverse the bytes.
             rng -> new HiLoCachingIntProvider(rng) {
                 @Override
                 public int next() {
                     return Integer.reverseBytes(super.next());
                 }
             },
             RngDataOutputTest::writeInt,
             RngDataOutput::ofLongAsHLInt, ByteOrder.LITTLE_ENDIAN);
     }
 
 
     @Test
     void testLongAsLHIntBigEndian() throws IOException {
         assertRngOutput(RandomSource.SPLIT_MIX_64,
             // Convert to an int provider so it is detected as requiring double the
             // length output.
             rng -> new LoHiCachingIntProvider(rng),
             RngDataOutputTest::writeInt,
             RngDataOutput::ofLongAsLHInt, ByteOrder.BIG_ENDIAN);
     }
 
     @Test
     void testLongAsLHIntLittleEndian() throws IOException {
         assertRngOutput(RandomSource.SPLIT_MIX_64,
             // Convert to an int provider so it is detected as requiring double the
             // length output. Then reverse the bytes.
             rng -> new LoHiCachingIntProvider(rng) {
                 @Override
                 public int next() {
                     return Integer.reverseBytes(super.next());
                 }
             },
             RngDataOutputTest::writeInt,
             RngDataOutput::ofLongAsLHInt, ByteOrder.LITTLE_ENDIAN);
     }
 
     private static void writeInt(DataOutputStream sink, UniformRandomProvider rng) {
         try {
             sink.writeInt(rng.nextInt());
         } catch (IOException e) {
             Assertions.fail();
         }
     }
 
     private static void writeLong(DataOutputStream sink, UniformRandomProvider rng) {
         try {
             sink.writeLong(rng.nextLong());
         } catch (IOException e) {
             Assertions.fail();
         }
     }
 
     /**
      * Assert the byte output from the source is the same. Creates two instances of the same
      * RNG with the same seed. The first is converted to a new RNG using the {@code rngConverter}.
      * This is used to output raw bytes using a {@link DataOutputStream} via the specified
      * {@code pipe}. The second is used to output raw bytes via the {@link RngDataOutput} class
      * created using the {@code factory}.
      *
      * <p>The random source should output native {@code int} or {@code long} values.
      *
      * @param source Random source.
      * @param rngConverter Converter for the raw RNG.
      * @param pipe Pipe to send data from the RNG to the DataOutputStream.
      * @param factory Factory for the RngDataOutput.
      * @param byteOrder Byte order for the RngDataOutput.
      * @throws IOException Signals that an I/O exception has occurred.
      */
     private static void assertRngOutput(RandomSource source,
         UnaryOperator<UniformRandomProvider> rngConverter,
         BiConsumer<DataOutputStream, UniformRandomProvider> pipe,
         RngDataOutputFactory factory,
         ByteOrder byteOrder) throws IOException {
         final long seed = RandomSource.createLong();
         UniformRandomProvider rng1 = source.create(seed);
         UniformRandomProvider rng2 = source.create(seed);
         final int size = 37;
         for (int repeats = 1; repeats <= 2; repeats++) {
             byte[] expected = createBytes(rng1, size, repeats, rngConverter, pipe);
             byte[] actual = writeRngOutput(rng2, size, repeats, byteOrder, factory);
             Assertions.assertArrayEquals(expected, actual);
         }
     }
 
     /**
      * Convert the RNG and then creates bytes from the RNG using the pipe.
      *
      * @param rng RNG.
      * @param size The number of values to send to the pipe.
      * @param repeats The number of repeat iterations.
      * @param rngConverter Converter for the raw RNG.
      * @param pipe Pipe to send data from the RNG to the DataOutputStream.
      * @return the bytes
      * @throws IOException Signals that an I/O exception has occurred.
      */
     private static byte[] createBytes(UniformRandomProvider rng, int size, int repeats,
         UnaryOperator<UniformRandomProvider> rngConverter,
         BiConsumer<DataOutputStream, UniformRandomProvider> pipe) throws IOException {
         UniformRandomProvider rng2 = rngConverter.apply(rng);
         // If the factory converts to an IntProvider then output twice the size
         if (rng instanceof RandomLongSource && rng2 instanceof RandomIntSource) {
             size *= 2;
         }
         ByteArrayOutputStream out = new ByteArrayOutputStream();
         try (DataOutputStream sink = new DataOutputStream(out)) {
             for (int j = 0; j < repeats; j++) {
                 for (int i = 0; i < size; i++) {
                     pipe.accept(sink, rng2);
                 }
             }
         }
         return out.toByteArray();
     }
 
     /**
      * Write the RNG to the RngDataOutput built by the factory.
      *
      * @param rng RNG.
      * @param size The number of values to send to the RngDataOutput.
      * @param repeats The number of repeat iterations.
      * @param byteOrder Byte order for the RngDataOutput.
      * @param factory Factory for the RngDataOutput.
      * @return the bytes
      * @throws IOException Signals that an I/O exception has occurred.
      */
     private static byte[] writeRngOutput(UniformRandomProvider rng, int size, int repeats,
         ByteOrder byteOrder, RngDataOutputFactory factory) throws IOException {
         ByteArrayOutputStream out = new ByteArrayOutputStream();
         try (RngDataOutput sink = factory.create(out, size, byteOrder)) {
             for (int j = 0; j < repeats; j++) {
                 sink.write(rng);
             }
         }
         return out.toByteArray();
     }
 
     private static class LoHiCachingIntProvider extends IntProvider {
         private long source = -1;
         private final UniformRandomProvider rng;
 
         LoHiCachingIntProvider(UniformRandomProvider rng) {
             this.rng = rng;
         }
 
         @Override
         public int next() {
             long next = source;
             if (next < 0) {
                 // refill
                 next = rng.nextLong();
                 // store hi
                 source = next >>> 32;
                 // extract low
                 return (int) next;
             }
             final int v = (int) next;
             // reset
             source = -1;
             return v;
         }
     }
 
     private static class HiLoCachingIntProvider extends IntProvider {
         private long source = -1;
         private final UniformRandomProvider rng;
 
         HiLoCachingIntProvider(UniformRandomProvider rng) {
             this.rng = rng;
         }
 
         @Override
         public int next() {
             long next = source;
             if (next < 0) {
                 // refill
                 next = rng.nextLong();
                 // store low
                 source = next & 0xffff_ffffL;
                 // extract hi
                 return (int) (next >>> 32);
             }
             final int v = (int) next;
             // reset
             source = -1;
             return v;
         }
     }
 }