use crate::{
    core::{
        bounds::FieldBounds,
        circuits::{
            arithmetic::ProjectiveEdwardsPointAdditionCircuit,
            boolean::{boolean_value::Boolean, utils::CircuitType},
            traits::general_circuit::GeneralCircuit,
        },
        expressions::expr::EvalFailure,
        global_value::value::FieldValue,
    },
    traits::{FromLeBytes, GetBit, Select},
    utils::{
        elliptic_curve::{
            EDWARDS25519_D,
            EDWARDS25519_D_TWIST,
            EDWARDS25519_G_X,
            EDWARDS25519_G_Y,
            F25519,
        },
        field::{BaseField, ScalarField},
    },
};
use ff::{Field, PrimeField};

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub struct ProjectiveEdwards25519MultiplicationCircuit {
    d: BaseField,
    d_twist: BaseField,
}

impl Default for ProjectiveEdwards25519MultiplicationCircuit {
    fn default() -> Self {
        Self::new()
    }
}

impl ProjectiveEdwards25519MultiplicationCircuit {
    #[allow(dead_code)]
    pub fn new() -> Self {
        Self {
            d: BaseField::from_le_bytes(EDWARDS25519_D),
            d_twist: BaseField::from_le_bytes(EDWARDS25519_D_TWIST),
        }
    }

    #[allow(clippy::type_complexity)]
    fn mul_rec<B: Boolean + Select<T, T, T>, T: F25519>(
        &self,
        inp: Vec<((T, T, T), (T, T, T), B)>,
    ) -> Vec<(T, T, T)> {
        let addition_circuit = ProjectiveEdwardsPointAdditionCircuit::<BaseField>::new(self.d);
        if inp.len() == 1 {
            let (acc, mul, b) = inp[0];
            let sum = addition_circuit.add(acc, mul);
            vec![(
                b.select(sum.0, acc.0),
                b.select(sum.1, acc.1),
                b.select(sum.2, acc.2),
            )]
        } else {
            let half = inp.len() / 2;
            self.mul_rec(inp[..half].to_vec())
                .into_iter()
                .zip(self.mul_rec(inp[half..].to_vec()))
                .map(|(lhs, rhs)| addition_circuit.add(lhs, rhs))
                .collect()
        }
    }

    #[allow(non_snake_case)]
    /// Performs the multiplication k * P, where k_bits is an array of secret-shared bits and P is a
    /// projective point on Edwards25519.
    pub fn mul_bits<B: Boolean + Select<T, T, T>, T: F25519>(
        &self,
        k_bits: Vec<B>,
        P: (T, T, T),
        circuit_type: CircuitType,
    ) -> (T, T, T) {
        let addition_circuit = ProjectiveEdwardsPointAdditionCircuit::<BaseField>::new(self.d);
        match circuit_type {
            CircuitType::DepthOpt => {
                let inps = (0..k_bits.len())
                    .scan(P, |tmp, _| {
                        let res = *tmp;
                        *tmp = addition_circuit.add(res, res);
                        Some(res)
                    })
                    .zip(k_bits.clone());
                self.mul_rec(
                    inps.map(|(mul, b)| {
                        (
                            (
                                T::from(BaseField::ZERO),
                                T::from(BaseField::ONE),
                                T::from(BaseField::ONE),
                            ),
                            mul,
                            b,
                        )
                    })
                    .collect::<Vec<((T, T, T), (T, T, T), B)>>(),
                )[0]
            }
            CircuitType::SizeOpt => k_bits.into_iter().rev().fold(
                (
                    T::from(BaseField::ZERO),
                    T::from(BaseField::ONE),
                    T::from(BaseField::ONE),
                ),
                |acc, b| {
                    let double = addition_circuit.add(acc, acc);
                    let sum = addition_circuit.add(P, double);
                    (
                        b.select(sum.0, double.0),
                        b.select(sum.1, double.1),
                        b.select(sum.2, double.2),
                    )
                },
            ),
        }
    }

    #[allow(non_snake_case)]
    /// Performs the multiplication k * G, where k_bits is an array of secret-shared bits and G is
    /// the fixed generator of the cyclic subgroup of Edwards25519 of prime order ell.
    pub fn mul_bits_generator<B: Boolean + Select<T, T, T>, T: F25519>(
        &self,
        k_bits: Vec<B>,
        circuit_type: CircuitType,
    ) -> (T, T, T) {
        let G = (
            BaseField::from_le_bytes(EDWARDS25519_G_X),
            BaseField::from_le_bytes(EDWARDS25519_G_Y),
            BaseField::from(1),
        );
        let addition_circuit = ProjectiveEdwardsPointAdditionCircuit::<BaseField>::new(self.d);
        match circuit_type {
            CircuitType::DepthOpt => {
                let inps = (0..k_bits.len())
                    .scan(G, |tmp, _| {
                        let res = *tmp;
                        *tmp = addition_circuit.add(res, res);
                        Some(res)
                    })
                    .zip(k_bits.clone());
                self.mul_rec(
                    inps.map(|(mul, b)| {
                        (
                            (
                                T::from(BaseField::ZERO),
                                T::from(BaseField::ONE),
                                T::from(BaseField::ONE),
                            ),
                            (T::from(mul.0), T::from(mul.1), T::from(mul.2)),
                            b,
                        )
                    })
                    .collect::<Vec<((T, T, T), (T, T, T), B)>>(),
                )[0]
            }
            CircuitType::SizeOpt => k_bits.into_iter().rev().fold(
                (
                    T::from(BaseField::ZERO),
                    T::from(BaseField::ONE),
                    T::from(BaseField::ONE),
                ),
                |acc, b| {
                    let double = addition_circuit.add(acc, acc);
                    let sum =
                        addition_circuit.add((T::from(G.0), T::from(G.1), T::from(G.2)), double);
                    (
                        b.select(sum.0, double.0),
                        b.select(sum.1, double.1),
                        b.select(sum.2, double.2),
                    )
                },
            ),
        }
    }

    #[allow(non_snake_case)]
    /// Performs the multiplication k * P, where k is an element of the scalar field of Edwards25519
    /// and P is a projective point on Edwards25519.
    pub fn mul<B: Boolean + Select<T, T, T>, S: GetBit<Output = B>, T: F25519>(
        &self,
        k: S,
        P: (T, T, T),
        circuit_type: CircuitType,
    ) -> (T, T, T) {
        self.mul_bits(
            (0..ScalarField::NUM_BITS as usize)
                .map(|i| k.get_bit(i, false))
                .collect(),
            P,
            circuit_type,
        )
    }

    #[allow(non_snake_case)]
    /// Performs the multiplication k * P, where k is the lsb-to-msb binary expansion of a number
    /// and P is a projective point on Edwards25519.
    pub fn mul_str<T: F25519>(
        &self,
        k: &str,
        P: (T, T, T),
        circuit_type: CircuitType,
    ) -> (T, T, T) {
        self.mul_bits(
            k.chars().map(|c| c == '1').collect::<Vec<bool>>(),
            P,
            circuit_type,
        )
    }

    #[allow(non_snake_case)]
    /// Performs the multiplication k * P_twist, where k is the lsb-to-msb binary expansion of a
    /// number and P_twist is a projective point on Edwards25519_twist.
    pub fn mul_twist_str<T: F25519>(&self, k: &str, P_twist: (T, T, T)) -> (T, T, T) {
        let addition_circuit =
            ProjectiveEdwardsPointAdditionCircuit::<BaseField>::new(self.d_twist);
        k.chars().rev().fold(
            (
                T::from(BaseField::ZERO),
                T::from(BaseField::ONE),
                T::from(BaseField::ONE),
            ),
            |acc, b| {
                let double = addition_circuit.add(acc, acc);
                let sum = addition_circuit.add(P_twist, double);
                if b == '1' {
                    sum
                } else {
                    double
                }
            },
        )
    }
}

impl GeneralCircuit for ProjectiveEdwards25519MultiplicationCircuit {
    fn eval(
        &self,
        scalars: Vec<ScalarField>,
        bases: Vec<BaseField>,
    ) -> Result<(Vec<ScalarField>, Vec<BaseField>), EvalFailure> {
        if scalars.len() != 1 || bases.len() != 3 {
            panic!(
                "Projective Edwards Point Multiplication requires exactly one scalar and one projective group point"
            );
        }
        let res: (BaseField, BaseField, BaseField) = Self::mul(
            self,
            scalars[0],
            (bases[0], bases[1], bases[2]),
            CircuitType::default(),
        );
        Ok((vec![], vec![res.0, res.1, res.2]))
    }

    fn bounds(
        &self,
        _scalar_bounds: Vec<FieldBounds<ScalarField>>,
        _base_bounds: Vec<FieldBounds<BaseField>>,
    ) -> (Vec<FieldBounds<ScalarField>>, Vec<FieldBounds<BaseField>>) {
        (
            vec![],
            vec![FieldBounds::All, FieldBounds::All, FieldBounds::All],
        )
    }

    fn run(
        &self,
        scalar_vals: Vec<FieldValue<ScalarField>>,
        base_vals: Vec<FieldValue<BaseField>>,
    ) -> (Vec<FieldValue<ScalarField>>, Vec<FieldValue<BaseField>>) {
        if scalar_vals.len() != 1 || base_vals.len() != 3 {
            panic!(
                "Projective Edwards Point Multiplication requires exactly one scalar and one projective group point"
            );
        }
        let res = Self::mul(
            self,
            scalar_vals[0],
            (base_vals[0], base_vals[1], base_vals[2]),
            CircuitType::default(),
        );
        (vec![], vec![res.0, res.1, res.2])
    }
}

#[allow(non_snake_case)]
#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        core::circuits::traits::general_circuit::tests::TestedGeneralCircuit,
        utils::elliptic_curve::AffineEdwardsPoint,
    };

    #[test]
    fn test_mul() {
        // each 5-tuple consists of the little-endian byte representation of
        // - a scalar k
        // - the x- and y-coordinates of a point P on Edwards25519
        // - the x- and y-coordinates of k * P
        // the data is generated with a small sagemath script
        #[allow(clippy::type_complexity)]
        let test_data: [([u8; 32], ([u8; 32], [u8; 32]), ([u8; 32], [u8; 32])); 5] = [
            (
                [
                    40, 208, 1, 207, 232, 255, 177, 192, 126, 174, 86, 54, 211, 109, 120, 129, 128,
                    109, 43, 115, 233, 75, 107, 245, 77, 232, 129, 194, 255, 192, 243, 9,
                ],
                (
                    [
                        5, 61, 57, 3, 5, 250, 36, 199, 208, 233, 160, 253, 219, 181, 92, 122, 178,
                        219, 22, 80, 36, 98, 68, 190, 113, 159, 26, 86, 3, 126, 194, 46,
                    ],
                    [
                        199, 83, 132, 17, 174, 78, 127, 231, 249, 58, 70, 26, 170, 115, 62, 98,
                        159, 140, 142, 135, 239, 51, 116, 32, 92, 223, 57, 41, 139, 26, 67, 4,
                    ],
                ),
                (
                    [
                        42, 139, 99, 164, 98, 81, 96, 226, 215, 191, 19, 168, 137, 188, 1, 160, 59,
                        102, 225, 172, 100, 241, 184, 202, 138, 26, 141, 116, 215, 60, 223, 98,
                    ],
                    [
                        29, 198, 62, 245, 66, 158, 208, 80, 139, 13, 242, 188, 76, 107, 116, 14,
                        154, 240, 184, 68, 152, 108, 219, 42, 58, 193, 26, 183, 235, 4, 137, 85,
                    ],
                ),
            ),
            (
                [
                    31, 54, 214, 162, 149, 215, 139, 164, 36, 248, 71, 93, 98, 139, 237, 145, 243,
                    97, 245, 197, 103, 176, 63, 176, 119, 61, 235, 202, 142, 219, 184, 7,
                ],
                (
                    [
                        95, 2, 244, 93, 36, 110, 67, 56, 237, 99, 49, 93, 216, 135, 195, 216, 0,
                        48, 48, 16, 223, 123, 163, 134, 146, 44, 147, 6, 131, 16, 202, 122,
                    ],
                    [
                        208, 62, 3, 17, 225, 173, 18, 43, 59, 40, 138, 249, 19, 124, 113, 172, 70,
                        53, 249, 253, 214, 141, 58, 76, 137, 182, 108, 179, 103, 230, 41, 29,
                    ],
                ),
                (
                    [
                        247, 239, 241, 23, 71, 28, 139, 135, 134, 80, 131, 174, 253, 94, 16, 180,
                        224, 190, 31, 52, 126, 244, 57, 40, 165, 239, 153, 183, 48, 103, 0, 34,
                    ],
                    [
                        21, 17, 23, 121, 52, 175, 114, 210, 231, 254, 110, 68, 134, 21, 38, 89,
                        190, 157, 234, 84, 109, 48, 122, 60, 37, 46, 79, 84, 82, 43, 120, 14,
                    ],
                ),
            ),
            (
                [
                    115, 229, 223, 224, 120, 85, 199, 10, 41, 49, 226, 171, 198, 154, 81, 202, 119,
                    71, 85, 248, 16, 174, 151, 41, 154, 40, 56, 73, 33, 209, 120, 14,
                ],
                (
                    [
                        182, 215, 219, 11, 162, 103, 67, 0, 192, 38, 188, 173, 160, 245, 249, 73,
                        105, 170, 13, 161, 225, 213, 157, 149, 215, 133, 40, 20, 70, 50, 234, 29,
                    ],
                    [
                        216, 186, 71, 193, 202, 56, 203, 19, 131, 38, 63, 194, 148, 55, 48, 246,
                        120, 245, 123, 92, 138, 212, 64, 250, 108, 4, 238, 125, 112, 184, 133, 60,
                    ],
                ),
                (
                    [
                        174, 243, 243, 159, 182, 175, 121, 62, 250, 93, 227, 91, 73, 211, 221, 81,
                        39, 83, 208, 207, 162, 47, 56, 15, 241, 245, 255, 122, 161, 183, 51, 31,
                    ],
                    [
                        19, 178, 138, 161, 241, 165, 131, 185, 20, 25, 75, 38, 200, 155, 85, 216,
                        42, 30, 32, 100, 208, 102, 110, 230, 193, 58, 51, 203, 98, 234, 246, 1,
                    ],
                ),
            ),
            (
                [
                    198, 221, 49, 95, 91, 219, 237, 177, 176, 228, 44, 163, 57, 150, 117, 67, 165,
                    233, 155, 56, 73, 237, 155, 140, 46, 147, 23, 22, 63, 186, 225, 5,
                ],
                (
                    [
                        244, 169, 143, 0, 29, 211, 88, 88, 248, 51, 174, 75, 161, 170, 55, 157,
                        229, 105, 84, 240, 67, 32, 101, 67, 201, 121, 29, 208, 43, 147, 183, 95,
                    ],
                    [
                        14, 49, 153, 175, 11, 244, 33, 10, 29, 169, 89, 134, 159, 75, 232, 229,
                        203, 65, 250, 238, 168, 252, 180, 41, 98, 51, 157, 72, 123, 125, 255, 59,
                    ],
                ),
                (
                    [
                        233, 31, 38, 201, 181, 22, 163, 112, 206, 17, 138, 158, 131, 41, 46, 254,
                        72, 223, 175, 28, 137, 12, 72, 213, 47, 206, 137, 176, 48, 117, 47, 49,
                    ],
                    [
                        185, 129, 138, 79, 246, 91, 69, 221, 216, 30, 119, 41, 167, 22, 3, 2, 24,
                        43, 8, 179, 251, 16, 245, 187, 235, 42, 30, 158, 59, 105, 20, 76,
                    ],
                ),
            ),
            (
                [
                    35, 234, 47, 201, 207, 163, 212, 38, 88, 130, 12, 80, 181, 242, 136, 227, 125,
                    242, 34, 165, 102, 62, 90, 50, 20, 116, 56, 231, 203, 25, 100, 14,
                ],
                (
                    [
                        3, 102, 189, 232, 21, 245, 10, 8, 64, 129, 244, 243, 211, 25, 180, 34, 56,
                        246, 201, 242, 53, 227, 99, 25, 82, 4, 69, 32, 221, 52, 209, 54,
                    ],
                    [
                        219, 84, 113, 246, 153, 202, 200, 173, 88, 62, 109, 23, 179, 118, 138, 116,
                        95, 240, 198, 83, 189, 139, 159, 236, 93, 190, 1, 57, 27, 180, 130, 59,
                    ],
                ),
                (
                    [
                        65, 203, 139, 53, 99, 121, 109, 101, 247, 46, 41, 112, 64, 217, 141, 45,
                        90, 37, 30, 19, 157, 17, 17, 81, 98, 84, 145, 204, 211, 185, 110, 1,
                    ],
                    [
                        105, 226, 81, 33, 35, 40, 235, 102, 249, 26, 77, 166, 68, 115, 251, 59, 58,
                        226, 158, 110, 227, 56, 230, 63, 17, 98, 167, 185, 12, 168, 111, 107,
                    ],
                ),
            ),
        ];

        test_data
            .into_iter()
            .for_each(|(k, (x, y), (expected_x, expected_y))| {
                let k = ScalarField::from_le_bytes(k);
                let P = AffineEdwardsPoint::new(
                    (BaseField::from_le_bytes(x), BaseField::from_le_bytes(y)),
                    true,
                    false,
                );
                let res = k * P;
                let expected = AffineEdwardsPoint::new(
                    (
                        BaseField::from_le_bytes(expected_x),
                        BaseField::from_le_bytes(expected_y),
                    ),
                    true,
                    false,
                );
                assert_eq!(res, expected)
            });
    }

    impl TestedGeneralCircuit for ProjectiveEdwards25519MultiplicationCircuit {
        fn gen_desc<R: rand::Rng + ?Sized>(_rng: &mut R) -> Self {
            ProjectiveEdwards25519MultiplicationCircuit::new()
        }

        fn gen_n_scalars<R: rand::Rng + ?Sized>(&self, _rng: &mut R) -> usize {
            1
        }

        fn gen_n_bases<R: rand::Rng + ?Sized>(&self, _rng: &mut R, _n_scalars: usize) -> usize {
            3
        }
    }

    #[test]
    fn tested_edwards_mul() {
        ProjectiveEdwards25519MultiplicationCircuit::test(1, 4)
    }
}
