use crate::simd::intrinsics;
use crate::simd::{LaneCount, Simd, SimdElement, SupportedLaneCount};

/// 通过将选定 lanes 中的元素复制到其他 vectors 中，创建一个新的 SIMD vector。
///
/// swizzling 一个 vector 时，lanes 由 `usize` 的 `const` 数组选择，如 [`Swizzle`]。
///
/// 当 swizzling 两个 vectors 时，lanes 由 [`Which`] 的 `const` 数组选择，如 [`Swizzle2`]。
///
/// # Examples
///
/// 使用单个 SIMD vector，const 数组指定 vector 中的 lane 索引:
///
/// ```
/// # #![feature(portable_simd)]
/// # use core::simd::{u32x2, u32x4, simd_swizzle};
/// let v = u32x4::from_array([10, 11, 12, 13]);
///
/// // 保持相同的大小
/// let r: u32x4 = simd_swizzle!(v, [3, 0, 1, 2]);
/// assert_eq!(r.to_array(), [13, 10, 11, 12]);
///
/// // 改变 lane 数
/// let r: u32x2 = simd_swizzle!(v, [3, 1]);
/// assert_eq!(r.to_array(), [13, 11]);
/// ```
///
/// 使用两个输入 SIMD vectors，const 数组使用 `Which` 来指定每个索引的来源:
///
/// ```
/// # #![feature(portable_simd)]
/// # use core::simd::{u32x2, u32x4, simd_swizzle, Which};
/// use Which::{First, Second};
/// let a = u32x4::from_array([0, 1, 2, 3]);
/// let b = u32x4::from_array([4, 5, 6, 7]);
///
/// // 保持相同的大小
/// let r: u32x4 = simd_swizzle!(a, b, [First(0), First(1), Second(2), Second(3)]);
/// assert_eq!(r.to_array(), [0, 1, 6, 7]);
///
/// // 改变 lane 数
/// let r: u32x2 = simd_swizzle!(a, b, [First(0), Second(0)]);
/// assert_eq!(r.to_array(), [0, 4]);
/// ```
#[allow(unused_macros)]
pub macro simd_swizzle {
    (
        $vector:expr, $index:expr $(,)?
    ) => {
        {
            use $crate::simd::Swizzle;
            struct Impl;
            impl<const LANES: usize> Swizzle<LANES, {$index.len()}> for Impl {
                const INDEX: [usize; {$index.len()}] = $index;
            }
            Impl::swizzle($vector)
        }
    },
    (
        $first:expr, $second:expr, $index:expr $(,)?
    ) => {
        {
            use $crate::simd::{Which, Swizzle2};
            struct Impl;
            impl<const LANES: usize> Swizzle2<LANES, {$index.len()}> for Impl {
                const INDEX: [Which; {$index.len()}] = $index;
            }
            Impl::swizzle2($first, $second)
        }
    }
}

/// 将 lane 索引指定为两个 SIMD vectors 之一。
///
/// 这是 [Swizzle2] 和辅助宏 (如 [simd_swizzle]) 的输入类型。
#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum Which {
    /// 第一个输入 SIMD vector 中 lane 的索引。
    First(usize),
    /// 第二个输入 SIMD vector 中 lane 的索引。
    Second(usize),
}

/// 从另一个 vector 的元素创建一个 vector。
pub trait Swizzle<const INPUT_LANES: usize, const OUTPUT_LANES: usize> {
    /// 从输入 vector 的 lanes 映射到输出 vector。
    const INDEX: [usize; OUTPUT_LANES];

    /// 从 `vector` 的 lanes 创建一个新的 vector。
    ///
    /// Lane `i` 的输出为 `vector[Self::INDEX[i]]`。
    #[inline]
    #[must_use = "method returns a new vector and does not mutate the original inputs"]
    fn swizzle<T>(vector: Simd<T, INPUT_LANES>) -> Simd<T, OUTPUT_LANES>
    where
        T: SimdElement,
        LaneCount<INPUT_LANES>: SupportedLaneCount,
        LaneCount<OUTPUT_LANES>: SupportedLaneCount,
    {
        // 安全性: `vector` 是 vector，`INDEX_IMPL` 是 u32 的 const 数组。
        unsafe { intrinsics::simd_shuffle(vector, vector, Self::INDEX_IMPL) }
    }
}

/// 从另外两个 vectors 的元素创建一个 vector。
pub trait Swizzle2<const INPUT_LANES: usize, const OUTPUT_LANES: usize> {
    /// 从输入 vectors 的 lanes 映射到输出 vector
    const INDEX: [Which; OUTPUT_LANES];

    /// 从 `first` 和 `second` 的 lanes 创建一个新的 vector。
    ///
    /// lane `i` 在 `Self::INDEX[i]` 为 `First(j)` 时为 `first[j]`，在 `Second(j)` 时为 `second[j]`。
    ///
    #[inline]
    #[must_use = "method returns a new vector and does not mutate the original inputs"]
    fn swizzle2<T>(
        first: Simd<T, INPUT_LANES>,
        second: Simd<T, INPUT_LANES>,
    ) -> Simd<T, OUTPUT_LANES>
    where
        T: SimdElement,
        LaneCount<INPUT_LANES>: SupportedLaneCount,
        LaneCount<OUTPUT_LANES>: SupportedLaneCount,
    {
        // 安全性: `first` 和 `second` 是 vectors，`INDEX_IMPL` 是 u32 的 const 数组。
        unsafe { intrinsics::simd_shuffle(first, second, Self::INDEX_IMPL) }
    }
}

/// `simd_shuffle` 内部函数需要 `u32`，因此这里需要进行错误检查和转换。
/// 这个 trait 对公共 API 隐藏了 `INDEX_IMPL`。
trait SwizzleImpl<const INPUT_LANES: usize, const OUTPUT_LANES: usize> {
    const INDEX_IMPL: [u32; OUTPUT_LANES];
}

impl<T, const INPUT_LANES: usize, const OUTPUT_LANES: usize> SwizzleImpl<INPUT_LANES, OUTPUT_LANES>
    for T
where
    T: Swizzle<INPUT_LANES, OUTPUT_LANES> + ?Sized,
{
    const INDEX_IMPL: [u32; OUTPUT_LANES] = {
        let mut output = [0; OUTPUT_LANES];
        let mut i = 0;
        while i < OUTPUT_LANES {
            let index = Self::INDEX[i];
            assert!(index as u32 as usize == index);
            assert!(index < INPUT_LANES, "source lane exceeds input lane count",);
            output[i] = index as u32;
            i += 1;
        }
        output
    };
}

/// `simd_shuffle` 内部函数需要 `u32`，因此这里需要进行错误检查和转换。
/// 这个 trait 对公共 API 隐藏了 `INDEX_IMPL`。
trait Swizzle2Impl<const INPUT_LANES: usize, const OUTPUT_LANES: usize> {
    const INDEX_IMPL: [u32; OUTPUT_LANES];
}

impl<T, const INPUT_LANES: usize, const OUTPUT_LANES: usize> Swizzle2Impl<INPUT_LANES, OUTPUT_LANES>
    for T
where
    T: Swizzle2<INPUT_LANES, OUTPUT_LANES> + ?Sized,
{
    const INDEX_IMPL: [u32; OUTPUT_LANES] = {
        let mut output = [0; OUTPUT_LANES];
        let mut i = 0;
        while i < OUTPUT_LANES {
            let (offset, index) = match Self::INDEX[i] {
                Which::First(index) => (false, index),
                Which::Second(index) => (true, index),
            };
            assert!(index < INPUT_LANES, "source lane exceeds input lane count",);

            // lanes 由第一个 vector 索引，然后是第二个 vector
            let index = if offset { index + INPUT_LANES } else { index };
            assert!(index as u32 as usize == index);
            output[i] = index as u32;
            i += 1;
        }
        output
    };
}

impl<T, const LANES: usize> Simd<T, LANES>
where
    T: SimdElement,
    LaneCount<LANES>: SupportedLaneCount,
{
    /// 反转 vector 中 lanes 的顺序。
    #[inline]
    #[must_use = "method returns a new vector and does not mutate the original inputs"]
    pub fn reverse(self) -> Self {
        const fn reverse_index<const LANES: usize>() -> [usize; LANES] {
            let mut index = [0; LANES];
            let mut i = 0;
            while i < LANES {
                index[i] = LANES - i - 1;
                i += 1;
            }
            index
        }

        struct Reverse;

        impl<const LANES: usize> Swizzle<LANES, LANES> for Reverse {
            const INDEX: [usize; LANES] = reverse_index::<LANES>();
        }

        Reverse::swizzle(self)
    }

    /// 旋转 vector，使切片的第一个 `OFFSET` 元素移动到末尾，而最后一个 `LANES - OFFSET` 元素移动到前面。
    ///
    /// 调用 `rotate_lanes_left` 后，先前在 `OFFSET` lane 中的元素将成为该 lane 中的第一个元素。
    #[inline]
    #[must_use = "method returns a new vector and does not mutate the original inputs"]
    pub fn rotate_lanes_left<const OFFSET: usize>(self) -> Self {
        const fn rotate_index<const OFFSET: usize, const LANES: usize>() -> [usize; LANES] {
            let offset = OFFSET % LANES;
            let mut index = [0; LANES];
            let mut i = 0;
            while i < LANES {
                index[i] = (i + offset) % LANES;
                i += 1;
            }
            index
        }

        struct Rotate<const OFFSET: usize>;

        impl<const OFFSET: usize, const LANES: usize> Swizzle<LANES, LANES> for Rotate<OFFSET> {
            const INDEX: [usize; LANES] = rotate_index::<OFFSET, LANES>();
        }

        Rotate::<OFFSET>::swizzle(self)
    }

    /// 旋转 vector，使 vector 的第一个 `LANES - OFFSET` 元素移动到末尾，而最后一个 `OFFSET` 元素移动到前面。
    ///
    /// 调用 `rotate_lanes_right` 后，之前位于索引 `LANES - OFFSET` 的元素将成为切片中的第一个元素。
    #[inline]
    #[must_use = "method returns a new vector and does not mutate the original inputs"]
    pub fn rotate_lanes_right<const OFFSET: usize>(self) -> Self {
        const fn rotate_index<const OFFSET: usize, const LANES: usize>() -> [usize; LANES] {
            let offset = LANES - OFFSET % LANES;
            let mut index = [0; LANES];
            let mut i = 0;
            while i < LANES {
                index[i] = (i + offset) % LANES;
                i += 1;
            }
            index
        }

        struct Rotate<const OFFSET: usize>;

        impl<const OFFSET: usize, const LANES: usize> Swizzle<LANES, LANES> for Rotate<OFFSET> {
            const INDEX: [usize; LANES] = rotate_index::<OFFSET, LANES>();
        }

        Rotate::<OFFSET>::swizzle(self)
    }

    /// 交错两个 vectors。
    ///
    /// 生成的 vectors 包含交替取自 `self` 和 `other` 的 lanes，首先填充第一个结果，然后填充第二个结果。
    ///
    ///
    /// 这个操作的反面是 [`Simd::deinterleave`]。
    ///
    /// ```
    /// # #![feature(portable_simd)]
    /// # use core::simd::Simd;
    /// let a = Simd::from_array([0, 1, 2, 3]);
    /// let b = Simd::from_array([4, 5, 6, 7]);
    /// let (x, y) = a.interleave(b);
    /// assert_eq!(x.to_array(), [0, 4, 1, 5]);
    /// assert_eq!(y.to_array(), [2, 6, 3, 7]);
    /// ```
    #[inline]
    #[must_use = "method returns a new vector and does not mutate the original inputs"]
    pub fn interleave(self, other: Self) -> (Self, Self) {
        const fn interleave<const LANES: usize>(high: bool) -> [Which; LANES] {
            let mut idx = [Which::First(0); LANES];
            let mut i = 0;
            while i < LANES {
                // 将源视为串联的 vector
                let dst_index = if high { i + LANES } else { i };
                let src_index = dst_index / 2 + (dst_index % 2) * LANES;
                idx[i] = if src_index < LANES {
                    Which::First(src_index)
                } else {
                    Which::Second(src_index % LANES)
                };
                i += 1;
            }
            idx
        }

        struct Lo;
        struct Hi;

        impl<const LANES: usize> Swizzle2<LANES, LANES> for Lo {
            const INDEX: [Which; LANES] = interleave::<LANES>(false);
        }

        impl<const LANES: usize> Swizzle2<LANES, LANES> for Hi {
            const INDEX: [Which; LANES] = interleave::<LANES>(true);
        }

        (Lo::swizzle2(self, other), Hi::swizzle2(self, other))
    }

    /// 去交错两个 vectors。
    ///
    /// 第一个结果取 `self` 的每隔一个 lane，然后是 `other`，从第一个 lane 开始。
    ///
    /// 第二个结果取 `self` 的每隔一个 lane，然后是 `other`，从第二个 lane 开始。
    ///
    ///
    /// 这个操作的反面是 [`Simd::interleave`]。
    ///
    /// ```
    /// # #![feature(portable_simd)]
    /// # use core::simd::Simd;
    /// let a = Simd::from_array([0, 4, 1, 5]);
    /// let b = Simd::from_array([2, 6, 3, 7]);
    /// let (x, y) = a.deinterleave(b);
    /// assert_eq!(x.to_array(), [0, 1, 2, 3]);
    /// assert_eq!(y.to_array(), [4, 5, 6, 7]);
    /// ```
    ///
    #[inline]
    #[must_use = "method returns a new vector and does not mutate the original inputs"]
    pub fn deinterleave(self, other: Self) -> (Self, Self) {
        const fn deinterleave<const LANES: usize>(second: bool) -> [Which; LANES] {
            let mut idx = [Which::First(0); LANES];
            let mut i = 0;
            while i < LANES {
                // 将源视为串联的 vector
                let src_index = i * 2 + second as usize;
                idx[i] = if src_index < LANES {
                    Which::First(src_index)
                } else {
                    Which::Second(src_index % LANES)
                };
                i += 1;
            }
            idx
        }

        struct Even;
        struct Odd;

        impl<const LANES: usize> Swizzle2<LANES, LANES> for Even {
            const INDEX: [Which; LANES] = deinterleave::<LANES>(false);
        }

        impl<const LANES: usize> Swizzle2<LANES, LANES> for Odd {
            const INDEX: [Which; LANES] = deinterleave::<LANES>(true);
        }

        (Even::swizzle2(self, other), Odd::swizzle2(self, other))
    }
}