use std::{
    f32::consts::PI,
    fmt::Display,
    ops::{Add, Div, Mul, Neg, Sub},
};

#[derive(Copy, Clone, Debug)]
pub struct Complex<T> {
    pub re: T,
    pub im: T,
}

impl<T> Complex<T> {
    pub fn new(re: T, im: T) -> Self {
        Complex { re, im }
    }
}

impl<T: Clone + Add<Output = T>> Add<Complex<T>> for Complex<T> {
    type Output = Complex<T>;

    #[inline]
    fn add(self, rhs: Self) -> Self::Output {
        Self::new(self.re + rhs.re, self.im + rhs.im)
    }
}

impl<T: Clone + Sub<Output = T>> Sub<Complex<T>> for Complex<T> {
    type Output = Complex<T>;

    #[inline]
    fn sub(self, rhs: Self) -> Self::Output {
        Self::new(self.re - rhs.re, self.im - rhs.im)
    }
}

impl<T: Clone + Add<Output = T>> Add<T> for Complex<T> {
    type Output = Complex<T>;

    #[inline]
    fn add(self, rhs: T) -> Self::Output {
        Self::new(self.re + rhs, self.im)
    }
}

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

    #[inline]
    fn mul(self, rhs: Complex<T>) -> Self::Output {
        self::Complex::new(
            self.re.clone() * rhs.re.clone() - self.im.clone() * rhs.im.clone(),
            self.re * rhs.im + self.im * rhs.re,
        )
    }
}

impl<T: Clone + Mul<Output = T>> Mul<T> for Complex<T> {
    type Output = Complex<T>;

    #[inline]
    fn mul(self, rhs: T) -> Self::Output {
        self::Complex::new(self.re * rhs.clone(), self.im * rhs)
    }
}

impl<T: Clone + Neg<Output = T>> Neg for Complex<T> {
    type Output = Complex<T>;

    fn neg(self) -> Self::Output {
        Self::new(-self.re.clone(), -self.im.clone())
    }
}

impl<T: Clone + Neg<Output = T>> Complex<T> {
    pub fn conj(&self) -> Complex<T> {
        Self::new(self.re.clone(), -self.im.clone())
    }
}

impl<T: Display> Display for Complex<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{} + i{}", self.re, self.im)
    }
}

pub type Complex32 = Complex<f32>;
impl Complex32 {
    pub fn zero() -> Self {
        Complex { re: 0.0, im: 0.0 }
    }

    pub fn cexp(angle: f32) -> Complex32 {
        Complex32 {
            re: angle.cos(),
            im: angle.sin(),
        }
    }

    pub fn mag(&self) -> f32 {
        (self.re * self.re + self.im * self.im).sqrt()
    }

    pub fn arg(&self) -> f32 {
        if self.im == 0. {
            if self.re >= 0. { 0. } else { PI }
        } else if self.re == 0. {
            if self.im >= 0. { PI / 2.0 } else { -PI / 2.0 }
        } else {
            (self.im / self.re).atan()
        }
    }
}