/*
    * 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.rng.UniformRandomProvider;
  
  /**
   * Discrete uniform distribution sampler generating values of type {@code long}.
   *
   * <p>Sampling uses {@link UniformRandomProvider#nextLong}.</p>
   *
   * <p>When the range is a power of two the number of calls is 1 per sample.
   * Otherwise a rejection algorithm is used to ensure uniformity. In the worst
   * case scenario where the range spans half the range of a {@code long}
   * (2<sup>63</sup> + 1) the expected number of calls is 2 per sample.</p>
   *
   * @since 1.4
   */
  public abstract class UniformLongSampler implements SharedStateLongSampler {
      /** Underlying source of randomness. */
      protected final UniformRandomProvider rng;
  
      /**
       * Discrete uniform distribution sampler when the sample value is fixed.
       */
      private static class FixedUniformLongSampler extends UniformLongSampler {
          /** The value. */
          private final long value;
  
          /**
           * @param value The value.
           */
          FixedUniformLongSampler(long value) {
              // No requirement for the RNG
              super(null);
              this.value = value;
          }
  
          @Override
          public long sample() {
              return value;
          }
  
          @Override
          public String toString() {
              // No RNG to include in the string
              return "Uniform deviate [X=" + value + "]";
          }
  
          @Override
          public UniformLongSampler withUniformRandomProvider(UniformRandomProvider rng) {
              // No requirement for the RNG
              return this;
          }
      }
  
      /**
       * Discrete uniform distribution sampler when the range is a power of 2 and greater than 1.
       * This sampler assumes the lower bound of the range is 0.
       *
       * <p>Note: This cannot be used when the range is 1 (2^0) as the shift would be 64-bits
       * which is ignored by the shift operator.</p>
       */
      private static class PowerOf2RangeUniformLongSampler extends UniformLongSampler {
          /** Bit shift to apply to the long sample. */
          private final int shift;
  
          /**
           * @param rng Generator of uniformly distributed random numbers.
           * @param range Maximum range of the sample (exclusive).
           * Must be a power of 2 greater than 2^0.
           */
          PowerOf2RangeUniformLongSampler(UniformRandomProvider rng,
                                          long range) {
              super(rng);
              this.shift = Long.numberOfLeadingZeros(range) + 1;
          }
  
          /**
           * @param rng Generator of uniformly distributed random numbers.
           * @param source Source to copy.
           */
          PowerOf2RangeUniformLongSampler(UniformRandomProvider rng,
                                          PowerOf2RangeUniformLongSampler source) {
             super(rng);
             this.shift = source.shift;
         }
 
         @Override
         public long sample() {
             // Use a bit shift to favour the most significant bits.
             return rng.nextLong() >>> shift;
         }
 
         @Override
         public UniformLongSampler withUniformRandomProvider(UniformRandomProvider rng) {
             return new PowerOf2RangeUniformLongSampler(rng, this);
         }
     }
 
     /**
      * Discrete uniform distribution sampler when the range is small
      * enough to fit in a positive long.
      * This sampler assumes the lower bound of the range is 0 and the range is
      * non-zero.
      */
     private static class SmallRangeUniformLongSampler extends UniformLongSampler {
         /** Maximum range of the sample (exclusive). */
         private final long n;
         /** Limit of the uniform range (inclusive) to sample a positive long.
          * This is the largest positive multiple of {@code n} minus 1:
          * {@code floor(2^63 / n) * n - 1}.
          * The -1 changes the limit to an inclusive bound and allows support
          * for a power of 2 range. */
         private final long limit;
 
         /**
          * @param rng Generator of uniformly distributed random numbers.
          * @param range Maximum range of the sample (exclusive).
          */
         SmallRangeUniformLongSampler(UniformRandomProvider rng,
                                      long range) {
             super(rng);
             this.n = range;
             // Set the upper limit for the positive long.
             // The sample must be selected from the largest multiple
             // of 'range' that fits within a positive value:
             // limit = floor(2^63 / range) * range
             //       = 2^63 - (2^63 % range)
             // Sample:
             //   X in [0, limit) or X in [0, limit - 1]
             //   return X % range
             // This is a one-off computation cost.
             // The divide will truncate towards zero (do not use Math.floorDiv).
             // Note: This is invalid if range is not strictly positive.
             limit = (Long.MIN_VALUE / range) * -range - 1;
         }
 
         /**
          * @param rng Generator of uniformly distributed random numbers.
          * @param source Source to copy.
          */
         SmallRangeUniformLongSampler(UniformRandomProvider rng,
                                      SmallRangeUniformLongSampler source) {
             super(rng);
             this.n = source.n;
             this.limit = source.limit;
         }
 
         @Override
         public long sample() {
             // Note:
             // This will have the same output as the rejection algorithm
             // from o.a.c.rng.core.BaseProvider. The limit for the uniform
             // positive value can be pre-computed. This ensures exactly
             // 1 modulus operation per call.
             for (;;) {
                 // bits in [0, limit]
                 final long bits = rng.nextLong() >>> 1;
                 if (bits <= limit) {
                     return bits % n;
                 }
             }
         }
 
         @Override
         public UniformLongSampler withUniformRandomProvider(UniformRandomProvider rng) {
             return new SmallRangeUniformLongSampler(rng, this);
         }
     }
 
     /**
      * Discrete uniform distribution sampler when the range between lower and upper is too large
      * to fit in a positive long.
      */
     private static class LargeRangeUniformLongSampler extends UniformLongSampler {
         /** Lower bound. */
         private final long lower;
         /** Upper bound. */
         private final long upper;
 
         /**
          * @param rng Generator of uniformly distributed random numbers.
          * @param lower Lower bound (inclusive) of the distribution.
          * @param upper Upper bound (inclusive) of the distribution.
          */
         LargeRangeUniformLongSampler(UniformRandomProvider rng,
                                      long lower,
                                      long upper) {
             super(rng);
             this.lower = lower;
             this.upper = upper;
         }
 
         @Override
         public long sample() {
             // Use a simple rejection method.
             // This is used when (upper-lower) >= Long.MAX_VALUE.
             // This will loop on average 2 times in the worst case scenario
             // when (upper-lower) == Long.MAX_VALUE.
             while (true) {
                 final long r = rng.nextLong();
                 if (r >= lower &&
                     r <= upper) {
                     return r;
                 }
             }
         }
 
         @Override
         public UniformLongSampler withUniformRandomProvider(UniformRandomProvider rng) {
             return new LargeRangeUniformLongSampler(rng, lower, upper);
         }
     }
 
     /**
      * Adds an offset to an underlying discrete sampler.
      */
     private static class OffsetUniformLongSampler extends UniformLongSampler {
         /** The offset. */
         private final long offset;
         /** The long sampler. */
         private final UniformLongSampler sampler;
 
         /**
          * @param offset The offset for the sample.
          * @param sampler The discrete sampler.
          */
         OffsetUniformLongSampler(long offset,
                                  UniformLongSampler sampler) {
             // No requirement for the RNG
             super(null);
             this.offset = offset;
             this.sampler = sampler;
         }
 
         @Override
         public long sample() {
             return offset + sampler.sample();
         }
 
         @Override
         public String toString() {
             return sampler.toString();
         }
 
         @Override
         public UniformLongSampler withUniformRandomProvider(UniformRandomProvider rng) {
             return new OffsetUniformLongSampler(offset, sampler.withUniformRandomProvider(rng));
         }
     }
 
     /**
      * @param rng Generator of uniformly distributed random numbers.
      */
     UniformLongSampler(UniformRandomProvider rng) {
         this.rng = rng;
     }
 
     /** {@inheritDoc} */
     @Override
     public String toString() {
         return "Uniform deviate [" + rng.toString() + "]";
     }
 
     /** {@inheritDoc} */
     // Redeclare the signature to return a UniformLongSampler not a SharedStateLongSampler
     @Override
     public abstract UniformLongSampler withUniformRandomProvider(UniformRandomProvider rng);
 
     /**
      * Creates a new discrete uniform distribution sampler.
      *
      * @param rng Generator of uniformly distributed random numbers.
      * @param lower Lower bound (inclusive) of the distribution.
      * @param upper Upper bound (inclusive) of the distribution.
      * @return the sampler
      * @throws IllegalArgumentException if {@code lower > upper}.
      */
     public static UniformLongSampler of(UniformRandomProvider rng,
                                         long lower,
                                         long upper) {
         if (lower > upper) {
             throw new IllegalArgumentException(lower  + " > " + upper);
         }
 
         // Choose the algorithm depending on the range
 
         // Edge case for no range.
         // This must be done first as the methods to handle lower == 0
         // do not handle upper == 0.
         if (upper == lower) {
             return new FixedUniformLongSampler(lower);
         }
 
         // Algorithms to ignore the lower bound if it is zero.
         if (lower == 0) {
             return createZeroBoundedSampler(rng, upper);
         }
 
         final long range = (upper - lower) + 1;
         // Check power of 2 first to handle range == 2^63.
         if (isPowerOf2(range)) {
             return new OffsetUniformLongSampler(lower,
                                                 new PowerOf2RangeUniformLongSampler(rng, range));
         }
         if (range <= 0) {
             // The range is too wide to fit in a positive long (larger
             // than 2^63); use a simple rejection method.
             // Note: if range == 0 then the input is [Long.MIN_VALUE, Long.MAX_VALUE].
             // No specialisation exists for this and it is handled as a large range.
             return new LargeRangeUniformLongSampler(rng, lower, upper);
         }
         // Use a sample from the range added to the lower bound.
         return new OffsetUniformLongSampler(lower,
                                             new SmallRangeUniformLongSampler(rng, range));
     }
 
     /**
      * Create a new sampler for the range {@code 0} inclusive to {@code upper} inclusive.
      *
      * <p>This can handle any positive {@code upper}.
      *
      * @param rng Generator of uniformly distributed random numbers.
      * @param upper Upper bound (inclusive) of the distribution. Must be positive.
      * @return the sampler
      */
     private static UniformLongSampler createZeroBoundedSampler(UniformRandomProvider rng,
                                                                long upper) {
         // Note: Handle any range up to 2^63 (which is negative as a signed
         // 64-bit long but handled as a power of 2)
         final long range = upper + 1;
         return isPowerOf2(range) ?
             new PowerOf2RangeUniformLongSampler(rng, range) :
             new SmallRangeUniformLongSampler(rng, range);
     }
 
     /**
      * Checks if the value is a power of 2.
      *
      * <p>This returns {@code true} for the value {@code Long.MIN_VALUE} which can be
      * handled as an unsigned long of 2^63.</p>
      *
      * @param value Value.
      * @return {@code true} if a power of 2
      */
     private static boolean isPowerOf2(final long value) {
         return value != 0 && (value & (value - 1)) == 0;
     }
 }