use crate::core::pbrt::Float;
use crate::utils::gpu_array_from_fn;
use crate::utils::interval::Interval;
use crate::utils::math::{next_float_down, next_float_up};
use core::ops::{Add, Div, DivAssign, Index, IndexMut, Mul, MulAssign, Neg, Sub};
use num_traits::{Float as NumFloat, FloatConst, Num, One, Signed, Zero};

pub trait Tuple<T, const N: usize>:
    Sized + Copy + Index<usize, Output = T> + IndexMut<usize>
{
    fn data(&self) -> &[T; N];

    fn data_mut(&mut self) -> &mut [T; N];

    fn from_array(arr: [T; N]) -> Self;

    #[inline]
    fn permute(&self, p: [usize; N]) -> Self
    where
        T: Copy,
    {
        let new_data = gpu_array_from_fn(|i| self[p[i]]);
        Self::from_array(new_data)
    }

    fn max_component_value(&self) -> T
    where
        T: PartialOrd + Copy,
    {
        self.data()
            .iter()
            .copied()
            .reduce(|a, b| if a > b { a } else { b })
            .expect("Cannot get max component of a zero-length tuple")
    }

    fn min_component_value(&self) -> T
    where
        T: PartialOrd + Copy,
    {
        self.data()
            .iter()
            .copied()
            .reduce(|a, b| if a < b { a } else { b })
            .expect("Cannot get min component of a zero-length tuple")
    }

    fn max_component_index(&self) -> usize
    where
        T: PartialOrd,
    {
        self.data()
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
            .map(|(index, _)| index)
            .unwrap_or(0)
    }

    fn min_component_index(&self) -> usize
    where
        T: PartialOrd,
    {
        self.data()
            .iter()
            .enumerate()
            .min_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
            .map(|(index, _)| index)
            .unwrap_or(0)
    }
}

pub trait VectorLike:
    Sized
    + Copy
    + Add<Output = Self>
    + Sub<Output = Self>
    + Div<Self::Scalar, Output = Self>
    + Mul<Self::Scalar, Output = Self>
{
    type Scalar: Copy + Zero + Add<Output = Self::Scalar> + Mul<Output = Self::Scalar> + SqrtExt;

    fn dot(self, rhs: Self) -> Self::Scalar;
    fn norm_squared(self) -> Self::Scalar {
        self.dot(self)
    }

    fn abs_dot(self, rhs: Self) -> Self::Scalar
    where
        Self::Scalar: Signed,
    {
        self.dot(rhs).abs()
    }

    fn gram_schmidt(self, rhs: Self) -> Self {
        self - rhs * self.dot(rhs)
    }

    fn norm(&self) -> Self::Scalar {
        self.norm_squared().sqrt_ext()
    }

    fn normalize(self) -> Self
    where
        Self::Scalar: NumFloat,
    {
        let n = self.norm();
        if n.is_zero() { self } else { self / n }
    }

    fn angle_between(self, rhs: Self) -> Self::Scalar
    where
        Self::Scalar: NumFloat,
    {
        let dot_product = self.normalize().dot(rhs.normalize());
        let clamped_dot = dot_product
            .min(Self::Scalar::one())
            .max(-Self::Scalar::one());
        clamped_dot.acos()
    }
}

pub trait SqrtExt {
    fn sqrt_ext(self) -> Self;
}

impl SqrtExt for Float {
    fn sqrt_ext(self) -> Self {
        <Self as num_traits::Float>::sqrt(self)
    }
}

impl SqrtExt for f64 {
    fn sqrt_ext(self) -> Self {
        <Self as num_traits::Float>::sqrt(self)
    }
}

impl SqrtExt for i32 {
    fn sqrt_ext(self) -> Self {
        self.isqrt()
    }
}

impl SqrtExt for u32 {
    fn sqrt_ext(self) -> Self {
        self.isqrt()
    }
}

impl SqrtExt for Interval {
    fn sqrt_ext(self) -> Self {
        let low = if self.low < 0.0 {
            0.0
        } else {
            next_float_down(self.low.sqrt())
        };

        let high = if self.high < 0.0 {
            0.0
        } else {
            next_float_up(self.high.sqrt())
        };

        Self { low, high }
    }
}

pub trait Lerp<Factor = Float>: Sized + Copy {
    fn lerp(t: Factor, a: Self, b: Self) -> Self;
}

impl<T, F, Diff> Lerp<F> for T
where
    T: Copy + Sub<Output = Diff> + Add<Diff, Output = T>,
    Diff: Mul<F, Output = Diff>,
    F: Copy,
{
    #[inline(always)]
    fn lerp(t: F, a: Self, b: Self) -> Self {
        a + (b - a) * t
    }
}
