use backend::{AddNormal, FillUniform, VecZnx, ZnxView, ZnxViewMut, alloc_aligned};
use sampling::source::Source;

use crate::{Infos, LWECiphertext, LWESecret, SIX_SIGMA, lwe::LWEPlaintext};

impl<DataSelf> LWECiphertext<DataSelf>
where
    DataSelf: AsMut<[u8]> + AsRef<[u8]>,
{
    pub fn encrypt_sk<DataPt, DataSk>(
        &mut self,
        pt: &LWEPlaintext<DataPt>,
        sk: &LWESecret<DataSk>,
        source_xa: &mut Source,
        source_xe: &mut Source,
        sigma: f64,
    ) where
        DataPt: AsRef<[u8]>,
        DataSk: AsRef<[u8]>,
    {
        #[cfg(debug_assertions)]
        {
            assert_eq!(self.n(), sk.n())
        }

        let basek: usize = self.basek();

        self.data.fill_uniform(basek, 0, self.size(), source_xa);
        let mut tmp_znx: VecZnx<Vec<u8>> = VecZnx::<Vec<u8>>::new::<i64>(1, 1, self.size());

        let min_size = self.size().min(pt.size());

        (0..min_size).for_each(|i| {
            tmp_znx.at_mut(0, i)[0] = pt.data.at(0, i)[0]
                - self.data.at(0, i)[1..]
                    .iter()
                    .zip(sk.data.at(0, 0))
                    .map(|(x, y)| x * y)
                    .sum::<i64>();
        });

        (min_size..self.size()).for_each(|i| {
            tmp_znx.at_mut(0, i)[0] -= self.data.at(0, i)[1..]
                .iter()
                .zip(sk.data.at(0, 0))
                .map(|(x, y)| x * y)
                .sum::<i64>();
        });

        tmp_znx.add_normal(basek, 0, self.k(), source_xe, sigma, sigma * SIX_SIGMA);

        let mut tmp_bytes: Vec<u8> = alloc_aligned(size_of::<i64>());

        tmp_znx.normalize(basek, 0, &mut tmp_bytes);

        (0..self.size()).for_each(|i| {
            self.data.at_mut(0, i)[0] = tmp_znx.at(0, i)[0];
        });
    }
}