//! Random floating point number generation with the `rand` crate.
//!
//! There are two new distributions for generating random floats. The first one is [Uniform01],
//! which supports generating floats between 0 and 1. This is also the underlying implementation
//! of the builtin `rand` distributions [Standard], [Open01], [OpenClosed01]. The other one is
//! [UniformFBig], which supports generating floats in a certain range. This is also the
//! backend for the [SampleUniform] trait.
//!
//! # Examples
//!
//! ```
//! use dashu_float::rand::Uniform01;
//! # use rand_v08::{distributions::uniform::Uniform, thread_rng, Rng};
//!
//! type FBig = dashu_float::FBig;
//!
//! // generate FBigs in a [0, 1) or [0, 1] with a given precision
//! let a: FBig = thread_rng().sample(Uniform01::new(10));
//! let b: FBig = thread_rng().sample(Uniform01::new_closed(10));
//! let c: FBig = thread_rng().gen(); // the default distribution generates in [0, 1)
//! assert!(a >= FBig::ZERO && a < FBig::ONE);
//! assert!(b >= FBig::ZERO && b <= FBig::ONE);
//! assert!(c >= FBig::ZERO && c < FBig::ONE);
//!
//! // generate FBigs in a range
//! let a = thread_rng().gen_range(FBig::from(3)..FBig::from(10));
//! let b = thread_rng().sample(Uniform::new(FBig::from(-5), &a));
//! assert!(a >= FBig::from(3) && a < FBig::from(10));
//! assert!(b >= FBig::from(-5) && b < a);
//! ```
//!
//! # Precision and Rounding
//!
//! The precision of a float generated by different distributions is explained below:
//! * [Uniform01] will generate floats with the precision decided by the constructor.
//! * [Standard], [Open01], [OpenClosed01] will generate floats with the max precision
//!   such that the significand fits in a [DoubleWord].
//! * [UniformFBig] (and therefore [Uniform][rand_v08::distributions::Uniform]) will generate floats
//!   the precision being the maximum between the interval boundaries.
//!
//! The rounding of the [FBig] type doesn't affect the number generation process.

use core::marker::PhantomData;

use crate::{
    fbig::FBig,
    repr::{Context, Repr, Word},
    round::{mode, Round},
};

use dashu_base::EstimatedLog2;
use dashu_int::{
    rand::{UniformBelow, UniformBits},
    DoubleWord, UBig,
};
use rand_v08::{
    distributions::{
        uniform::{SampleBorrow, SampleUniform, UniformSampler},
        Open01, OpenClosed01, Standard,
    },
    prelude::Distribution,
    Rng,
};

/// The back-end implementing [UniformSampler] for [FBig] (and [DBig][crate::DBig]).
///
/// See the module ([rand][crate::rand]) level documentation for examples.
pub struct UniformFBig<R: Round, const B: Word> {
    sampler: Uniform01<B>,

    // XXX: there should be a way to prevent storing a context in this struct
    //      (just like a complex number). It might require binop between FBig and Repr.
    scale: FBig<mode::Down, B>,
    offset: FBig<mode::Down, B>,

    /// This field is used to distinguish between uniform distributions
    /// with different rounding modes, no actual effect on the sampling.
    _marker: PhantomData<R>,
}

impl<R: Round, const B: Word> UniformFBig<R, B> {
    /// Same as [UniformSampler::new] but with an additional argument to specify
    /// the precision of the sampled [FBig].
    #[inline]
    pub fn new(low: &FBig<R, B>, high: &FBig<R, B>, precision: usize) -> Self {
        assert!(low <= high);

        Self {
            sampler: Uniform01::new(precision),
            scale: (high - low).with_rounding(),
            offset: low.clone().with_rounding(),
            _marker: PhantomData,
        }
    }

    /// Same as [UniformSampler::new_inclusive] but with an additional argument to
    /// specify the precision of the sampled [FBig].
    #[inline]
    pub fn new_inclusive(low: &FBig<R, B>, high: &FBig<R, B>, precision: usize) -> Self {
        assert!(low <= high);

        Self {
            sampler: Uniform01::new_closed(precision),
            scale: (high - low).with_rounding(),
            offset: low.clone().with_rounding(),
            _marker: PhantomData,
        }
    }
}

impl<R: Round, const B: Word> UniformSampler for UniformFBig<R, B> {
    type X = FBig<R, B>;

    #[inline]
    fn new<B1, B2>(low: B1, high: B2) -> Self
    where
        B1: SampleBorrow<Self::X> + Sized,
        B2: SampleBorrow<Self::X> + Sized,
    {
        let precision = low.borrow().precision().min(high.borrow().precision());
        UniformFBig::new(low.borrow(), high.borrow(), precision)
    }

    #[inline]
    fn new_inclusive<B1, B2>(low: B1, high: B2) -> Self
    where
        B1: SampleBorrow<Self::X> + Sized,
        B2: SampleBorrow<Self::X> + Sized,
    {
        let precision = low.borrow().precision().min(high.borrow().precision());
        UniformFBig::new_inclusive(low.borrow(), high.borrow(), precision)
    }

    #[inline]
    fn sample<RNG: Rng + ?Sized>(&self, rng: &mut RNG) -> Self::X {
        <Self as Distribution<FBig<R, B>>>::sample(self, rng)
    }
}

impl<R: Round, const B: Word> Distribution<FBig<R, B>> for UniformFBig<R, B> {
    fn sample<RNG: Rng + ?Sized>(&self, rng: &mut RNG) -> FBig<R, B> {
        // After we have a sample in [0, 1), all the following operations are rounded down
        // so that we can ensure we don't reach the right bound.
        let unit: FBig<mode::Down, B> = self.sampler.sample(rng);
        (unit * &self.scale + &self.offset).with_rounding()
    }
}

impl<R: Round, const B: Word> SampleUniform for FBig<R, B> {
    type Sampler = UniformFBig<R, B>;
}

// when sampling with the builtin distributions, the precision is choosen
// such that the significand fits in a double word and no allocation is required.
#[inline]
fn get_inline_precision<const B: Word>() -> usize {
    (DoubleWord::BITS as f32 / B.log2_bounds().1) as _
}

impl<R: Round, const B: Word> Distribution<FBig<R, B>> for Standard {
    fn sample<RNG: Rng + ?Sized>(&self, rng: &mut RNG) -> FBig<R, B> {
        Uniform01::<B>::new(get_inline_precision::<B>()).sample(rng)
    }
}

impl<R: Round, const B: Word> Distribution<FBig<R, B>> for Open01 {
    #[inline]
    fn sample<RNG: Rng + ?Sized>(&self, rng: &mut RNG) -> FBig<R, B> {
        Uniform01::<B>::new_open(get_inline_precision::<B>()).sample(rng)
    }
}

impl<R: Round, const B: Word> Distribution<FBig<R, B>> for OpenClosed01 {
    #[inline]
    fn sample<RNG: Rng + ?Sized>(&self, rng: &mut RNG) -> FBig<R, B> {
        Uniform01::<B>::new_open_closed(get_inline_precision::<B>()).sample(rng)
    }
}

/// A uniform distribution between 0 and 1. It can be used to replace the [Standard], [Open01],
/// [OpenClosed01] distributions from the `rand` crate when you want to customize the precision
/// of the generated float number.
pub struct Uniform01<const BASE: Word> {
    precision: usize,
    range: Option<UBig>, // BASE ^ precision (±1 if necessary)
    include_zero: bool,  // whether include the zero
    include_one: bool,   // whether include the one
}

impl<const B: Word> Uniform01<B> {
    /// Create a uniform distribution in `[0, 1)` with a given precision.
    #[inline]
    pub fn new(precision: usize) -> Self {
        let range = match B {
            2 => None,
            _ => Some(UBig::from_word(B).pow(precision)),
        };
        Self {
            precision,
            range,
            include_zero: true,
            include_one: false,
        }
    }

    /// Create a uniform distribution in `[0, 1]` with a given precision.
    #[inline]
    pub fn new_closed(precision: usize) -> Self {
        let range = Some(UBig::from_word(B).pow(precision) + UBig::ONE);
        Self {
            precision,
            range,
            include_zero: true,
            include_one: true,
        }
    }

    /// Create a uniform distribution in `(0, 1)` with a given precision.
    #[inline]
    pub fn new_open(precision: usize) -> Self {
        let range = match B {
            2 => None,
            _ => Some(UBig::from_word(B).pow(precision) - UBig::ONE),
        };
        Self {
            precision,
            range,
            include_zero: false,
            include_one: false,
        }
    }

    /// Create a uniform distribution in `(0, 1]` with a given precision.
    #[inline]
    pub fn new_open_closed(precision: usize) -> Self {
        let range = match B {
            2 => None,
            _ => Some(UBig::from_word(B).pow(precision)),
        };
        Self {
            precision,
            range,
            include_zero: false,
            include_one: true,
        }
    }
}

impl<R: Round, const B: Word> Distribution<FBig<R, B>> for Uniform01<B> {
    fn sample<RNG: Rng + ?Sized>(&self, rng: &mut RNG) -> FBig<R, B> {
        let repr = match (self.include_zero, self.include_one) {
            (true, false) => {
                // sample in [0, 1)
                let signif: UBig = if B == 2 {
                    UniformBits::new(self.precision).sample(rng)
                } else {
                    UniformBelow::new(self.range.as_ref().unwrap()).sample(rng)
                };
                Repr::<B>::new(signif.into(), -(self.precision as isize))
            }
            (true, true) => {
                // sample in [0, 1]
                let signif: UBig = UniformBelow::new(self.range.as_ref().unwrap()).sample(rng);
                Repr::new(signif.into(), -(self.precision as isize))
            }
            (false, false) => {
                // sample in (0, 1)
                let signif = if B == 2 {
                    loop {
                        // simply reject zero
                        let n: UBig = UniformBits::new(self.precision).sample(rng);
                        if !n.is_zero() {
                            break n;
                        }
                    }
                } else {
                    let n: UBig = UniformBelow::new(self.range.as_ref().unwrap()).sample(rng);
                    n + UBig::ONE
                };
                Repr::<B>::new(signif.into(), -(self.precision as isize))
            }
            (false, true) => {
                // sample in (0, 1]
                let signif: UBig = if B == 2 {
                    UniformBits::new(self.precision).sample(rng)
                } else {
                    UniformBelow::new(self.range.as_ref().unwrap()).sample(rng)
                };
                Repr::<B>::new((signif + UBig::ONE).into(), -(self.precision as isize))
            }
        };

        let context = Context::<mode::Down>::new(self.precision);
        FBig::new(repr, context).with_rounding()
    }
}