/*
 * 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 accord.utils;

import java.lang.reflect.Array;
import java.math.BigDecimal;
import java.math.RoundingMode;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.EnumMap;
import java.util.EnumSet;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.SortedMap;
import java.util.SortedSet;
import java.util.function.BooleanSupplier;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.IntStream;
import java.util.stream.LongStream;
import java.util.stream.Stream;

import com.google.common.collect.Iterables;

import accord.utils.random.Picker;

@SuppressWarnings("unused")
public class Gens
{
    private Gens() {
    }

    public static <T> Gen<T> flatten(Gen<Gen<T>> gen)
    {
        return rs -> gen.next(rs).next(rs);
    }

    public static <T> Gen<T> constant(T constant)
    {
        return ignore -> constant;
    }

    public static <T> Gen<T> constant(Supplier<T> constant)
    {
        return ignore -> constant.get();
    }

    /**
     * Creates a generator that randomly selects one of the provided generators on each execution, then calls {@link Gen#next(RandomSource)} on that chosen generator.
     * <p>
     * This method distributes the selection with equal probability. For weighted selection,
     * see {@link #oneOf(Map)}.
     *
     * @param <T> the type of values to generate
     * @param gens the generators to select from
     * @return a generator that uses one of the provided generators for each value
     * @throws IllegalArgumentException if gens is empty
     */
    @SuppressWarnings("unchecked")
    @SafeVarargs
    public static <T> Gen<T> oneOf(Gen<? extends T>... gens)
    {
        switch (gens.length)
        {
            case 0: throw new IllegalArgumentException("Unable to select oneOf an empty list");
            case 1: return (Gen<T>) gens[0];
        }
        return oneOf(Arrays.asList(gens));
    }

    /**
     * Creates a generator that randomly selects one of the provided generators from a list on each execution, then calls {@link Gen#next(RandomSource)} on that chosen generator.
     * <p>
     * This method distributes the selection with equal probability. For weighted selection,
     * see {@link #oneOf(Map)}.
     *
     * @param <T> the type of values to generate
     * @param gens the list of generators to select from
     * @return a generator that uses one of the provided generators for each value
     * @throws IllegalArgumentException if gens is empty
     */
    @SuppressWarnings("unchecked")
    public static <T> Gen<T> oneOf(List<Gen<? extends T>> gens)
    {
        switch (gens.size())
        {
            case 0: throw new IllegalArgumentException("Unable to select oneOf an empty list");
            case 1: return (Gen<T>) gens.get(0);
        }
        return rs -> rs.pick(gens).next(rs);
    }

    /**
     * Creates a generator that randomly selects one of the provided generators based on weight, then calls {@link Gen#next(RandomSource)} on that chosen generator.
     * <p>
     * Each generator is selected with a probability proportional to its associated weight.
     * Higher weights increase the likelihood of selection.
     *
     * @param <T> the type of values to generate
     * @param values a map of generators to their respective weights
     * @return a generator that uses one of the provided generators according to their weights
     * @throws IllegalArgumentException if values is empty
     */
    public static <T> Gen<T> oneOf(Map<Gen<T>, Integer> values)
    {
        Gen<Gen<T>> gen = pick(values);
        return rs -> gen.next(rs).next(rs);
    }

    /**
     * A builder for creating generators that randomly select from multiple generators with customizable weights.
     *
     * @see OneOfBuilder
     * @return a oneOf builder to aid in creating a {@code Gen} of type {@code T}
     * @param <T> the generator will return
     */
    public static <T> OneOfBuilder<T> oneOf()
    {
        return new OneOfBuilder<>();
    }

    /**
     * A builder for creating generators that randomly select from multiple generators with customizable weights.
     * <p>
     * This builder allows constructing complex random selection logic with:
     * <ul>
     * <li>Weighted probabilities for different generators</li>
     * <li>Dynamically determined weights at generation time (if {@link #build} is called, the weight will be {@code 1})</li>
     * </ul>
     * <p>
     * The {@link #buildWithDynamicWeights()} method is particularly powerful as it creates
     * a meta-generator that recalculates the weights for unweighted generators on each invocation,
     * producing more varied randomization patterns than static weights.
     *
     * @param <T> the type of values the generators will produce
     */
    public static class OneOfBuilder<T>
    {
        private final Map<Gen<T>, Integer> weighted = new LinkedHashMap<>();
        private final Set<Gen<T>> unweighted = new LinkedHashSet<>();
        private Gen.IntGen unknownWeightGen = Gens.ints().between(1, 10);

        public OneOfBuilder<T> add(Gen<T> gen)
        {
            unweighted.add(gen);
            return this;
        }

        public OneOfBuilder<T> add(int weight, Gen<T> gen)
        {
            weighted.put(gen, weight);
            return this;
        }

        public OneOfBuilder<T> unknownWeights(Gen.IntGen gen)
        {
            this.unknownWeightGen = gen;
            return this;
        }

        public Gen<Gen<T>> buildWithDynamicWeights()
        {
            if (unweighted.isEmpty())
            {
                Gen<T> gen = build();
                return i -> gen;
            }
            return rs -> {
                Map<Gen<T>, Integer> commands = new LinkedHashMap<>(weighted);
                for (var gen : unweighted)
                    commands.put(gen, unknownWeightGen.nextInt(rs));
                var top = pick(commands);
                return rs2 -> top.next(rs2).next(rs2);
            };
        }

        public Gen<T> build()
        {
            Map<Gen<T>, Integer> commands = new LinkedHashMap<>(weighted);
            for (var gen : unweighted)
                commands.put(gen, 1);
            var top = pick(commands);
            return rs -> top.next(rs).next(rs);
        }
    }

    /**
     * Creates a generator that randomly selects one value from an array.
     * <p>
     * Each value has an equal probability of being selected.
     *
     * @param ts the array of values to select from
     * @return a generator that produces random values from the given array
     * @throws IllegalArgumentException if the array is empty
     */
    public static Gen.IntGen pickInt(int... ts)
    {
        Invariants.requireArgument(ts.length > 0, "Unable to pick from an empty array");
        return rs -> ts[rs.nextInt(0, ts.length)];
    }

    /**
     * Creates a generator that randomly selects one value from an array.
     * <p>
     * Each value has an equal probability of being selected.
     *
     * @param <T> the type of values to generate
     * @param ts the array of values to select from
     * @return a generator that produces random values from the given array
     * @throws IllegalArgumentException if the array is empty
     */
    @SafeVarargs
    public static <T> Gen<T> pick(T... ts)
    {
        Invariants.requireArgument(ts.length > 0, "Unable to pick from an empty array");
        return pick(Arrays.asList(ts));
    }

    /**
     * Creates a generator that randomly selects one value from a list.
     * <p>
     * Each value has an equal probability of being selected.
     *
     * @param <T> the type of values to generate
     * @param ts the list of values to select from
     * @return a generator that produces random values from the given list
     * @throws IllegalArgumentException if the list is empty
     */
    public static <T> Gen<T> pick(List<T> ts)
    {
        Invariants.requireArgument(!ts.isEmpty(), "Unable to pick from an empty collection");
        Gen.IntGen offset = ints().between(0, ts.size() - 1);
        return rs -> ts.get(offset.nextInt(rs));
    }

    /**
     * Creates a generator that randomly selects one value from a set. When the set does not have deterministic
     * iteration ordering, then the values will be sorted based off their natural ordering first.
     * <p>
     * Each value has an equal probability of being selected.
     *
     * @param <T> the type of values to generate
     * @param set the set of values to select from
     * @return a generator that produces random values from the given set
     * @throws IllegalArgumentException if the set is empty
     */
    public static <T extends Comparable<T>> Gen<T> pick(Set<T> set)
    {
        Invariants.requireArgument(!set.isEmpty(), "Unable to pick from an empty collection");
        List<T> list = new ArrayList<>(set);
        // Non-ordered sets may have different iteration order on different environments, which would make a seed produce different histories!
        // To avoid such a problem, make sure to apply a deterministic function (sort).
        if (!isIterationSafe(set))
            list.sort(Comparator.naturalOrder());
        return pick(list);
    }

    /**
     * Creates a generator that randomly selects one value from a weighted map.
     * <p>
     * Each value is selected with a probability proportional to its associated weight.
     * Higher weights increase the likelihood of selection.
     * <p>
     * For deterministic behavior, this method requires maps with deterministic iteration order
     * like {@link java.util.EnumMap} or {@link java.util.LinkedHashMap}.
     *
     * @param <T> the type of values to generate
     * @param values a map of values to their respective weights
     * @return a generator that produces values according to their weights
     * @throws IllegalArgumentException if values is empty or doesn't have deterministic iteration order
     */
    public static <T> Gen<T> pick(Map<T, Integer> values)
    {
        if (values == null || values.isEmpty())
            throw new IllegalArgumentException("values is empty");
        checkIterationSafe("pick(Map)", values);
        if (values.size() == 1)
            return constant(Objects.requireNonNull(Iterables.getFirst(values.keySet(), null)));
        double totalWeight = values.values().stream().mapToDouble(Integer::intValue).sum();
        List<Weight<T>> list = new ArrayList<>(values.size());
        Iterator<Map.Entry<T, Integer>> it = values.entrySet().iterator();
        for (int i = 0; it.hasNext(); i++)
        {
            Map.Entry<T, Integer> e = it.next();
            list.add(new Weight<>(e.getKey(), e.getValue(), i));
        }
        Collections.sort(list);
        return rs -> {
            double value = rs.nextDouble() * totalWeight;
            for (Weight<T> w : list)
            {
                value -= w.weight;
                if (value <= 0)
                    return w.value;
            }
            return list.get(list.size() - 1).value;
        };
    }

    private static boolean isIterationSafe(Map<?, ?> values)
    {
        return values instanceof EnumMap
                || values instanceof LinkedHashMap
                || values instanceof SortedMap;
    }

    private static boolean isIterationSafe(Set<?> values)
    {
        return values instanceof EnumSet
               || values instanceof LinkedHashSet
               || values instanceof SortedSet;
    }

    private static void checkIterationSafe(String name, Map<?, ?> values)
    {
        if (!isIterationSafe(values))
            throw new IllegalArgumentException(name + " requires a map with deterministic iteration; given " + values.getClass());
    }

    /**
     * Creates a generator that selects values from an array using a Zipfian distribution.
     * <p>
     * In a Zipfian distribution, the probability of selecting an element decreases as its index increases.
     * The first element has the highest probability of being selected, the second element has the second-highest
     * probability, and so on. This is useful for generating skewed data where some values are much more common than
     * others.
     * <p>
     * The weights follow a power law where {@code weight = base / (i+1)}, creating a harmonic series.
     *
     * @see <a href="https://en.wikipedia.org/wiki/Zipf%27s_law">Wikipedia</a>
     * @param array the array to select from
     * @return a generator that produces values from the array following a Zipfian distribution
     * @throws IllegalArgumentException if the array is null or empty
     */
    public static Gen.IntGen pickZipf(int[] array)
    {
        if (array == null || array.length == 0)
            throw new IllegalArgumentException("Empty array given");
        if (array.length == 1)
            return ignore -> array[0];
        BigDecimal[] weights = new BigDecimal[array.length];
        BigDecimal base = BigDecimal.valueOf(Math.pow(2, array.length));
        weights[0] = base;
        for (int i = 1; i < array.length; i++)
            weights[i] = base.divide(BigDecimal.valueOf(i + 1), RoundingMode.UP);
        BigDecimal totalWeights = Stream.of(weights).reduce(BigDecimal.ZERO, BigDecimal::add);

        return rs -> {
            BigDecimal value = BigDecimal.valueOf(rs.nextDouble()).multiply(totalWeights);
            for (int i = 0; i < weights.length; i++)
            {
                value = value.subtract(weights[i]);
                if (value.compareTo(BigDecimal.ZERO) <= 0)
                    return array[i];
            }
            return array[array.length - 1];
        };
    }

    /**
     * Creates a generator that selects values from an array using a Zipfian distribution.
     * <p>
     * In a Zipfian distribution, the probability of selecting an element decreases as its index increases.
     * The first element has the highest probability of being selected, the second element has the second-highest
     * probability, and so on. This is useful for generating skewed data where some values are much more common than
     * others.
     * <p>
     * The weights follow a power law where {@code weight = base / (i+1)}, creating a harmonic series.
     *
     * @see <a href="https://en.wikipedia.org/wiki/Zipf%27s_law">Wikipedia</a>
     * @param array the array to select from
     * @return a generator that produces values from the array following a Zipfian distribution
     * @throws IllegalArgumentException if the array is null or empty
     */
    public static Gen.LongGen pickZipf(long[] array)
    {
        if (array == null || array.length == 0)
            throw new IllegalArgumentException("Empty array given");
        if (array.length == 1)
            return ignore -> array[0];
        BigDecimal[] weights = new BigDecimal[array.length];
        BigDecimal base = BigDecimal.valueOf(Math.pow(2, array.length));
        weights[0] = base;
        for (int i = 1; i < array.length; i++)
            weights[i] = base.divide(BigDecimal.valueOf(i + 1), RoundingMode.UP);
        BigDecimal totalWeights = Stream.of(weights).reduce(BigDecimal.ZERO, BigDecimal::add);

        return rs -> {
            BigDecimal value = BigDecimal.valueOf(rs.nextDouble()).multiply(totalWeights);
            for (int i = 0; i < weights.length; i++)
            {
                value = value.subtract(weights[i]);
                if (value.compareTo(BigDecimal.ZERO) <= 0)
                    return array[i];
            }
            return array[array.length - 1];
        };
    }

    /**
     * Creates a generator that selects values from an array using a Zipfian distribution.
     * <p>
     * In a Zipfian distribution, the probability of selecting an element decreases as its index increases.
     * The first element has the highest probability of being selected, the second element has the second-highest
     * probability, and so on. This is useful for generating skewed data where some values are much more common than
     * others.
     * <p>
     * The weights follow a power law where {@code weight = base / (i+1)}, creating a harmonic series.
     *
     * @see <a href="https://en.wikipedia.org/wiki/Zipf%27s_law">Wikipedia</a>
     * @param array the array to select from
     * @return a generator that produces values from the array following a Zipfian distribution
     * @throws IllegalArgumentException if the array is null or empty
     */
    @SafeVarargs
    public static <T> Gen<T> pickZipf(T... array)
    {
        return pickZipf(Arrays.asList(array));
    }

    /**
     * Creates a generator that selects values from a list using a Zipfian distribution.
     * <p>
     * In a Zipfian distribution, the probability of selecting an element decreases as its index increases.
     * The first element has the highest probability of being selected, the second element has the second-highest
     * probability, and so on. This is useful for generating skewed data where some values are much more common than
     * others.
     * <p>
     * The weights follow a power law where {@code weight = base / (i+1)}, creating a harmonic series.
     *
     * @see <a href="https://en.wikipedia.org/wiki/Zipf%27s_law">Wikipedia</a>
     * @param array the list to select from
     * @return a generator that produces values from the list following a Zipfian distribution
     * @throws IllegalArgumentException if the list is null or empty
     */
    public static <T> Gen<T> pickZipf(List<T> array)
    {
        if (array == null || array.isEmpty())
            throw new IllegalArgumentException("Empty array given");
        if (array.size() == 1)
            return ignore -> array.get(0);
        BigDecimal[] weights = new BigDecimal[array.size()];
        BigDecimal base = BigDecimal.valueOf(Math.pow(2, array.size()));
        weights[0] = base;
        for (int i = 1; i < array.size(); i++)
            weights[i] = base.divide(BigDecimal.valueOf(i + 1), RoundingMode.UP);
        BigDecimal totalWeights = Stream.of(weights).reduce(BigDecimal.ZERO, BigDecimal::add);

        return rs -> {
            BigDecimal value = BigDecimal.valueOf(rs.nextDouble()).multiply(totalWeights);
            for (int i = 0; i < weights.length; i++)
            {
                value = value.subtract(weights[i]);
                if (value.compareTo(BigDecimal.ZERO) <= 0)
                    return array.get(i);
            }
            return array.get(array.size() - 1);
        };
    }

    /**
     * Creates a generator that produces a subset of the input list.
     * <p>
     * Elements are selected randomly from the input list without modifying it; the subset returned may not reflect
     * the original order.
     *
     * @param <T> the type of elements in the list
     * @param input the source list to select elements from
     * @return a generator that produces variable-sized random subsets of the input list
     * @throws IndexOutOfBoundsException if sizeGen produces a value outside the list's range
     */
    public static <T> Gen<List<T>> select(List<T> input)
    {
        return select(input, rs -> rs.nextInt(0, input.size() + 1));
    }

    /**
     * Creates a generator that produces a subset of the input list.
     * <p>
     * Elements are selected randomly from the input list without modifying it; the subset returned may not reflect
     * the original order.
     *
     * @param <T> the type of elements in the list
     * @param input the source list to select elements from
     * @param size fixed length of how many elements to return
     * @return a generator that produces fixed-sized random subsets of the input list
     * @throws IndexOutOfBoundsException if sizeGen produces a value outside the list's range
     */
    public static <T> Gen<List<T>> select(List<T> input, int size)
    {
        return select(input, ignore -> size);
    }

    /**
     * Creates a generator that produces a subset of the input list.
     * <p>
     * Elements are selected randomly from the input list without modifying it; the subset returned may not reflect
     * the original order.
     *
     * @param <T> the type of elements in the list
     * @param input the source list to select elements from
     * @param sizeGen a generator that determines how many elements to select
     * @return a generator that produces variable-sized random subsets of the input list
     * @throws IndexOutOfBoundsException if sizeGen produces a value outside the list's range
     */
    public static <T> Gen<List<T>> select(List<T> input, Gen.IntGen sizeGen)
    {
        if (input.isEmpty()) return constant(input);
        return rs -> {
            int size = sizeGen.nextInt(rs);
            if (size == 0)
                return List.of();
            Invariants.requireIndexInBounds(input.size(), 0, size);
            List<T> remaining = new ArrayList<>(input);
            List<T> list = new ArrayList<>(size);
            for (int i = 0; i < size; i++)
                list.add(remaining.remove(rs.nextInt(0, remaining.size())));
            return list;
        };
    }

    public static Gen<Gen.IntGen> randomWeights(int[] array)
    {
        return rs -> {
            float[] weights = Picker.randomWeights(rs, array.length);
            return r -> array[index(r, weights)];
        };
    }

    public static Gen<Gen.LongGen> randomWeights(long[] array)
    {
        return rs -> {
            float[] weights = Picker.randomWeights(rs, array.length);
            return r -> array[index(r, weights)];
        };
    }

    public static <T> Gen<Gen<T>> randomWeights(T[] array)
    {
        return rs -> {
            float[] weights = Picker.randomWeights(rs, array.length);
            return r -> array[index(r, weights)];
        };
    }

    public static <T> Gen<Gen<T>> randomWeights(List<T> array)
    {
        return rs -> {
            float[] weights = Picker.randomWeights(rs, array.size());
            return r -> array.get(index(r, weights));
        };
    }

    private static int index(RandomSource rs, float[] weights)
    {
        int i = Arrays.binarySearch(weights, rs.nextFloat());
        if (i < 0) i = -1 - i;
        return i;
    }

    /**
     * Creates a generator that uses different distribution strategies for selecting items from a range.  The range is
     * broken up into {@code numBuckets} sub ranges, and the bucket will be selected using the distribution strategies.
     * <p>
     * When the top level generator is called it selects what distribution to use, and returns a generator of {@code int}
     * that selects from the input using that distribution.
     * <p>
     * This method does not document what distributions it can select from, that way it is free to change (add/remove)
     * over time; it only maintains the property that the distribution may change between calls to the top level generator.
     * <p>
     * Some sample distribution strategies:
     * <ul>
     *     <li>Uniform distribution - Each item has equal chance of selection</li>
     *     <li>Zipfian distribution - Early items are heavily favored, with probability decreasing by index</li>
     * </ul>
     *
     * @param minInclusive value inclusive
     * @param maxExclusive value exclusive
     * @param numBuckets how many sub-ranges to create for the top level range
     * @return a generator of different distribution strategies
     */
    public static Gen<Gen.IntGen> mixedDistribution(int minInclusive, int maxExclusive, int numBuckets)
    {
        int domainSize = (maxExclusive - minInclusive);
        if (domainSize < 0)
            throw new IllegalArgumentException("Range is too large; min=" + minInclusive + ", max=" + maxExclusive);
        if (numBuckets <= 0 || numBuckets > domainSize)
            throw new IllegalArgumentException("Num buckets must be between 1 and " + domainSize + "; given " + numBuckets);
        int[] bucket, indexes;
        bucket = new int[numBuckets];
        int delta = domainSize / numBuckets;
        for (int i = 0; i < numBuckets; i++)
            bucket[i] = minInclusive + i * delta;
        indexes = IntStream.range(0, bucket.length).toArray();
        Gen<Gen.IntGen> indexDistro = mixedDistribution(indexes);
        return rs -> {
            Gen.IntGen indexGen = indexDistro.next(rs);
            switch (rs.nextInt(0, 2))
            {
                case 0: // uniform
                {
                    return r -> {
                        int idx = indexGen.next(rs);
                        int start = bucket[idx];
                        int end = idx == bucket.length - 1 ? maxExclusive : bucket[idx + 1];
                        return r.nextInt(start, end);
                    };
                }
                case 1: // median biased
                {
                    int[] medians = new int[bucket.length];
                    for (int i = 0; i < medians.length; i++)
                    {
                        int start = bucket[i];
                        int end = i == bucket.length - 1 ? maxExclusive : bucket[i + 1];
                        medians[i] = rs.nextInt(start, end);
                    }
                    return r -> {
                        int idx = indexGen.next(rs);
                        int start = bucket[idx];
                        int end = idx == bucket.length - 1 ? maxExclusive : bucket[idx + 1];
                        int median = medians[idx];
                        return r.nextBiasedInt(start, median, end);
                    };
                }
                default:
                    throw new AssertionError();
            }
        };
    }

    /**
     * Creates a generator that uses different distribution strategies for selecting items from a range.
     * <p>
     * When the top level generator is called it selects what distribution to use, and returns a generator of {@code int}
     * that selects from the input using that distribution.
     * <p>
     * This method does not document what distributions it can select from, that way it is free to change (add/remove)
     * over time; it only maintains the property that the distribution may change between calls to the top level generator.
     * <p>
     * Some sample distribution strategies:
     * <ul>
     *     <li>Uniform distribution - Each item has equal chance of selection</li>
     *     <li>Zipfian distribution - Early items are heavily favored, with probability decreasing by index</li>
     * </ul>
     *
     * @param minInclusive value inclusive
     * @param maxExclusive value exclusive
     * @return a generator of different distribution strategies
     */
    public static Gen<Gen.IntGen> mixedDistribution(int minInclusive, int maxExclusive)
    {
        int domainSize = (maxExclusive - minInclusive + 1);
        if (domainSize < 0)
            throw new IllegalArgumentException("Range is too large; min=" + minInclusive + ", max=" + maxExclusive);
        int[] array, indexes;
        if (domainSize > 200) // randomly selected
        {
            int numBuckets = 10;
            int delta = domainSize / numBuckets;
            array = new int[numBuckets];
            for (int i = 0; i < numBuckets; i++)
                array[i] = minInclusive + i * delta;
            indexes = IntStream.range(0, array.length).toArray();
        }
        else
        {
            array = IntStream.range(minInclusive, maxExclusive).toArray();
            indexes = null;
        }
        return rs -> {
            switch (rs.nextInt(0, 4))
            {
                case 0: // uniform
                    return r -> r.nextInt(minInclusive, maxExclusive);
                case 1: // median biased
                    int median = rs.nextInt(minInclusive, maxExclusive);
                    return r -> r.nextBiasedInt(minInclusive, median, maxExclusive);
                case 2: // zipf
                    if (indexes == null)
                        return Gens.pickZipf(rs.nextBoolean() ? reverseAndCopy(array) : array);
                    return Gens.pickZipf(rs.nextBoolean() ? reverseAndCopy(indexes) : indexes).mapAsInt((r, index) -> {
                        int start = array[index];
                        int end = index == array.length - 1 ? maxExclusive : array[index + 1];
                        return r.nextInt(start, end);
                    });
                case 3: // random weight
                    if (indexes == null)
                        return randomWeights(array).next(rs);
                    return randomWeights(indexes).next(rs).mapAsInt((r, index) -> {
                        int start = array[index];
                        int end = index == array.length - 1 ? maxExclusive : array[index + 1];
                        return r.nextInt(start, end);
                    });
                default:
                    throw new AssertionError();
            }
        };
    }

    private static int[] reverseAndCopy(int[] array)
    {
        array = Arrays.copyOf(array, array.length);
        for (int i = 0, mid = array.length / 2, j = array.length - 1; i < mid; i++, j--)
        {
            int tmp = array[i];
            array[i] = array[j];
            array[j] = tmp;
        }
        return array;
    }

    /**
     * Creates a generator that uses different distribution strategies for selecting items from a range.
     * <p>
     * When the top level generator is called it selects what distribution to use, and returns a generator of {@code int}
     * that selects from the input using that distribution.
     * <p>
     * This method does not document what distributions it can select from, that way it is free to change (add/remove)
     * over time; it only maintains the property that the distribution may change between calls to the top level generator.
     * <p>
     * Some sample distribution strategies:
     * <ul>
     *     <li>Uniform distribution - Each item has equal chance of selection</li>
     *     <li>Zipfian distribution - Early items are heavily favored, with probability decreasing by index</li>
     * </ul>
     *
     * @param minInclusive value inclusive
     * @param maxExclusive value exclusive
     * @return a generator of different distribution strategies
     */
    public static Gen<Gen.LongGen> mixedDistribution(long minInclusive, long maxExclusive)
    {
        long domainSize = (maxExclusive - minInclusive + 1);
        if (domainSize < 0)
            throw new IllegalArgumentException("Range is too large; min=" + minInclusive + ", max=" + maxExclusive);
        long[] array;
        int[] indexes;
        if (domainSize > 200) // randomly selected
        {
            int numBuckets = 10;
            long delta = domainSize / numBuckets;
            array = new long[numBuckets];
            for (int i = 0; i < numBuckets; i++)
                array[i] = minInclusive + i * delta;
            indexes = IntStream.range(0, array.length).toArray();
        }
        else
        {
            array = LongStream.range(minInclusive, maxExclusive).toArray();
            indexes = null;
        }
        return rs -> {
            switch (rs.nextInt(0, 4))
            {
                case 0: // uniform
                    return r -> r.nextLong(minInclusive, maxExclusive);
                case 1: // median biased
                    long median = rs.nextLong(minInclusive, maxExclusive);
                    return r -> r.nextBiasedLong(minInclusive, median, maxExclusive);
                case 2: // zipf
                    if (indexes == null)
                        return Gens.pickZipf(rs.nextBoolean() ? reverseAndCopy(array) : array);
                    return Gens.pickZipf(rs.nextBoolean() ? reverseAndCopy(indexes) : indexes).mapAsLong((r, index) -> {
                        long start = array[index];
                        long end = index == array.length - 1 ? maxExclusive : array[index + 1];
                        return r.nextLong(start, end);
                    });
                case 3: // random weight
                    if (indexes == null)
                        return randomWeights(array).next(rs);
                    return randomWeights(indexes).next(rs).mapAsLong((r, index) -> {
                        long start = array[index];
                        long end = index == array.length - 1 ? maxExclusive : array[index + 1];
                        return r.nextLong(start, end);
                    });
                default:
                    throw new AssertionError();
            }
        };
    }

    private static long[] reverseAndCopy(long[] array)
    {
        array = Arrays.copyOf(array, array.length);
        for (int i = 0, mid = array.length / 2, j = array.length - 1; i < mid; i++, j--)
        {
            long tmp = array[i];
            array[i] = array[j];
            array[j] = tmp;
        }
        return array;
    }

    /**
     * Creates a generator that uses different distribution strategies for selecting items from an array.
     * <p>
     * When the top level generator is called it selects what distribution to use, and returns a generator of {@code T}
     * that selects from the input using that distribution.
     * <p>
     * This method does not document what distributions it can select from, that way it is free to change (add/remove)
     * over time; it only maintains the property that the distribution may change between calls to the top level generator.
     * <p>
     * Some sample distribution strategies:
     * <ul>
     *     <li>Uniform distribution - Each item has equal chance of selection</li>
     *     <li>Zipfian distribution - Early items are heavily favored, with probability decreasing by index</li>
     * </ul>
     *
     * @param <T> the type of elements in the array
     * @param list the array of values to select from
     * @return a generator of different distribution strategies
     */
    @SafeVarargs
    public static <T> Gen<Gen<T>> mixedDistribution(T... list)
    {
        return mixedDistribution(Arrays.asList(list));
    }

    /**
     * Creates a generator that uses different distribution strategies for selecting items from a list.
     * <p>
     * When the top level generator is called it selects what distribution to use, and returns a generator of {@code T}
     * that selects from the input using that distribution.
     * <p>
     * This method does not document what distributions it can select from, that way it is free to change (add/remove)
     * over time; it only maintains the property that the distribution may change between calls to the top level generator.
     * <p>
     * Some sample distribution strategies:
     * <ul>
     *     <li>Uniform distribution - Each item has equal chance of selection</li>
     *     <li>Zipfian distribution - Early items are heavily favored, with probability decreasing by index</li>
     * </ul>
     *
     * @param <T> the type of elements in the list
     * @param list the list of values to select from
     * @return a generator of different distribution strategies
     */
    public static <T> Gen<Gen<T>> mixedDistribution(List<T> list)
    {
        Invariants.require(!list.isEmpty(), "can not pick from an empty collection");
        if (list.size() == 1)
            return i -> constant(list.get(0));
        return rs -> {
            switch (rs.nextInt(0, 4))
            {
                case 0: // uniform
                    return r -> list.get(rs.nextInt(0, list.size()));
                case 1: // median biased
                    int median = rs.nextInt(0, list.size());
                    return r -> list.get(r.nextBiasedInt(0, median, list.size()));
                case 2: // zipf
                    List<T> array = list;
                    if (rs.nextBoolean())
                    {
                        array = new ArrayList<>(list);
                        Collections.reverse(array);
                    }
                    return pickZipf(array);
                case 3: // random weight
                    return randomWeights(list).next(rs);
                default:
                    throw new AssertionError();
            }
        };
    }

    /**
     * Creates a generator that uses different distribution strategies for selecting items from an array.
     * <p>
     * When the top level generator is called it selects what distribution to use, and returns a generator of {@code int}
     * that selects from the input using that distribution.
     * <p>
     * This method does not document what distributions it can select from, that way it is free to change (add/remove)
     * over time; it only maintains the property that the distribution may change between calls to the top level generator.
     * <p>
     * Some sample distribution strategies:
     * <ul>
     *     <li>Uniform distribution - Each item has equal chance of selection</li>
     *     <li>Zipfian distribution - Early items are heavily favored, with probability decreasing by index</li>
     * </ul>
     *
     * @param list the array of values to select from
     * @return a generator of different distribution strategies
     */
    public static Gen<Gen.IntGen> mixedDistribution(int[] list)
    {
        return rs -> {
            switch (rs.nextInt(0, 4))
            {
                case 0: // uniform
                    return r -> list[rs.nextInt(0, list.length)];
                case 1: // median biased
                    int median = rs.nextInt(0, list.length);
                    return r -> list[r.nextBiasedInt(0, median, list.length)];
                case 2: // zipf
                    int[] array = list;
                    if (rs.nextBoolean())
                    {
                        array = Arrays.copyOf(array, array.length);
                        Utils.reverse(array);
                    }
                    return pickZipf(array);
                case 3: // random weight
                    return randomWeights(list).next(rs);
                default:
                    throw new AssertionError();
            }
        };
    }

    public static Gen<char[]> charArray(Gen.IntGen sizes, char[] domain)
    {
        return charArray(sizes, domain, (a, b) -> true);
    }

    public interface IntCharBiPredicate
    {
        boolean test(int a, char b);
    }

    public static Gen<char[]> charArray(Gen.IntGen sizes, char[] domain, IntCharBiPredicate fn)
    {
        Gen.IntGen indexGen = ints().between(0, domain.length - 1);
        return rs -> {
            int size = sizes.nextInt(rs);
            char[] is = new char[size];
            for (int i = 0; i != size; i++)
            {
                char c;
                do
                {
                    c = domain[indexGen.nextInt(rs)];
                }
                while (!fn.test(i, c));
                is[i] = c;
            }
            return is;
        };
    }

    public static Gen<RandomSource> random() {
        return r -> r;
    }

    public static BooleanDSL bools()
    {
        return new BooleanDSL();
    }

    public static IntDSL ints()
    {
        return new IntDSL();
    }

    public static LongDSL longs() {
        return new LongDSL();
    }

    public static <T> ListDSL<T> lists(Gen<T> fn) {
        return new ListDSL<>(fn);
    }

    public static <T> ArrayDSL<T> arrays(Class<T> type, Gen<T> fn) {
        return new ArrayDSL<>(type, fn);
    }

    public static IntArrayDSL arrays(Gen.IntGen fn) {
        return new IntArrayDSL(fn);
    }

    public static LongArrayDSL arrays(Gen.LongGen fn) {
        return new LongArrayDSL(fn);
    }

    public static EnumDSL enums()
    {
        return new EnumDSL();
    }

    public static StringDSL strings()
    {
        return new StringDSL();
    }

    public static BooleanSupplier supplier(Gen<Boolean> gen, RandomSource rs)
    {
        return () -> gen.next(rs);
    }

    public static class BooleanDSL
    {
        public Gen<Boolean> all()
        {
            return RandomSource::nextBoolean;
        }

        /**
         * Creates a generator that produces boolean values with a target ratio of true values,
         * and with true values appearing in contiguous runs of random length.
         * <p>
         * This is a stateful generator (not thread safe) that maintains an approximate ratio of true values by
         * tracking the counts of true and false values it has generated. It will produce runs of consecutive true
         * values up to the specified maximum length to create a more realistic pattern than
         * independent random choices.
         *
         * @param ratio the target ratio of true values (between 0 and 1)
         * @param maxRuns the maximum length of consecutive true values in a run
         * @return a generator that produces boolean values matching the specified distribution pattern
         * @throws IllegalArgumentException if ratio is not between 0 and 1
         */
        public Gen<Boolean> biasedRepeatingRuns(double ratio, int maxRuns)
        {
            Invariants.requireArgument(ratio > 0 && ratio <= 1, "Expected %d to be larger than 0 and <= 1", ratio);
            double lower = ratio * .8;
            double upper = ratio * 1.2;
            return new Gen<>()
            {
                // run represents how many consecutaive true values should be returned; -1 implies no active "run" exists
                private int run = -1;
                private long falseCount = 0, trueCount = 0;

                @Override
                public Boolean next(RandomSource rs)
                {
                    if (run != -1)
                    {
                        run--;
                        trueCount++;
                        return true;
                    }
                    double currentRatio = trueCount / (double) (falseCount + trueCount);
                    if (currentRatio < lower)
                    {
                        // not enough true
                        trueCount++;
                        return true;
                    }
                    if (currentRatio > upper)
                    {
                        // not enough false
                        falseCount++;
                        return false;
                    }
                    if (rs.decide(ratio))
                    {
                        run = rs.nextInt(maxRuns);
                        run--;
                        trueCount++;
                        return true;
                    }
                    falseCount++;
                    return false;
                }
            };
        }

        public Gen<Gen<Boolean>> mixedDistribution()
        {
            return rs -> {
                int selection = rs.nextInt(0, 4);
                switch (selection)
                {
                    case 0: // uniform 50/50
                        return RandomSource::nextBoolean;
                    case 1: // variable frequency
                        var freq = rs.nextFloat();
                        return r -> r.decide(freq);
                    case 2: // fixed result
                        boolean result = rs.nextBoolean();
                        return ignore -> result;
                    case 3: // biased repeating runs
                        return biasedRepeatingRuns(rs.nextDouble(), rs.nextInt(1, 100));
                    default:
                        throw new IllegalStateException("Unexpected int for bool selection: " + selection);
                }
            };
        }
    }

    public static class IntDSL
    {
        public Gen.IntGen of(int value)
        {
            return r -> value;
        }

        public Gen.IntGen all()
        {
            return RandomSource::nextInt;
        }

        public Gen.IntGen between(int min, int max)
        {
            Invariants.requireArgument(max >= min, "max (%d) < min (%d)", max, min);
            if (min == max)
                return of(min);
            // since bounds is exclusive, if max == max_value unable to do +1 to include... so will return a gen
            // that does not include
            if (max == Integer.MAX_VALUE)
                return r -> r.nextInt(min, max);
            return r -> r.nextInt(min, max + 1);
        }


        /**
         * @see Gens#mixedDistribution(int, int)
         */
        public Gen<Gen.IntGen> mixedDistribution(int minInclusive, int maxExclusive)
        {
            return Gens.mixedDistribution(minInclusive, maxExclusive);
        }

        /**
         * @see Gens#mixedDistribution(int, int, int)
         */
        public Gen<Gen.IntGen> mixedDistribution(int minInclusive, int maxExclusive, int numBuckets)
        {
            return Gens.mixedDistribution(minInclusive, maxExclusive, numBuckets);
        }
    }

    public static class LongDSL {
        public Gen.LongGen of(long value)
        {
            return r -> value;
        }

        public Gen.LongGen all() {
            return RandomSource::nextLong;
        }

        public Gen.LongGen between(long min, long max) {
            Invariants.requireArgument(max >= min);
            if (min == max)
                return of(min);
            // since bounds is exclusive, if max == max_value unable to do +1 to include... so will return a gen
            // that does not include
            if (max == Long.MAX_VALUE)
                return r -> r.nextLong(min, max);
            return r -> r.nextLong(min, max + 1);
        }
    }

    public static class EnumDSL
    {
        public <T extends Enum<T>> Gen<T> all(Class<T> klass)
        {
            return pick(klass.getEnumConstants());
        }

        /**
         * A mixedDistribution using the enum classes values as input to {@link Gens#mixedDistribution(Object[])}
         *
         * @see Gens#mixedDistribution(Object[])
         */
        public <T extends Enum<T>> Gen<Gen<T>> allMixedDistribution(Class<T> klass)
        {
            return mixedDistribution(klass.getEnumConstants());
        }

        /**
         * Creates a generator that randomly selects one value from the enum values based off a set of weights
         *
         * @see #pick(Map)
         */
        public <T extends Enum<T>> Gen<T> allWithWeights(Class<T> klass, int... weights)
        {
            T[] constants = klass.getEnumConstants();
            if (constants.length != weights.length)
                throw new IllegalArgumentException(String.format("Total number of weights (%s) does not match the enum (%s)", Arrays.toString(weights), Arrays.toString(constants)));
            Map<T, Integer> values = new EnumMap<>(klass);
            for (int i = 0; i < constants.length; i++)
                values.put(constants[i], weights[i]);
            return pick(values);
        }
    }

    public static class StringDSL
    {
        public Gen<String> of(Gen.IntGen sizes, char[] domain)
        {
            // note, map is overloaded so String::new is ambugious to javac, so need a lambda here
            return charArray(sizes, domain).map(c -> new String(c));
        }

        public SizeBuilder<String> of(char[] domain)
        {
            return new SizeBuilder<>(sizes -> of(sizes, domain));
        }

        public Gen<String> of(Gen.IntGen sizes, char[] domain, IntCharBiPredicate fn)
        {
            // note, map is overloaded so String::new is ambugious to javac, so need a lambda here
            return charArray(sizes, domain, fn).map(c -> new String(c));
        }

        public SizeBuilder<String> of(char[] domain, IntCharBiPredicate fn)
        {
            return new SizeBuilder<>(sizes -> of(sizes, domain, fn));
        }

        public Gen<String> all(Gen.IntGen sizes)
        {
            return betweenCodePoints(sizes, Character.MIN_CODE_POINT, Character.MAX_CODE_POINT);
        }

        public SizeBuilder<String> all()
        {
            return new SizeBuilder<>(this::all);
        }

        public Gen<String> ascii(Gen.IntGen sizes)
        {
            return betweenCodePoints(sizes, 0, 127);
        }

        public SizeBuilder<String> ascii()
        {
            return new SizeBuilder<>(this::ascii);
        }

        public Gen<String> betweenCodePoints(Gen.IntGen sizes, int min, int max)
        {
            Gen.IntGen codePointGen = ints().between(min, max).filter(Character::isDefined);
            return rs -> {
                int[] array = new int[sizes.nextInt(rs)];
                for (int i = 0; i < array.length; i++)
                    array[i] = codePointGen.nextInt(rs);
                return new String(array, 0, array.length);
            };
        }

        public SizeBuilder<String> betweenCodePoints(int min, int max)
        {
            return new SizeBuilder<>(sizes -> betweenCodePoints(sizes, min, max));
        }
    }

    public static class SizeBuilder<T>
    {
        private final Function<Gen.IntGen, Gen<T>> fn;

        public SizeBuilder(Function<Gen.IntGen, Gen<T>> fn)
        {
            this.fn = fn;
        }

        public Gen<T> ofLength(int fixed)
        {
            return ofLengthBetween(fixed, fixed);
        }

        public Gen<T> ofLengthBetween(int min, int max)
        {
            return fn.apply(ints().between(min, max));
        }
    }

    public static class ListDSL<T> implements BaseSequenceDSL<ListDSL<T>, List<T>> {
        private final Gen<T> fn;

        public ListDSL(Gen<T> fn) {
            this.fn = Objects.requireNonNull(fn);
        }

        @Override
        public ListDSL<T> unique()
        {
            return new ListDSL<>(new GenReset<>(fn, false));
        }

        public ListDSL<T> uniqueBestEffort()
        {
            return new ListDSL<>(new GenReset<>(fn, true));
        }

        @Override
        public Gen<List<T>> ofSizeBetween(int minSize, int maxSize) {
            Gen.IntGen sizeGen = ints().between(minSize, maxSize);
            return r ->
            {
                Reset.tryReset(fn);
                int size = sizeGen.nextInt(r);
                List<T> list = new ArrayList<>(size);
                for (int i = 0; i < size; i++)
                {
                    try
                    {
                        list.add(fn.next(r));
                    }
                    catch (IgnoreGenResult e)
                    {
                        // ignore
                    }
                }
                return list;
            };
        }
    }

    public static class ArrayDSL<T> implements BaseSequenceDSL<ArrayDSL<T>, T[]>
    {
        private final Class<T> type;
        private final Gen<T> fn;

        public ArrayDSL(Class<T> type, Gen<T> fn)
        {
            this.type = Objects.requireNonNull(type);
            this.fn = Objects.requireNonNull(fn);
        }

        @Override
        public ArrayDSL<T> unique()
        {
            return new ArrayDSL<>(type, new GenReset<>(fn, false));
        }

        public ArrayDSL<T> uniqueBestEffort()
        {
            return new ArrayDSL<>(type, new GenReset<>(fn, true));
        }

        @Override
        public Gen<T[]> ofSizeBetween(int minSize, int maxSize) {
            Gen.IntGen sizeGen = ints().between(minSize, maxSize);
            return r ->
            {
                Reset.tryReset(fn);
                int size = sizeGen.nextInt(r);
                //noinspection unchecked
                T[] list = (T[]) Array.newInstance(type, size);
                for (int i = 0; i < size; i++)
                    list[i] = fn.next(r);
                return list;
            };
        }
    }

    public static class IntArrayDSL implements BaseSequenceDSL<IntArrayDSL, int[]> {
        private final Gen.IntGen fn;

        public IntArrayDSL(Gen.IntGen fn) {
            this.fn = Objects.requireNonNull(fn);
        }

        @Override
        public IntArrayDSL unique()
        {
            return new IntArrayDSL(new IntGenReset(fn));
        }

        @Override
        public Gen<int[]> ofSizeBetween(int minSize, int maxSize) {
            Gen.IntGen sizeGen = ints().between(minSize, maxSize);
            return r ->
            {
                int size = sizeGen.nextInt(r);
                int[] list = new int[size];
                for (int i = 0; i < size; i++)
                    list[i] = fn.nextInt(r);
                return list;
            };
        }
    }

    public static class LongArrayDSL implements BaseSequenceDSL<LongArrayDSL, long[]> {
        private final Gen.LongGen fn;

        public LongArrayDSL(Gen.LongGen fn) {
            this.fn = Objects.requireNonNull(fn);
        }

        @Override
        public LongArrayDSL unique()
        {
            return new LongArrayDSL(new LongGenReset(fn));
        }

        @Override
        public Gen<long[]> ofSizeBetween(int minSize, int maxSize) {
            Gen.IntGen sizeGen = ints().between(minSize, maxSize);
            return r ->
            {
                int size = sizeGen.nextInt(r);
                long[] list = new long[size];
                for (int i = 0; i < size; i++)
                    list[i] = fn.nextLong(r);
                return list;
            };
        }
    }

    public interface BaseSequenceDSL<A extends BaseSequenceDSL<A, B>, B>
    {
        A unique();

        Gen<B> ofSizeBetween(int min, int max);

        default Gen<B> ofSize(int size) {
            return ofSizeBetween(size, size);
        }
    }

    protected interface Reset {
        static void tryReset(Object o)
        {
            if (o instanceof Reset)
                ((Reset) o).reset();
        }

        void reset();
    }

    private static final class IgnoreGenResult extends RuntimeException
    {
        private static final IgnoreGenResult INSTANCE = new IgnoreGenResult();
        private IgnoreGenResult()
        {
            super(null, null, false, false);
        }
    }

    private static class GenReset<T> implements Gen<T>, Reset
    {
        private final Set<T> seen = new HashSet<>();
        private final Gen<T> fn;
        private final boolean bestEffort;

        private GenReset(Gen<T> fn, boolean bestEffort)
        {
            this.fn = fn;
            this.bestEffort = bestEffort;
        }

        @Override
        public T next(RandomSource random)
        {
            if (!bestEffort)
            {
                T value;
                // 10k attempts
                for (int i = 0; i < 10_000; i++)
                {
                    if (seen.add((value = fn.next(random))))
                        return value;
                }

                throw new IllegalArgumentException("Could not generate a unique value after 10k attempts");
            }
            else
            {
                T value = null;
                int i;
                //noinspection StatementWithEmptyBody
                for (i = 0; i < 42 && !seen.add((value = fn.next(random))); i++) {}
                if (i == 42) throw IgnoreGenResult.INSTANCE;
                return value;
            }
        }

        @Override
        public void reset()
        {
            seen.clear();
        }
    }

    private static class IntGenReset implements Gen.IntGen, Reset
    {
        private final GenReset<Integer> base;

        private IntGenReset(Gen.IntGen fn)
        {
            this.base = new GenReset<>(fn, false);
        }
        @Override
        public int nextInt(RandomSource random) {
            return base.next(random);
        }

        @Override
        public void reset() {
            base.reset();
        }
    }

    private static class LongGenReset implements Gen.LongGen, Reset
    {
        private final GenReset<Long> base;

        private LongGenReset(Gen.LongGen fn)
        {
            this.base = new GenReset<>(fn, false);
        }
        @Override
        public long nextLong(RandomSource random) {
            return base.next(random);
        }

        @Override
        public void reset() {
            base.reset();
        }
    }

    private static class Weight<T> implements Comparable<Weight<T>>
    {
        private final T value;
        private final double weight;
        private final int index;

        private Weight(T value, double weight, int index) {
            this.value = value;
            this.weight = weight;
            this.index = index;
        }

        @Override
        public int compareTo(Weight<T> o) {
            int rc = Double.compare(weight, o.weight);
            if (rc == 0)
                rc = Integer.compare(index, o.index);
            return rc;
        }
    }
}
