use crate::{
    core::{
        actually_used_field::ActuallyUsedField,
        bounds::FieldBounds,
        circuits::{
            boolean::boolean_value::BooleanValue,
            traits::arithmetic_circuit::ArithmeticCircuit,
        },
        expressions::expr::EvalFailure,
        global_value::value::FieldValue,
    },
    traits::{FromLeBits, GetBit, GreaterThan, RandomBit},
    utils::used_field::UsedField,
};

/// A classical C-style loop iterator for (i = begin; pred(i); i = iter(i)) { yield i }
fn for_loop_iter<'a>(
    begin: usize,
    iter: impl Fn(usize) -> usize + 'a,
    pred: impl Fn(usize) -> bool + 'a,
) -> impl Iterator<Item = usize> + 'a {
    std::iter::successors(if pred(begin) { Some(begin) } else { None }, move |i| {
        let new_i = iter(*i);
        if pred(new_i) {
            Some(new_i)
        } else {
            None
        }
    })
}

pub trait MPCSwappable<F: ActuallyUsedField> {
    fn conditional_swap(
        &mut self,
        condition: BooleanValue,
        first_index: usize,
        second_index: usize,
    );
    fn len(&self) -> usize;
    #[allow(dead_code)]
    fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

pub trait MPCSortable<F: ActuallyUsedField>: MPCSwappable<F> {
    fn get_val(&mut self, index: usize) -> FieldValue<F>;
}

impl<F: ActuallyUsedField> MPCSwappable<F> for [FieldValue<F>] {
    fn conditional_swap(
        &mut self,
        condition: BooleanValue,
        first_index: usize,
        second_index: usize,
    ) {
        let first = self[first_index];
        let second = self[second_index];
        let union_bounds = first.bounds().union(second.bounds());
        let signed = union_bounds.has_negatives();
        let circuit_size = union_bounds.signed_bin_size();
        let first_bits = (0..circuit_size)
            .map(|i| first.get_bit(i, signed))
            .collect::<Vec<BooleanValue>>();
        let second_bits = (0..circuit_size)
            .map(|i| second.get_bit(i, signed))
            .collect::<Vec<BooleanValue>>();
        let (first_swapped_bits, second_swapped_bits): (Vec<BooleanValue>, Vec<BooleanValue>) =
            first_bits
                .into_iter()
                .zip(second_bits)
                .map(|(bit_first, bit_second)| {
                    let tmp = condition & (bit_first ^ bit_second);
                    (bit_first ^ tmp, bit_second ^ tmp)
                })
                .unzip();

        self[first_index] = FieldValue::<F>::from_le_bits(first_swapped_bits, signed);
        self[second_index] = FieldValue::<F>::from_le_bits(second_swapped_bits, signed);
    }

    fn len(&self) -> usize {
        self.len()
    }
}
impl<F: ActuallyUsedField> MPCSortable<F> for [FieldValue<F>] {
    fn get_val(&mut self, index: usize) -> FieldValue<F> {
        self[index]
    }
}

pub fn random_number<F: ActuallyUsedField>(randomness_width: usize) -> FieldValue<F> {
    FieldValue::<F>::from_le_bits(
        (0..randomness_width)
            .map(|_| BooleanValue::random())
            .collect::<Vec<BooleanValue>>(),
        false,
    )
}

pub struct Shuffle<F: ActuallyUsedField, T: MPCSwappable<F>> {
    shuffled: T,
    random: Vec<FieldValue<F>>,
}

impl<F: ActuallyUsedField, T: MPCSwappable<F>> Shuffle<F, T> {
    #[allow(dead_code)]
    pub fn new(shuffled: T) -> Self {
        let n = shuffled.len();
        // This needs to be the same as in arcis/src/rng.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 = Vec::with_capacity(n);
        for _ in 0..n {
            random.push(random_number(randomness_width));
        }
        Self { shuffled, random }
    }

    #[allow(dead_code)]
    pub fn mpc_shuffle(&mut self) {
        batcher_odd_even_mergesort(self, false)
    }
}

impl<F: ActuallyUsedField, T: MPCSwappable<F>> MPCSwappable<F> for Shuffle<F, T> {
    fn conditional_swap(
        &mut self,
        condition: BooleanValue,
        first_index: usize,
        second_index: usize,
    ) {
        self.shuffled
            .conditional_swap(condition, first_index, second_index);
        self.random
            .conditional_swap(condition, first_index, second_index);
    }

    fn len(&self) -> usize {
        self.shuffled.len()
    }
}
impl<F: ActuallyUsedField, T: MPCSwappable<F>> MPCSortable<F> for Shuffle<F, T> {
    fn get_val(&mut self, index: usize) -> FieldValue<F> {
        self.random[index]
    }
}

fn cmp<F: ActuallyUsedField>(
    signed: bool,
    first_val: FieldValue<F>,
    second_val: FieldValue<F>,
) -> BooleanValue {
    if signed {
        first_val.signed_gt(second_val)
    } else {
        first_val.gt(second_val)
    }
}

/// Sorts in increasing order or according to perm. Fills perm in the second case.
pub fn batcher_odd_even_mergesort<F: ActuallyUsedField>(
    data: &mut (impl MPCSortable<F> + ?Sized),
    signed: bool,
) {
    let n = data.len();
    let p_iter = for_loop_iter(1, |p| p * 2, |p| p < n);
    for p in p_iter {
        let k_iter = for_loop_iter(p, |k| k / 2, |k| k >= 1);
        for k in k_iter {
            let j_iter = for_loop_iter(k % p, |j| j + 2 * k, |j| j + k < n);
            for j in j_iter {
                let i_iter = for_loop_iter(0, |i| i + 1, |i| i < k && i + j + k < n);
                for i in i_iter {
                    let first = i + j;
                    let second = i + j + k;
                    if first / (2 * p) == second / (2 * p) {
                        let cmp = cmp(signed, data.get_val(first), data.get_val(second));
                        data.conditional_swap(cmp, first, second);
                    }
                }
            }
        }
    }
}

#[derive(Debug, Clone)]
#[allow(dead_code)]
pub struct SignedSort;

impl<F: UsedField> ArithmeticCircuit<F> for SignedSort {
    fn eval(&self, x: Vec<F>) -> Result<Vec<F>, EvalFailure> {
        let mut temp: Vec<_> = x.into_iter().map(|x| x - F::TWO_INV).collect();
        temp.sort_by(|a, b| a.partial_cmp(b).unwrap());
        let res = temp.into_iter().map(|x| x + F::TWO_INV).collect();
        Ok(res)
    }

    fn bounds(&self, bounds: Vec<FieldBounds<F>>) -> Vec<FieldBounds<F>> {
        let n = bounds.len();
        let res = bounds
            .into_iter()
            .fold(FieldBounds::Empty, |b, t| b.union(t));
        vec![res; n]
    }

    fn run(&self, mut vals: Vec<FieldValue<F>>) -> Vec<FieldValue<F>>
    where
        F: ActuallyUsedField,
    {
        batcher_odd_even_mergesort(vals.as_mut_slice(), true);
        vals
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        core::circuits::traits::arithmetic_circuit::tests::TestedArithmeticCircuit,
        ArcisField,
    };
    use rand::Rng;
    impl TestedArithmeticCircuit<ArcisField> for SignedSort {
        fn gen_desc<R: Rng + ?Sized>(_rng: &mut R) -> Self {
            SignedSort
        }

        fn gen_n_inputs<R: Rng + ?Sized>(&self, rng: &mut R) -> usize {
            let mut size = 0;
            while rng.gen_bool(0.875) {
                size += 1;
            }
            size
        }
    }
    #[test]
    fn tested() {
        SignedSort::test(1, 16);
    }
}
