idsp 0.23.0

use super::{atan2, cossin};
use core::num::Wrapping;
use core::ops::{Add, Deref, DerefMut, Div, Mul, Sub};
use dsp_fixedpoint::{Accu, Q, Shift};
use num_traits::{AsPrimitive, One, Zero};

#[cfg(not(feature = "std"))]
#[allow(unused_imports)]
use num_traits::float::Float as _;

/// A complex number in cartesian coordinates
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord)]
#[cfg_attr(feature = "serde", derive(::serde::Serialize, ::serde::Deserialize))]
#[cfg_attr(feature = "bytemuck", derive(bytemuck::Zeroable, bytemuck::Pod))]
#[repr(transparent)]
#[cfg_attr(feature = "serde", serde(transparent))]
pub struct Complex<T>(
    /// Real and imaginary parts
    pub [T; 2],
);

impl<T> Deref for Complex<T> {
    type Target = [T; 2];

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<T> DerefMut for Complex<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

impl<T: Copy> Complex<T> {
    /// Create a new `Complex<T>`
    pub const fn new(re: T, im: T) -> Self {
        Self([re, im])
    }

    /// The real part
    pub fn re(&self) -> T {
        self.0[0]
    }

    /// The imaginary part
    pub fn im(&self) -> T {
        self.0[1]
    }
}

impl<T: Copy + core::ops::Neg<Output = T>> Complex<T> {
    /// Conjugate
    pub fn conj(self) -> Self {
        Self([self.0[0], -self.0[1]])
    }
}

macro_rules! fwd_binop {
    ($tr:ident::$meth:ident) => {
        impl<T: Copy + core::ops::$tr<Output = T>> core::ops::$tr for Complex<T> {
            type Output = Self;
            fn $meth(self, rhs: Self) -> Self {
                Self([self.0[0].$meth(rhs.0[0]), self.0[1].$meth(rhs.0[1])])
            }
        }
    };
}
fwd_binop!(Add::add);
fwd_binop!(Sub::sub);
fwd_binop!(BitAnd::bitand);
fwd_binop!(BitOr::bitor);
fwd_binop!(BitXor::bitxor);

macro_rules! fwd_binop_inner {
    ($tr:ident::$meth:ident) => {
        impl<T: Copy + core::ops::$tr<Output = T>> core::ops::$tr<T> for Complex<T> {
            type Output = Self;
            fn $meth(self, rhs: T) -> Self {
                Self([self.0[0].$meth(rhs), self.0[1].$meth(rhs)])
            }
        }
    };
}
fwd_binop_inner!(Mul::mul);
fwd_binop_inner!(Div::div);
fwd_binop_inner!(Rem::rem);
fwd_binop_inner!(BitAnd::bitand);
fwd_binop_inner!(BitOr::bitor);
fwd_binop_inner!(BitXor::bitxor);

macro_rules! fwd_unop {
    ($tr:ident::$meth:ident) => {
        impl<T: Copy + core::ops::$tr<Output = T>> core::ops::$tr for Complex<T> {
            type Output = Self;
            fn $meth(self) -> Self {
                Self([self.0[0].$meth(), self.0[1].$meth()])
            }
        }
    };
}
fwd_unop!(Not::not);
fwd_unop!(Neg::neg);

impl<T: Copy + Mul<Output = T> + Add<Output = T> + Sub<Output = T>> Mul for Complex<T> {
    type Output = Self;

    fn mul(self, rhs: Self) -> Self {
        Self([
            self.0[0] * rhs.0[0] - self.0[1] * rhs.0[1],
            self.0[0] * rhs.0[1] + self.0[1] * rhs.0[0],
        ])
    }
}

impl<
    T: 'static + Copy,
    A: Copy + Add<Output = A> + Sub<Output = A> + AsPrimitive<T>,
    B,
    const F: i8,
> Mul<Complex<T>> for Complex<Q<T, B, F>>
where
    Q<T, B, F>: Mul<T, Output = A>,
{
    type Output = Complex<T>;

    fn mul(self, rhs: Complex<T>) -> Complex<T> {
        Complex([
            (self.0[0] * rhs.0[0] - self.0[1] * rhs.0[1]).as_(),
            (self.0[0] * rhs.0[1] + self.0[1] * rhs.0[0]).as_(),
        ])
    }
}

impl<
    T: 'static + Copy,
    A: Copy + Add<Output = A> + Sub<Output = A> + AsPrimitive<T>,
    B,
    const F: i8,
> Mul<Complex<Q<T, B, F>>> for Complex<T>
where
    T: Mul<Q<T, B, F>, Output = A>,
{
    type Output = Complex<T>;

    fn mul(self, rhs: Complex<Q<T, B, F>>) -> Complex<T> {
        Complex([
            (self.0[0] * rhs.0[0] - self.0[1] * rhs.0[1]).as_(),
            (self.0[0] * rhs.0[1] + self.0[1] * rhs.0[0]).as_(),
        ])
    }
}

impl<T> core::iter::Sum for Complex<T>
where
    Self: Default + Add<Output = Self>,
{
    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
        iter.fold(Default::default(), |c, i| c + i)
    }
}

impl<T> core::iter::Product for Complex<T>
where
    T: One + Zero,
    Self: Mul<Output = Self>,
{
    fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
        iter.fold(Self([T::one(), T::zero()]), |c, i| c * i)
    }
}

macro_rules! impl_float {
    ($t:ty) => {
        impl Complex<$t> {
            /// Return a unit complex number from a floating-point angle in radians.
            pub fn from_angle(angle: $t) -> Self {
                let (s, c) = angle.sin_cos();
                Self::new(c, s)
            }

            /// Return the angle.
            pub fn arg(&self) -> $t {
                self.im().atan2(self.re())
            }

            /// Return the squared norm.
            pub fn norm_sqr(&self) -> $t {
                self.re() * self.re() + self.im() * self.im()
            }
        }

        impl Div for Complex<$t> {
            type Output = Self;
            fn div(self, rhs: Self) -> Self::Output {
                (self * rhs.conj()) / rhs.norm_sqr()
            }
        }
    };
}
impl_float!(f32);
impl_float!(f64);

impl Complex<i32> {
    /// Return the absolute square (the squared magnitude).
    ///
    /// Panics: on `Complex::new(i32::MIN, i32::MIN)`
    /// Example:
    ///
    /// ```
    /// use idsp::Complex;
    /// assert_eq!(Complex::new(i32::MIN, 0).norm_sqr(), 1 << 62);
    /// ```
    pub fn norm_sqr(&self) -> i64 {
        let [x, y] = self.0.map(|x| x as i64 * x as i64);
        x + y
    }

    /// Return the integer log base 2
    ///
    /// Panics: on [i32::MIN; 2]
    pub fn ilog2(&self) -> u32 {
        self.norm_sqr().ilog2()
    }

    /// Return a Complex on the unit circle given an angle.
    ///
    /// Example:
    ///
    /// ```
    /// use core::num::Wrapping as W;
    /// use idsp::Complex;
    /// Complex::<i32>::from_angle(W(0));
    /// Complex::<i32>::from_angle(W(1 << 30)); // pi/2
    /// Complex::<i32>::from_angle(W(-1 << 30)); // -pi/2
    /// ```
    pub fn from_angle(angle: Wrapping<i32>) -> Self {
        let (c, s) = cossin(angle.0);
        Self::new(c, s)
    }

    /// Return the angle.
    ///
    /// Note: Normalization is `1 << 31 == pi`.
    ///
    /// Example:
    ///
    /// ```
    /// use core::num::Wrapping as W;
    /// use idsp::Complex;
    /// assert_eq!(Complex::new(0, 0).arg(), W(0));
    /// assert_eq!(Complex::new(0, 1).arg(), W(0x3fff_ffff));
    /// ```
    pub fn arg(&self) -> Wrapping<i32> {
        Wrapping(atan2(self.im(), self.re()))
    }
}

impl<A: Shift, T: Accu<A> + Copy, const F: i8> Complex<Q<A, T, F>> {
    /// Quantize each lane back to the base integer type.
    pub fn quantize(self) -> Complex<T> {
        Complex::new(self.re().quantize(), self.im().quantize())
    }
}

impl<T: Copy, A, const F: i8> Complex<Q<T, A, F>> {
    /// Return the raw fixed-point representation lane-wise.
    pub fn into_bits(self) -> Complex<T> {
        Complex::new(self.re().into_bits(), self.im().into_bits())
    }

    /// Convert from raw to fixedpoint
    pub fn from_bits(value: Complex<T>) -> Self {
        Self::new(Q::from_bits(value.re()), Q::from_bits(value.im()))
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use dsp_fixedpoint::Q32;

    #[test]
    fn products() {
        let one = Complex::new(1i32, 0);
        assert_eq!(core::iter::empty().product::<Complex<i32>>(), one);
        let a = Complex::new(2, 3);
        let b = Complex::new(-4, 5);
        assert_eq!([a].into_iter().product::<Complex<i32>>(), a);
        assert_eq!(
            [a, b].into_iter().product::<Complex<i32>>(),
            Complex::new(-23, -2)
        );
        let q = Complex::new(Q32::<28>::from_f32(0.5), Q32::from_int(0));
        assert_eq!(
            [q, q].into_iter().product::<Complex<Q32<28>>>(),
            Complex::new(Q32::from_f32(0.25), Q32::from_int(0))
        );
    }

    #[test]
    fn float_angles() {
        for i in -8..=8 {
            let angle = i as f64 * core::f64::consts::PI / 8.0;
            assert!((Complex::<f64>::from_angle(angle).arg() - angle).abs() < 1e-15);
            let angle = angle as f32;
            let error = Complex::<f32>::from_angle(angle).arg() - angle;
            assert!(error.sin().abs() < 1e-6);
        }
        assert_eq!(Complex::new(1.0f64, 0.0).arg(), 0.0);
        assert!(Complex::new(1.0f64, -0.0).arg().is_sign_negative());
        assert_eq!(Complex::new(-1.0f32, -0.0).arg(), -core::f32::consts::PI);
    }

    #[test]
    fn fixedpoint_into_bits_exposes_raw_representation() {
        let x = Complex::new(
            Q32::<32>::from_bits(0x1234_5678),
            Q32::<32>::from_bits(-0x2345_6789),
        );
        assert_eq!(x.into_bits(), Complex::new(0x1234_5678, -0x2345_6789));
    }

    #[test]
    fn fixedpoint_arg_matches_i32_arg() {
        let z = Complex::new(
            Q::<i64, i32, 32>::from_bits(1 << 34),
            Q::<i64, i32, 32>::from_bits(1 << 34),
        );
        assert_eq!(z.quantize().arg(), Complex::new(4, 4).arg());
    }
}