use crate::modulus::WordOps;
use crate::poly::{Poly, PolyRNS};
use crate::ring::{Ring, RingRNS};
use num::ToPrimitive;
use rand_distr::{Normal, Distribution};
use sampling::source::Source;

impl Ring<u64> {
    pub fn fill_uniform(&self, source: &mut Source, a: &mut Poly<u64>) {
        let max: u64 = self.modulus.q;
        let mask: u64 = max.mask();
        a.0.iter_mut()
            .for_each(|a| *a = source.next_u64n(max, mask));
    }

    pub fn fill_dist_f64<T: Distribution<f64>>(&self, source: &mut Source, dist: T, bound: f64, a: &mut Poly<u64>) {
        let max: u64 = self.modulus.q;
        a.0.iter_mut()
            .for_each(|a| {

                let mut dist_f64: f64 = dist.sample(source);
                
                while dist_f64.abs() > bound{
                    dist_f64 = dist.sample(source)
                }

                let dist_u64: u64 = (dist_f64+0.5).abs().to_u64().unwrap();
                let sign: u64 = dist_f64.to_bits()>>63;

                *a = (dist_u64 * sign) | (max-dist_u64)*(sign^1)
        });
    }
}

impl RingRNS<u64> {
    pub fn fill_uniform(&self, source: &mut Source, a: &mut PolyRNS<u64>) {
        self.0
            .iter()
            .enumerate()
            .for_each(|(i, r)| r.fill_uniform(source, a.at_mut(i)));
    }

    pub fn fill_dist_f64<T: Distribution<f64>>(&self, source: &mut Source, dist: T, bound: f64, a: &mut PolyRNS<u64>) {
        (0..a.n()).for_each(|j|{
            let mut dist_f64: f64 = dist.sample(source);
                
            while dist_f64.abs() > bound{
                dist_f64 = dist.sample(source)
            }

            let dist_u64: u64 = (dist_f64+0.5).abs().to_u64().unwrap();
            let sign: u64 = dist_f64.to_bits()>>63;

            self.0.iter().enumerate().for_each(|(i, r)|{
                a.at_mut(i).0[j] = (dist_u64 * sign) | (r.modulus.q-dist_u64)*(sign^1);
            })
        })
    }
}