1use poulpy_hal::{
  2    api::{TakeSlice, ZnxInfos, ZnxSliceSize, ZnxView, ZnxViewMut},
  3    layouts::{Scratch, Zn, ZnToMut},
  4    oep::{
  5        TakeSliceImpl, ZnAddDistF64Impl, ZnAddNormalImpl, ZnFillDistF64Impl, ZnFillNormalImpl, ZnFillUniformImpl,
  6        ZnNormalizeInplaceImpl,
  7    },
  8    source::Source,
  9};
 10use rand_distr::Normal;
 11
 12use crate::cpu_spqlios::{FFT64, ffi::zn64};
 13
 14unsafe impl ZnNormalizeInplaceImpl<Self> for FFT64
 15where
 16    Self: TakeSliceImpl<Self>,
 17{
 18    fn zn_normalize_inplace_impl<A>(n: usize, basek: usize, a: &mut A, a_col: usize, scratch: &mut Scratch<Self>)
 19    where
 20        A: ZnToMut,
 21    {
 22        let mut a: Zn<&mut [u8]> = a.to_mut();
 23
 24        let (tmp_bytes, _) = scratch.take_slice(n * size_of::<i64>());
 25
 26        unsafe {
 27            zn64::zn64_normalize_base2k_ref(
 28                n as u64,
 29                basek as u64,
 30                a.at_mut_ptr(a_col, 0),
 31                a.size() as u64,
 32                a.sl() as u64,
 33                a.at_ptr(a_col, 0),
 34                a.size() as u64,
 35                a.sl() as u64,
 36                tmp_bytes.as_mut_ptr(),
 37            );
 38        }
 39    }
 40}
 41
 42unsafe impl ZnFillUniformImpl<Self> for FFT64 {
 43    fn zn_fill_uniform_impl<R>(n: usize, basek: usize, res: &mut R, res_col: usize, k: usize, source: &mut Source)
 44    where
 45        R: ZnToMut,
 46    {
 47        let mut a: Zn<&mut [u8]> = res.to_mut();
 48        let base2k: u64 = 1 << basek;
 49        let mask: u64 = base2k - 1;
 50        let base2k_half: i64 = (base2k >> 1) as i64;
 51        (0..k.div_ceil(basek)).for_each(|j| {
 52            a.at_mut(res_col, j)[..n]
 53                .iter_mut()
 54                .for_each(|x| *x = (source.next_u64n(base2k, mask) as i64) - base2k_half);
 55        })
 56    }
 57}
 58
 59unsafe impl ZnFillDistF64Impl<Self> for FFT64 {
 60    fn zn_fill_dist_f64_impl<R, D: rand::prelude::Distribution<f64>>(
 61        n: usize,
 62        basek: usize,
 63        res: &mut R,
 64        res_col: usize,
 65        k: usize,
 66        source: &mut Source,
 67        dist: D,
 68        bound: f64,
 69    ) where
 70        R: ZnToMut,
 71    {
 72        let mut a: Zn<&mut [u8]> = res.to_mut();
 73        assert!(
 74            (bound.log2().ceil() as i64) < 64,
 75            "invalid bound: ceil(log2(bound))={} > 63",
 76            (bound.log2().ceil() as i64)
 77        );
 78
 79        let limb: usize = k.div_ceil(basek) - 1;
 80        let basek_rem: usize = (limb + 1) * basek - k;
 81
 82        if basek_rem != 0 {
 83            a.at_mut(res_col, limb)[..n].iter_mut().for_each(|a| {
 84                let mut dist_f64: f64 = dist.sample(source);
 85                while dist_f64.abs() > bound {
 86                    dist_f64 = dist.sample(source)
 87                }
 88                *a = (dist_f64.round() as i64) << basek_rem;
 89            });
 90        } else {
 91            a.at_mut(res_col, limb)[..n].iter_mut().for_each(|a| {
 92                let mut dist_f64: f64 = dist.sample(source);
 93                while dist_f64.abs() > bound {
 94                    dist_f64 = dist.sample(source)
 95                }
 96                *a = dist_f64.round() as i64
 97            });
 98        }
 99    }
100}
101
102unsafe impl ZnAddDistF64Impl<Self> for FFT64 {
103    fn zn_add_dist_f64_impl<R, D: rand::prelude::Distribution<f64>>(
104        n: usize,
105        basek: usize,
106        res: &mut R,
107        res_col: usize,
108        k: usize,
109        source: &mut Source,
110        dist: D,
111        bound: f64,
112    ) where
113        R: ZnToMut,
114    {
115        let mut a: Zn<&mut [u8]> = res.to_mut();
116        assert!(
117            (bound.log2().ceil() as i64) < 64,
118            "invalid bound: ceil(log2(bound))={} > 63",
119            (bound.log2().ceil() as i64)
120        );
121
122        let limb: usize = k.div_ceil(basek) - 1;
123        let basek_rem: usize = (limb + 1) * basek - k;
124
125        if basek_rem != 0 {
126            a.at_mut(res_col, limb)[..n].iter_mut().for_each(|a| {
127                let mut dist_f64: f64 = dist.sample(source);
128                while dist_f64.abs() > bound {
129                    dist_f64 = dist.sample(source)
130                }
131                *a += (dist_f64.round() as i64) << basek_rem;
132            });
133        } else {
134            a.at_mut(res_col, limb)[..n].iter_mut().for_each(|a| {
135                let mut dist_f64: f64 = dist.sample(source);
136                while dist_f64.abs() > bound {
137                    dist_f64 = dist.sample(source)
138                }
139                *a += dist_f64.round() as i64
140            });
141        }
142    }
143}
144
145unsafe impl ZnFillNormalImpl<Self> for FFT64
146where
147    Self: ZnFillDistF64Impl<Self>,
148{
149    fn zn_fill_normal_impl<R>(
150        n: usize,
151        basek: usize,
152        res: &mut R,
153        res_col: usize,
154        k: usize,
155        source: &mut Source,
156        sigma: f64,
157        bound: f64,
158    ) where
159        R: ZnToMut,
160    {
161        Self::zn_fill_dist_f64_impl(
162            n,
163            basek,
164            res,
165            res_col,
166            k,
167            source,
168            Normal::new(0.0, sigma).unwrap(),
169            bound,
170        );
171    }
172}
173
174unsafe impl ZnAddNormalImpl<Self> for FFT64
175where
176    Self: ZnAddDistF64Impl<Self>,
177{
178    fn zn_add_normal_impl<R>(
179        n: usize,
180        basek: usize,
181        res: &mut R,
182        res_col: usize,
183        k: usize,
184        source: &mut Source,
185        sigma: f64,
186        bound: f64,
187    ) where
188        R: ZnToMut,
189    {
190        Self::zn_add_dist_f64_impl(
191            n,
192            basek,
193            res,
194            res_col,
195            k,
196            source,
197            Normal::new(0.0, sigma).unwrap(),
198            bound,
199        );
200    }
201}