// Copyright © 2024 Mikhail Hogrefe
//
// This file is part of Malachite.
//
// Malachite is free software: you can redistribute it and/or modify it under the terms of the GNU
// Lesser General Public License (LGPL) as published by the Free Software Foundation; either version
// 3 of the License, or (at your option) any later version. See <https://www.gnu.org/licenses/>.

use crate::bools::random::{weighted_random_bools, WeightedRandomBools};
use crate::random::Seed;

/// Generates random [`Option`]s except `None`, with values from a given random iterator.
///
/// This `struct` is created by [`random_somes`]; see its documentation for more.
#[derive(Clone, Debug)]
pub struct RandomSomes<I: Iterator> {
    xs: I,
}

impl<I: Iterator> Iterator for RandomSomes<I> {
    type Item = Option<I::Item>;

    #[inline]
    fn next(&mut self) -> Option<Option<I::Item>> {
        Some(self.xs.next())
    }
}

/// Generates random [`Option`]s except `None`, with values from a given random iterator.
///
/// The values have the same distribution as the values generated by the given iterator: If $Q(x)$
/// is the probability of $x$ being generated by `xs`, then
///
/// $P(\operatorname{Some}(x)) = Q(x)$.
///
/// `xs` must be infinite.
///
/// The output length is infinite.
///
/// # Examples
/// ```
/// use malachite_base::iterators::prefix_to_string;
/// use malachite_base::num::random::random_primitive_ints;
/// use malachite_base::options::random::random_somes;
/// use malachite_base::random::EXAMPLE_SEED;
/// use malachite_base::strings::ToDebugString;
///
/// assert_eq!(
///     prefix_to_string(
///         random_somes(random_primitive_ints::<u8>(EXAMPLE_SEED)).map(|x| x.to_debug_string()),
///         5
///     ),
///     "[Some(113), Some(239), Some(69), Some(108), Some(228), ...]",
/// )
/// ```
pub const fn random_somes<I: Iterator>(xs: I) -> RandomSomes<I> {
    RandomSomes { xs }
}

/// Generates random [`Option`]s with values from a given random iterator.
///
/// We don't use [`WithSpecialValue`](crate::iterators::WithSpecialValue) here because that requires
/// `I::Item` to be cloneable. The "special value" in this case, `None`, can be produced on demand
/// without any cloning.
///
/// This `struct` is created by [`random_options`]; see its documentation for more.
#[derive(Clone, Debug)]
pub struct RandomOptions<I: Iterator> {
    bs: WeightedRandomBools,
    xs: I,
}

impl<I: Iterator> Iterator for RandomOptions<I> {
    type Item = Option<I::Item>;

    #[inline]
    fn next(&mut self) -> Option<Option<I::Item>> {
        Some(if self.bs.next().unwrap() {
            self.xs.next()
        } else {
            None
        })
    }
}

/// Generates random [`Option`]s with values from a given random iterator.
///
/// The probability of generating `None` is specified by $p$ = `none_p_numerator /
/// none_p_denominator`. If a `Some` is generated, its values have the same distribution as the
/// values generated by the given iterator.
///
/// If $Q(x)$ is the probability of $x$ being generated by `xs`, then
///
/// $P(\text{None}) = p$
///
/// $P(\operatorname{Some}(x)) = (1-p)Q(x)$
///
/// `xs` must be infinite.
///
/// The output length is infinite.
///
/// # Panics
/// Panics if `none_p_denominator` is 0 or `none_p_numerator > none_p_denominator`.
///
/// # Examples
/// ```
/// use malachite_base::iterators::prefix_to_string;
/// use malachite_base::num::random::random_primitive_ints;
/// use malachite_base::options::random::random_options;
/// use malachite_base::random::EXAMPLE_SEED;
/// use malachite_base::strings::ToDebugString;
///
/// assert_eq!(
///     prefix_to_string(
///         random_options(EXAMPLE_SEED, 1, 2, &random_primitive_ints::<u8>)
///             .map(|x| x.to_debug_string()),
///         10
///     ),
///     "[Some(85), Some(11), Some(136), None, Some(200), None, Some(235), Some(134), Some(203), \
///     None, ...]"
/// )
/// ```
pub fn random_options<I: Iterator>(
    seed: Seed,
    none_p_numerator: u64,
    none_p_denominator: u64,
    xs_gen: &dyn Fn(Seed) -> I,
) -> RandomOptions<I> {
    RandomOptions {
        bs: weighted_random_bools(
            seed.fork("bs"),
            none_p_denominator - none_p_numerator,
            none_p_denominator,
        ),
        xs: xs_gen(seed.fork("xs")),
    }
}
