use crate::prelude::*;
use std::ops::{Add, Div, Index, Mul, Neg, Sub};

#[derive(Debug, Clone, Copy)]
pub struct Dir3 {
    pub(crate) x: Float,
    pub(crate) y: Float,
    pub(crate) z: Float,
}

impl Dir3 {
    pub fn new(x: Float, y: Float, z: Float) -> Self {
        Self { x, y, z }
    }

    pub fn x(self) -> Float {
        self.x
    }

    pub fn y(self) -> Float {
        self.y
    }

    pub fn z(self) -> Float {
        self.z
    }

    pub fn zero() -> Self {
        Self {
            x: 0.0,
            y: 0.0,
            z: 0.0,
        }
    }

    pub fn up() -> Self {
        Self::new(0.0, 1.0, 0.0)
    }

    pub fn normalize(self) -> Self {
        self / self.length()
    }

    pub fn length_squared(self) -> Float {
        self.x * self.x + self.y * self.y + self.z * self.z
    }

    pub fn length(self) -> Float {
        Float::sqrt(self.length_squared())
    }

    pub fn dot(self, rhs: Self) -> Float {
        self.x * rhs.x + self.y * rhs.y + self.z * rhs.z
    }

    pub fn cross(self, rhs: Self) -> Self {
        Self::new(
            self.y * rhs.z - self.z * rhs.y,
            self.z * rhs.x - self.x * rhs.z,
            self.x * rhs.y - self.y * rhs.x,
        )
    }

    pub fn sample_uniform_hemisphere<R: Rng>(rng: &mut R) -> Self {
        let theta = Float::acos(1.0 - rng.gen::<Float>());
        let phi = 2.0 * FloatConsts::PI * rng.gen::<Float>();
        let r = Float::sin(theta);
        Dir3::new(Float::sin(phi) * r, Float::cos(theta), Float::cos(phi) * r)
    }

    pub fn sample_uniform_sphere<R: Rng>(rng: &mut R) -> Self {
        let theta = Float::acos(1.0 - 2.0 * rng.gen::<Float>());
        let phi = 2.0 * FloatConsts::PI * rng.gen::<Float>();
        let r = Float::sin(theta);
        Dir3::new(Float::sin(phi) * r, Float::cos(theta), Float::cos(phi) * r)
    }

    pub fn sample_cosine_hemisphere<R: Rng>(rng: &mut R) -> Self {
        let mut d = Self::sample_disc(rng);
        d.y = Float::sqrt(Float::max(0.0, 1.0 - d.x * d.x - d.z * d.z));
        d
    }

    /// https://www.pbr-book.org/3ed-2018/Monte_Carlo_Integration/2D_Sampling_with_Multidimensional_Transformations#ConcentricSampleDisk
    pub fn sample_disc<R: Rng>(rng: &mut R) -> Self {
        let x = 2.0 * rng.gen::<Float>() - 1.0;
        let y = 2.0 * rng.gen::<Float>() - 1.0;

        if x == 0.0 && y == 0.0 {
            return Self::zero();
        }

        let r;
        let theta;
        if x.abs() > y.abs() {
            r = x;
            theta = FloatConsts::FRAC_PI_4 * (y / x);
        } else {
            r = y;
            theta = FloatConsts::FRAC_PI_2 - FloatConsts::FRAC_PI_4 * (x / y);
        }

        r * Dir3::new(Float::cos(theta), 0.0, Float::sin(theta))
    }
}

impl Add for Dir3 {
    type Output = Dir3;

    fn add(self, rhs: Self) -> Self::Output {
        Dir3 {
            x: self.x + rhs.x,
            y: self.y + rhs.y,
            z: self.z + rhs.z,
        }
    }
}

impl Add<pos3::Pos3> for Dir3 {
    type Output = pos3::Pos3;

    fn add(self, rhs: pos3::Pos3) -> Self::Output {
        pos3::Pos3 {
            x: self.x + rhs.x,
            y: self.y + rhs.y,
            z: self.z + rhs.z,
        }
    }
}

impl Sub for Dir3 {
    type Output = Dir3;

    fn sub(self, rhs: Self) -> Self::Output {
        Dir3 {
            x: self.x - rhs.x,
            y: self.y - rhs.y,
            z: self.z - rhs.z,
        }
    }
}

impl Mul<Float> for Dir3 {
    type Output = Dir3;

    fn mul(self, rhs: Float) -> Self::Output {
        Dir3 {
            x: self.x * rhs,
            y: self.y * rhs,
            z: self.z * rhs,
        }
    }
}

impl Mul<Dir3> for Float {
    type Output = Dir3;

    fn mul(self, rhs: Dir3) -> Self::Output {
        Dir3 {
            x: self * rhs.x,
            y: self * rhs.y,
            z: self * rhs.z,
        }
    }
}

impl Div<Float> for Dir3 {
    type Output = Dir3;

    fn div(self, rhs: Float) -> Self::Output {
        Dir3 {
            x: self.x / rhs,
            y: self.y / rhs,
            z: self.z / rhs,
        }
    }
}

impl Neg for Dir3 {
    type Output = Dir3;

    fn neg(self) -> Self::Output {
        Dir3 {
            x: -self.x,
            y: -self.y,
            z: -self.z,
        }
    }
}

impl Index<usize> for Dir3 {
    type Output = Float;

    fn index(&self, index: usize) -> &Self::Output {
        match index {
            0 => &self.x,
            1 => &self.y,
            2 => &self.z,
            _ => panic!("Index out of bound {}", index),
        }
    }
}
