use crate::{
    core::{
        bounds::FieldBounds,
        circuits::{arithmetic::SqrtCircuit, traits::arithmetic_circuit::ArithmeticCircuit},
        expressions::expr::EvalFailure,
        global_value::value::FieldValue,
    },
    traits::{Invert, Reveal},
    utils::field::BaseField,
};

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

impl ArithmeticCircuit<BaseField> for RefreshShares {
    fn eval(&self, x: Vec<BaseField>) -> Result<Vec<BaseField>, EvalFailure> {
        Ok(x)
    }

    fn bounds(&self, bounds: Vec<FieldBounds<BaseField>>) -> Vec<FieldBounds<BaseField>> {
        bounds
    }

    fn run(&self, vals: Vec<FieldValue<BaseField>>) -> Vec<FieldValue<BaseField>> {
        vals.into_iter()
            .map(|v| {
                if v.is_plaintext() {
                    v
                } else {
                    let r = FieldValue::<BaseField>::random();
                    let square_revealed = (r * r).reveal();
                    let sqrt = SqrtCircuit.run(vec![square_revealed])[0];
                    // on plaintext values we simply set is_expected_non_zero = true
                    let inv_sqrt = sqrt.invert(true);
                    let one_or_minus_one = r * inv_sqrt;
                    let one = one_or_minus_one * one_or_minus_one;
                    let zero = one - FieldValue::<BaseField>::from(1);
                    let res = v + zero;
                    res.with_bounds(v.bounds());
                    res
                }
            })
            .collect::<Vec<FieldValue<BaseField>>>()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::circuits::traits::arithmetic_circuit::tests::TestedArithmeticCircuit;
    use rand::Rng;

    impl TestedArithmeticCircuit<BaseField> for RefreshShares {
        fn gen_desc<R: Rng + ?Sized>(_rng: &mut R) -> Self {
            Self
        }

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