use rand::Rng;

#[derive(Debug, Clone)]
pub struct WithProbability<T>
where
    T: Clone,
{
    value: T,
    probability: f64,
}

impl<T> WithProbability<T>
where
    T: Clone,
{
    pub fn new(value: T, probability: f64) -> WithProbability<T> {
        assert!(probability > 0.0 && probability <= 1.0);
        WithProbability { value, probability }
    }

    pub fn get_probability(&self) -> f64 {
        self.probability
    }

    pub fn get_value(&self) -> &T {
        &(self.value)
    }
}

#[derive(Debug, Clone)]
pub struct DiscreteRandom<T, R>
where
    T: Clone,
    R: Rng,
{
    values: Vec<WithProbability<T>>,
    rng: R,
    biased_coins: Vec<(usize, usize, f64)>,
}

impl<T, R> DiscreteRandom<T, R>
where
    T: Clone,
    R: Rng,
{
    pub fn new(values: &[WithProbability<T>], rng: R) -> DiscreteRandom<T, R> {
        assert!(
            (values
                .iter()
                .map(WithProbability::get_probability)
                .sum::<f64>()
                - 1.0)
                .abs()
                <= std::f64::EPSILON
        );
        DiscreteRandom {
            values: values.to_vec(),
            rng,
            biased_coins: Self::build_biased_coins(
                values
                    .iter()
                    .map(WithProbability::get_probability)
                    .collect::<Vec<_>>(),
            ),
        }
    }

    pub fn random(&mut self) -> &T {
        let n = self.biased_coins.len();
        let idx = self.rng.gen_range(0, n);

        self.flip_biased_coin(idx)
    }

    fn build_biased_coins(probabilities: Vec<f64>) -> Vec<(usize, usize, f64)> {
        let n = probabilities.len();
        let scaled = probabilities.into_iter().map(|p| p*n).collect::<Vec<_>>();
        let small: Vec<f64> = Vec::with_capacity(n);
        let big: Vec<f64> = Vec::with_capacity(n);

        Vec::new()
    }

    fn flip_biased_coin(&mut self, idx: usize) -> &T {
        let (index_1, index_2, probability) = self.biased_coins[idx];
        if self.rng.gen::<f64>() < probability {
            self.values[index_1].get_value()
        } else {
            self.values[index_2].get_value()
        }
    }
}

fn main() {
    let dr = DiscreteRandom::new(
        &[
            WithProbability::new(String::from("test1"), 1.0 / 3.0),
            WithProbability::new(String::from("test2"), 1.0 / 3.0),
            WithProbability::new(String::from("test3"), 1.0 / 3.0),
        ],
        rand::thread_rng(),
    );
    println!("{:?}", dr);
}