# Source code for dask.bag.random

    from __future__ import annotations

    import heapq
    import math
    import random as rnd
    from functools import partial
    from itertools import islice

    from dask.bag.core import Bag

    def sample(population, k, split_every=None):
        """Chooses k unique random elements from a bag.

        Returns a new bag containing elements from the population while
        leaving the original population unchanged.

        Parameters
        ----------
        population: Bag
            Elements to sample.
        k: integer, optional
            Number of elements to sample.
        split_every: int (optional)
            Group partitions into groups of this size while performing reduction.
            Defaults to 8.

        Examples
        --------
        >>> import dask.bag as db
        >>> from dask.bag import random
        >>> b = db.from_sequence(range(5), npartitions=2)
        >>> list(random.sample(b, 3).compute())  # doctest: +SKIP
        [1, 3, 5]
        """
        res = _sample(population=population, k=k, split_every=split_every)
        return res.map_partitions(_finalize_sample, k)

    def choices(population, k=1, split_every=None):
        """
        Return a k sized list of elements chosen with replacement.

        Parameters
        ----------
        population: Bag
            Elements to sample.
        k: integer, optional
            Number of elements to sample.
        split_every: int (optional)
            Group partitions into groups of this size while performing reduction.
            Defaults to 8.

        Examples
        --------
        >>> import dask.bag as db
        >>> from dask.bag import random
        >>> b = db.from_sequence(range(5), npartitions=2)
        >>> list(random.choices(b, 3).compute())  # doctest: +SKIP
        [1, 1, 5]
        """
        res = _sample_with_replacement(population=population, k=k, split_every=split_every)
        return res.map_partitions(_finalize_sample, k)

    def _sample_reduce(reduce_iter, k, replace):
        """
        Reduce function used on the sample and choice functions.

        Parameters
        ----------
        reduce_iter : iterable
            Each element is a tuple coming generated by the _sample_map_partitions function.
        replace: bool
            If True, sample with replacement. If False, sample without replacement.

        Returns a sequence of uniformly distributed samples;
        """
        ns_ks = []
        s = []
        n = 0
        # unfolding reduce outputs
        for i in reduce_iter:
            (s_i, n_i) = i
            s.extend(s_i)
            n += n_i
            k_i = len(s_i)
            ns_ks.append((n_i, k_i))

        if k > n and not replace:
            return s, n

        # creating the probability array
        p = []
        for n_i, k_i in ns_ks:
            if k_i > 0:
                p_i = n_i / (k_i * n)
                p += [p_i] * k_i

        sample_func = rnd.choices if replace else _weighted_sampling_without_replacement
        return sample_func(population=s, weights=p, k=k), n

    def _weighted_sampling_without_replacement(population, weights, k):
        """
        Source:
            Weighted random sampling with a reservoir, Pavlos S. Efraimidis, Paul G. Spirakis
        """
        elt = [(math.log(rnd.random()) / weights[i], i) for i in range(len(weights))]
        return [population[x[1]] for x in heapq.nlargest(k, elt)]

    def _sample(population, k, split_every):
        if k < 0:
            raise ValueError("Cannot take a negative number of samples")
        return population.reduction(
            partial(_sample_map_partitions, k=k),
            partial(_sample_reduce, k=k, replace=False),
            out_type=Bag,
            split_every=split_every,
        )

    def _finalize_sample(reduce_iter, k):
        sample = reduce_iter[0]
        if len(sample) < k:
            raise ValueError("Sample larger than population")
        return sample

    def _sample_map_partitions(population, k):
        """
        Reservoir sampling strategy based on the L algorithm
        See https://en.wikipedia.org/wiki/Reservoir_sampling#An_optimal_algorithm
        """

        reservoir, stream_length = [], 0
        stream = iter(population)
        for e in islice(stream, k):
            reservoir.append(e)
            stream_length += 1

        w = math.exp(math.log(rnd.random()) / k)
        nxt = (k - 1) + _geometric(w)

        for i, e in enumerate(stream, k):
            if i == nxt:
                reservoir[rnd.randrange(k)] = e
                w *= math.exp(math.log(rnd.random()) / k)
                nxt += _geometric(w)
            stream_length += 1

        return reservoir, stream_length

    def _sample_with_replacement(population, k, split_every):
        return population.reduction(
            partial(_sample_with_replacement_map_partitions, k=k),
            partial(_sample_reduce, k=k, replace=True),
            out_type=Bag,
            split_every=split_every,
        )

    def _sample_with_replacement_map_partitions(population, k):
        """
        Reservoir sampling with replacement, the main idea is to use k reservoirs of size 1
        See Section Applications in http://utopia.duth.gr/~pefraimi/research/data/2007EncOfAlg.pdf
        """

        stream = iter(population)
        e = next(stream)
        reservoir, stream_length = [e for _ in range(k)], 1

        w = [rnd.random() for _ in range(k)]
        nxt = [_geometric(wi) for wi in w]
        min_nxt = min(nxt)

        for i, e in enumerate(stream, 1):
            if i == min_nxt:
                for j, n in enumerate(nxt):
                    if n == min_nxt:
                        reservoir[j] = e
                        w[j] *= rnd.random()
                        nxt[j] += _geometric(w[j])
                min_nxt = min(nxt)

            stream_length += 1

        return reservoir, stream_length

    def _geometric(p):
        return int(math.log(rnd.uniform(0, 1)) / math.log(1 - p)) + 1