use crate::Float;
use num_traits::Float as NumFloat;

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

impl Complex {
    pub fn new(re: Float, im: Float) -> Self {
        Self { re, im }
    }
    pub fn norm(&self) -> Float {
        (self.re * self.re + self.im * self.im).sqrt()
    }
    pub fn sqrt(self) -> Self {
        let r = self.norm();
        let re = ((r + self.re) / 2.0).sqrt();
        let im = self.im.signum() * ((r - self.re) / 2.0).sqrt();
        Self { re, im }
    }
}

impl core::ops::Add for Complex {
    type Output = Self;
    fn add(self, rhs: Self) -> Self {
        Self {
            re: self.re + rhs.re,
            im: self.im + rhs.im,
        }
    }
}

impl core::ops::Mul for Complex {
    type Output = Self;
    fn mul(self, rhs: Self) -> Self {
        Self {
            re: self.re * rhs.re - self.im * rhs.im,
            im: self.re * rhs.im + self.im * rhs.re,
        }
    }
}

impl core::ops::Div for Complex {
    type Output = Self;
    fn div(self, rhs: Self) -> Self {
        let denom = rhs.re * rhs.re + rhs.im * rhs.im;
        Self {
            re: (self.re * rhs.re + self.im * rhs.im) / denom,
            im: (self.im * rhs.re - self.re * rhs.im) / denom,
        }
    }
}

impl core::ops::Sub for Complex {
    type Output = Self;
    fn sub(self, rhs: Self) -> Self {
        Self {
            re: self.re - rhs.re,
            im: self.im - rhs.im,
        }
    }
}

impl core::ops::Mul<Complex> for Float {
    type Output = Complex;
    fn mul(self, rhs: Complex) -> Complex {
        Complex {
            re: self * rhs.re,
            im: self * rhs.im,
        }
    }
}

impl core::ops::Mul<Float> for Complex {
    type Output = Complex;
    fn mul(self, rhs: Float) -> Complex {
        Complex {
            re: self.re * rhs,
            im: self.im * rhs,
        }
    }
}

impl core::ops::Sub<Complex> for Float {
    type Output = Complex;
    fn sub(self, rhs: Complex) -> Complex {
        Complex {
            re: self - rhs.re,
            im: -rhs.im,
        }
    }
}

impl core::ops::Add<Complex> for Float {
    type Output = Complex;
    fn add(self, rhs: Complex) -> Complex {
        Complex {
            re: self + rhs.re,
            im: rhs.im,
        }
    }
}

impl core::ops::Div<Complex> for Float {
    type Output = Complex;
    fn div(self, rhs: Complex) -> Complex {
        let denom = rhs.re * rhs.re + rhs.im * rhs.im;
        Complex {
            re: self * rhs.re / denom,
            im: -self * rhs.im / denom,
        }
    }
}
