use crate::{ArcisRNG, ArcisType, Reveal};
use rand::{thread_rng, Rng};
use std::cell::Cell;

thread_local! {
    static GENERATED_BOOLS: Cell<Vec<bool>> = const {Cell::new(Vec::new()) };
}

pub fn generated_bools() -> Vec<bool> {
    GENERATED_BOOLS.take()
}

impl ArcisRNG {
    /// Generates a boolean, with 1/2 probability to be `true`, 1/2 probability to be `false`.
    pub fn bool() -> bool {
        let mut rng = thread_rng();
        let b = rng.gen_bool(0.5);
        let mut v = GENERATED_BOOLS.take();
        v.push(b);
        GENERATED_BOOLS.replace(v);
        b
    }
    /// Generates a secret integer between 0 and 2^width - 1 included.
    /// Note that `width` must be compile-time known.
    pub fn gen_integer_from_width(width: usize) -> u128 {
        assert!(width <= u128::BITS as usize);
        let mut integer: u128 = 0;
        for i in 0..width {
            integer += (Self::bool() as u128) << i;
        }
        integer
    }
    /// Generates a public integer between 0 and 2^width - 1 included.
    /// Note that `width` must be compile-time known.
    pub fn gen_public_integer_from_width(width: usize) -> u128 {
        Self::gen_integer_from_width(width).reveal()
    }
    /// Uniformly and secretly shuffles the array in argument.
    /// Complexity is in `O(n*log³(n) + n*log²(n)*sizeof(T))`.
    pub fn shuffle<T>(data: &mut [T]) {
        // We want to end up with the same shuffle as the MPC version
        // if we start from the same generated booleans.
        let n = data.len();
        // This needs to be the same as in arcis/src/core/circuits/sort.rs
        let randomness_width = if n < 2 {
            // No need for any randomness to shuffle 0 or 1 items.
            0
        } else {
            // We want P(two generated numbers are identical) ~<= 2^-64
            // P(two generated numbers are identical) < E(pair of identical numbers)
            // = n(n-1)/2 * 2^-randomness_width
            // With this value, we have from 2^-65 to 2^-63
            // depending on how close to a power of two n is.
            2 * (n.ilog2() as usize) + 64
        };
        let mut random_nums = Vec::with_capacity(n);
        for i in 0..n {
            random_nums.push((Self::gen_integer_from_width(randomness_width), i));
        }
        random_nums.sort();
        let final_perm = random_nums.into_iter().map(|(_, x)| x).collect::<Vec<_>>();
        apply_perm(data, &final_perm);
    }

    /// Generates uniformly a `T`.
    /// Works for every supported type except types containing public keys, floats, and slices.
    /// Our interpreter does not have type inference so you will need to explicit `T`.
    /// For instance: `ArcisRNG::gen_uniform::<[u8; 1024]>()`.
    pub fn gen_uniform<T: ArcisType>() -> T {
        let bools = (0..T::n_bools())
            .map(|_| ArcisRNG::bool())
            .collect::<Vec<_>>();
        T::from_bools(&bools)
    }
}

/// Applies perm, ie changes data
/// from data[i] = init_data[i]
/// to data[i] = init_data[perm[i]].
fn apply_perm<T>(data: &mut [T], perm: &[usize]) {
    let n = data.len();
    let mut current_perm: Vec<_> = (0..n).collect();
    let mut positions: Vec<_> = (0..n).collect();
    for i in 0..(n - 1) {
        let val_to_transfer = perm[i];
        let other_i = positions[val_to_transfer];
        if i == other_i {
            continue;
        }
        let i_val = current_perm[i];
        assert_eq!(current_perm[i], i_val);
        assert_eq!(current_perm[other_i], val_to_transfer);
        assert_eq!(positions[i_val], i);
        assert_eq!(positions[val_to_transfer], other_i);
        data.swap(i, other_i);
        current_perm.swap(i, other_i);
        positions.swap(i_val, val_to_transfer);
    }
    assert_eq!(current_perm.as_slice(), perm);
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_apply_perm() {
        let final_perm = [1, 2, 0, 5, 4, 3];
        let init_data = [10, 11, 12, 13, 14, 15];
        let mut data = init_data;
        apply_perm(&mut data, &final_perm);
        assert_eq!(data, [11, 12, 10, 15, 14, 13]);
        for i in 0..init_data.len() {
            assert_eq!(data[i], init_data[final_perm[i]]);
        }
    }
}
