use std::mem::transmute;

use crate::memory_manager::{
    MemoryStorable,
    type_traits::{FromFloat, UpcastDestTraits},
};

#[inline]
pub fn add_all<M>(sum: &[M]) -> M
where
    M: UpcastDestTraits<M> + MemoryStorable + FromFloat,
{
    if size_of::<M>() == 4 {
        let (prefix, slice, suffix) = unsafe { sum.align_to::<f32>() };
        return M::from_f32(add_all_f32_impl(slice));
    } else if size_of::<M>() == 8 {
        let (prefix, slice, suffix) = unsafe { sum.align_to::<f64>() };
        return M::from_f64(add_all_f64_impl(slice));
    }
    panic!("Size not found");
}

#[inline]
pub fn add_all_f64_impl(sum: &[f64]) -> f64 {
    let len = sum.len();
    if len > 64 {
        let mut acc = 0f64;
        let mut y = 0;
        #[cfg(all(
            any(target_arch = "x86", target_arch = "x86_64"),
            target_feature = "avx512f"
        ))]
        {
            while y + 64 < len {
                unsafe {
                    use std::arch::x86_64::_mm512_load_pd;
                    use std::arch::x86_64::_mm512_reduce_add_pd;

                    let tmp = _mm512_load_pd(&sum[0 + y]);
                    let a = _mm512_reduce_add_pd(tmp);
                    let tmp = _mm512_load_pd(&sum[8 + y]);
                    let b = _mm512_reduce_add_pd(tmp);
                    let tmp = _mm512_load_pd(&sum[16 + y]);
                    let c = _mm512_reduce_add_pd(tmp);
                    let tmp = _mm512_load_pd(&sum[24 + y]);
                    let d = _mm512_reduce_add_pd(tmp);
                    let tmp = _mm512_load_pd(&sum[32 + y]);
                    let e = _mm512_reduce_add_pd(tmp);
                    let tmp = _mm512_load_pd(&sum[40 + y]);
                    let f = _mm512_reduce_add_pd(tmp);
                    let tmp = _mm512_load_pd(&sum[48 + y]);
                    let g = _mm512_reduce_add_pd(tmp);
                    let tmp = _mm512_load_pd(&sum[54 + y]);
                    let h = _mm512_reduce_add_pd(tmp);

                    acc = acc + a + b + c + d + e + f + g + h;
                }
                y += 64;
            }
        }
        #[cfg(all(
            any(target_arch = "x86", target_arch = "x86_64"),
            not(target_feature = "avx512f")
        ))]
        {
            while y + 64 < len {
                let a = sum[y]
                    + sum[1 + y]
                    + sum[2 + y]
                    + sum[3 + y]
                    + sum[4 + y]
                    + sum[5 + y]
                    + sum[6 + y]
                    + sum[7 + y]
                    + sum[8 + y]
                    + sum[9 + y]
                    + sum[10 + y]
                    + sum[11 + y]
                    + sum[12 + y]
                    + sum[13 + y]
                    + sum[14 + y]
                    + sum[15 + y];

                let b = sum[16 + y]
                    + sum[17 + y]
                    + sum[18 + y]
                    + sum[19 + y]
                    + sum[20 + y]
                    + sum[21 + y]
                    + sum[22 + y]
                    + sum[23 + y]
                    + sum[24 + y]
                    + sum[25 + y]
                    + sum[26 + y]
                    + sum[27 + y]
                    + sum[28 + y]
                    + sum[29 + y]
                    + sum[30 + y]
                    + sum[31 + y];

                let c = sum[32 + y]
                    + sum[33 + y]
                    + sum[34 + y]
                    + sum[35 + y]
                    + sum[36 + y]
                    + sum[37 + y]
                    + sum[38 + y]
                    + sum[39 + y]
                    + sum[40 + y]
                    + sum[41 + y]
                    + sum[42 + y]
                    + sum[43 + y]
                    + sum[44 + y]
                    + sum[45 + y]
                    + sum[46 + y]
                    + sum[47 + y];

                let d = sum[48 + y]
                    + sum[49 + y]
                    + sum[50 + y]
                    + sum[51 + y]
                    + sum[52 + y]
                    + sum[53 + y]
                    + sum[54 + y]
                    + sum[55 + y]
                    + sum[56 + y]
                    + sum[57 + y]
                    + sum[58 + y]
                    + sum[59 + y]
                    + sum[60 + y]
                    + sum[61 + y]
                    + sum[62 + y]
                    + sum[63 + y];
                acc = acc + a + b + c + d;
                y += 64;
            }
        }

        for x in y..len {
            acc += sum[x];
        }

        acc
    } else {
        let mut acc = 0f64;
        for y in sum {
            acc += *y;
        }
        acc
    }
}

#[inline]
pub fn add_all_f32_impl(sum: &[f32]) -> f32 {
    let len = sum.len();
    let rem_4 = len % 4;
    let rem_4 = len % 16;
    if len > 64 {
        let mut acc = 0f32;
        let mut y = 0;
        #[cfg(all(
            any(target_arch = "x86", target_arch = "x86_64"),
            target_feature = "avx512f"
        ))]
        {
            while y + 64 < len {
                unsafe {
                    use std::arch::x86_64::{_mm512_load_ps, _mm512_reduce_add_ps};
                    let a = _mm512_load_ps(&sum[0 + y]);
                    let a = _mm512_reduce_add_ps(a);
                    let b = _mm512_load_ps(&sum[16 + y]);
                    let b = _mm512_reduce_add_ps(b);
                    let c = _mm512_load_ps(&sum[32 + y]);
                    let c = _mm512_reduce_add_ps(c);
                    let d = _mm512_load_ps(&sum[48 + y]);
                    let d = _mm512_reduce_add_ps(d);
                    acc = acc + a + b + c + d;
                }
                y += 64;
            }
        }
        #[cfg(all(
            any(target_arch = "x86", target_arch = "x86_64"),
            not(target_feature = "avx512f")
        ))]
        {
            while y + 64 < len {
                let a = sum[y]
                    + sum[1 + y]
                    + sum[2 + y]
                    + sum[3 + y]
                    + sum[4 + y]
                    + sum[5 + y]
                    + sum[6 + y]
                    + sum[7 + y]
                    + sum[8 + y]
                    + sum[9 + y]
                    + sum[10 + y]
                    + sum[11 + y]
                    + sum[12 + y]
                    + sum[13 + y]
                    + sum[14 + y]
                    + sum[15 + y];

                let b = sum[16 + y]
                    + sum[17 + y]
                    + sum[18 + y]
                    + sum[19 + y]
                    + sum[20 + y]
                    + sum[21 + y]
                    + sum[22 + y]
                    + sum[23 + y]
                    + sum[24 + y]
                    + sum[25 + y]
                    + sum[26 + y]
                    + sum[27 + y]
                    + sum[28 + y]
                    + sum[29 + y]
                    + sum[30 + y]
                    + sum[31 + y];

                let c = sum[32 + y]
                    + sum[33 + y]
                    + sum[34 + y]
                    + sum[35 + y]
                    + sum[36 + y]
                    + sum[37 + y]
                    + sum[38 + y]
                    + sum[39 + y]
                    + sum[40 + y]
                    + sum[41 + y]
                    + sum[42 + y]
                    + sum[43 + y]
                    + sum[44 + y]
                    + sum[45 + y]
                    + sum[46 + y]
                    + sum[47 + y];

                let d = sum[48 + y]
                    + sum[49 + y]
                    + sum[50 + y]
                    + sum[51 + y]
                    + sum[52 + y]
                    + sum[53 + y]
                    + sum[54 + y]
                    + sum[55 + y]
                    + sum[56 + y]
                    + sum[57 + y]
                    + sum[58 + y]
                    + sum[59 + y]
                    + sum[60 + y]
                    + sum[61 + y]
                    + sum[62 + y]
                    + sum[63 + y];
                acc = acc + a + b + c + d;
                y += 64;
            }
        }

        for x in y..len {
            acc += sum[x];
        }

        acc
    } else {
        let mut acc = 0f32;
        for y in sum {
            acc += *y;
        }
        acc
    }
}
