use rand::{RngExt, SeedableRng, distr::Uniform, rngs::SmallRng, seq::IndexedRandom};
use rand_distr::{Bernoulli, Normal, Pareto};

use crate::Jiffies;

/// Represents seed type for deterministic random number generators.
pub type Seed = u64;

/// Probability Distributions used to sample network latencies.
#[derive(Copy, Clone, Debug)]
pub enum Distr {
    /// Uniform distribution over `[low, high]`.
    Uniform { low: Jiffies, high: Jiffies },
    /// Bernoulli trial: with probability `p` the latency is the given value, otherwise 0.
    Bernoulli { p: f64, value: Jiffies },
    /// Truncated normal distribution clamped to `[low, high]`.
    /// See <https://en.wikipedia.org/wiki/Truncated_normal_distribution>.
    Normal {
        mean: Jiffies,
        std_dev: Jiffies,
        low: Jiffies,
        high: Jiffies,
    },
    /// Pareto distribution with minimal value - `scale`.
    Pareto { scale: Jiffies, shape: f64 },
}

impl Distr {
    pub(super) fn safe_window(&self) -> Jiffies {
        match self.clone() {
            Self::Uniform { low, .. } => low,
            Self::Bernoulli { .. } => Jiffies(1),
            Self::Normal { low, .. } => low,
            Self::Pareto { scale, .. } => scale,
        }
    }
}

#[derive(Debug)]
pub(crate) struct Randomizer {
    rnd: SmallRng,
}

impl Default for Randomizer {
    fn default() -> Self {
        Self {
            rnd: SmallRng::seed_from_u64(0),
        }
    }
}

impl Randomizer {
    pub(crate) fn new(seed: Seed) -> Self {
        Self {
            rnd: SmallRng::seed_from_u64(seed),
        }
    }

    pub(crate) fn random_usize(&mut self, d: Distr) -> usize {
        match d {
            Distr::Uniform { low, high } => {
                let distr = Uniform::new_inclusive(low.0, high.0)
                    .expect("invalid bounds for uniform distribution");
                self.rnd.sample(distr)
            }
            Distr::Bernoulli { p, value } => {
                let distr =
                    Bernoulli::new(p).expect("invalid probability for bernoulli distribution");
                if self.rnd.sample(distr) { value.0 } else { 0 }
            }
            Distr::Normal {
                mean: Jiffies(mean),
                std_dev: Jiffies(std_dev),
                low: Jiffies(low),
                high: Jiffies(high),
            } => {
                let distr = Normal::new(mean as f64, std_dev as f64)
                    .expect("invalid parameters for normal distribution");
                loop {
                    let sample: f64 = self.rnd.sample(distr);
                    let rounded = sample.round() as isize;
                    if rounded >= low as isize && rounded <= high as isize {
                        return rounded as usize;
                    }
                }
            }
            Distr::Pareto { scale, shape } => {
                let distr = Pareto::new(scale.0 as f64, shape)
                    .expect("invalid parameters for pareto distribution");
                self.rnd.sample(distr) as usize
            }
        }
    }

    pub(crate) fn choose_from_slice<'a, T: Copy>(&mut self, from: &[T]) -> T {
        from.choose(&mut self.rnd)
            .copied()
            .expect("choose from empty slice")
    }
}
