use crate::core::geometry::{
    Bounds2i, Bounds3f, Bounds3i, Point2i, Point3f, Point3i, Vector2i, Vector3f, Vector3i,
};
use crate::utils::math::lerp;
use crate::{gvec, gvec_from_slice, Float, GVec};
use core::ops::{Add, Index, IndexMut, Mul, Sub};

pub trait Interpolatable:
    Copy + Default + Add<Output = Self> + Sub<Output = Self> + Mul<Float, Output = Self>
{
}

impl<T> Interpolatable for T where
    T: Copy + Default + Add<Output = T> + Sub<Output = T> + Mul<Float, Output = T>
{
}

#[derive(Debug, Clone)]
pub struct Array2D<T> {
    extent: Bounds2i,
    values: GVec<T>,
}

impl<T> Array2D<T> {
    pub fn extent(&self) -> Bounds2i {
        self.extent
    }
    pub fn stride(&self) -> i32 {
        self.extent.p_max.x() - self.extent.p_min.x()
    }
    pub fn x_size(&self) -> usize {
        self.stride() as usize
    }
    pub fn y_size(&self) -> usize {
        (self.extent.p_max.y() - self.extent.p_min.y()) as usize
    }
    pub fn as_slice(&self) -> &[T] {
        &self.values
    }
    pub fn as_mut_slice(&mut self) -> &mut [T] {
        &mut self.values
    }
    pub fn as_ptr(&self) -> *const T {
        self.values.as_ptr()
    }
    pub fn as_mut_ptr(&mut self) -> *mut T {
        self.values.as_mut_ptr()
    }
    pub fn len(&self) -> usize {
        self.values.len()
    }
}

unsafe impl<T: Send> Send for Array2D<T> {}
unsafe impl<T: Sync> Sync for Array2D<T> {}

impl<T: Default + Clone> Array2D<T> {
    pub fn new(extent: Bounds2i) -> Self {
        let n = extent.area() as usize;
        let mut values = gvec();
        values.resize(n, T::default());
        Self { extent, values }
    }

    pub fn new_dims(nx: i32, ny: i32) -> Self {
        Self::new(Bounds2i::from_points(
            Point2i::new(0, 0),
            Point2i::new(nx, ny),
        ))
    }

    pub fn new_filled(nx: i32, ny: i32, val: T) -> Self {
        let extent = Bounds2i::from_points(Point2i::new(0, 0), Point2i::new(nx, ny));
        let n = (nx * ny) as usize;
        let mut values = gvec();
        values.resize(n, val);
        Self { extent, values }
    }

      pub fn from_slice(extent: Bounds2i, slice: &[T]) -> Self {
        assert_eq!(slice.len(), extent.area() as usize);
        Self {
            extent,
            values: gvec_from_slice(slice),
        }
    }

    pub fn get(&self, p: Point2i) -> &T {
        &self[p]
    }

    pub fn get_linear_mut(&mut self, index: usize) -> &mut T {
        unsafe { &mut *self.values.as_mut_ptr().add(index) }
    }
}

impl<T> Index<(i32, i32)> for Array2D<T> {
    type Output = T;
    fn index(&self, (x, y): (i32, i32)) -> &T {
        let offset = (y - self.extent.p_min.y()) * self.stride() + (x - self.extent.p_min.x());
        &self.values[offset as usize]
    }
}

impl<T> IndexMut<(i32, i32)> for Array2D<T> {
    fn index_mut(&mut self, (x, y): (i32, i32)) -> &mut T {
        let offset = (y - self.extent.p_min.y()) * self.stride() + (x - self.extent.p_min.x());
        &mut self.values[offset as usize]
    }
}

impl<T> Index<Point2i> for Array2D<T> {
    type Output = T;
    fn index(&self, p: Point2i) -> &T {
        let offset =
            (p.y() - self.extent.p_min.y()) * self.stride() + (p.x() - self.extent.p_min.x());
        &self.values[offset as usize]
    }
}

impl<T> IndexMut<Point2i> for Array2D<T> {
    fn index_mut(&mut self, p: Point2i) -> &mut T {
        let offset =
            (p.y() - self.extent.p_min.y()) * self.stride() + (p.x() - self.extent.p_min.x());
        &mut self.values[offset as usize]
    }
}

#[repr(C)]
#[derive(Debug, Clone)]
pub struct SampledGrid<T> {
    pub values: GVec<T>,
    pub values_len: u32,
    pub nx: i32,
    pub ny: i32,
    pub nz: i32,
}

unsafe impl<T: Sync> Sync for SampledGrid<T> {}
unsafe impl<T: Send> Send for SampledGrid<T> {}

impl<T: core::clone::Clone> SampledGrid<T> {
    pub fn new(slice: &[T], nx: i32, ny: i32, nz: i32) -> Self {
        assert_eq!(slice.len(), (nx * ny * nz) as usize);
        Self {
            values: gvec_from_slice(slice),
            values_len: (nx * ny * nz) as u32,
            nx,
            ny,
            nz,
        }
    }

    pub fn empty() -> Self {
        Self {
            values: gvec(),
            values_len: 0,
            nx: 0,
            ny: 0,
            nz: 0,
        }
    }

    pub fn is_valid(&self) -> bool {
        !self.values.is_empty() && self.nx > 0 && self.ny > 0 && self.nz > 0
    }

    pub fn bytes_allocated(&self) -> u32 {
        self.values_len * core::mem::size_of::<T>() as u32
    }

    pub fn x_size(&self) -> i32 {
        self.nx
    }

    pub fn y_size(&self) -> i32 {
        self.ny
    }

    pub fn z_size(&self) -> i32 {
        self.nz
    }

    pub fn lookup_int_convert<F, U>(&self, p: Point3i, convert: F) -> U
    where
        F: Fn(&T) -> U,
        U: Default,
    {
        if !self.is_valid() {
            return U::default();
        }

        let sample_bounds = Bounds3i::from_points(
            Point3i::new(0, 0, 0),
            Point3i::new(self.nx, self.ny, self.nz),
        );

        if !sample_bounds.contains_exclusive(p) {
            return U::default();
        }

        let idx = (p.z() * self.ny + p.y()) * self.nx + p.x();
        unsafe {
            let val = &*self.values.as_ptr().add(idx as usize);
            convert(val)
        }
    }

    pub fn lookup_int(&self, p: Point3i) -> T
    where
        T: Clone + Default,
    {
        self.lookup_int_convert(p, |v| v.clone())
    }

    pub fn lookup_convert<F, U>(&self, p: Point3f, convert: F) -> U
    where
        F: Fn(&T) -> U + Copy,
        U: Interpolatable + Default,
    {
        if !self.is_valid() {
            return U::default();
        }

        let p_samples = Point3f::new(
            p.x() * self.nx as Float - 0.5,
            p.y() * self.ny as Float - 0.5,
            p.z() * self.nz as Float - 0.5,
        );

        let pi = Point3i::from(p_samples.floor());
        let d = p_samples - Point3f::from(pi);

        // Helper to retrieve corners with conversion
        let lk = |offset: Vector3i| -> U { self.lookup_int_convert(pi + offset, convert) };

        // Trilinear interpolation
        let d00 = lerp(
            d.x(),
            lk(Vector3i::new(0, 0, 0)),
            lk(Vector3i::new(1, 0, 0)),
        );
        let d10 = lerp(
            d.x(),
            lk(Vector3i::new(0, 1, 0)),
            lk(Vector3i::new(1, 1, 0)),
        );
        let d01 = lerp(
            d.x(),
            lk(Vector3i::new(0, 0, 1)),
            lk(Vector3i::new(1, 0, 1)),
        );
        let d11 = lerp(
            d.x(),
            lk(Vector3i::new(0, 1, 1)),
            lk(Vector3i::new(1, 1, 1)),
        );

        lerp(d.z(), lerp(d.y(), d00, d10), lerp(d.y(), d01, d11))
    }

    pub fn lookup(&self, p: Point3f) -> T
    where
        T: Interpolatable + Clone,
    {
        self.lookup_convert(p, |v| *v)
    }

    pub fn max_value_convert<F, U>(&self, bounds: Bounds3f, convert: F) -> U
    where
        F: Fn(&T) -> U,
        U: PartialOrd + Default,
    {
        if !self.is_valid() {
            return U::default();
        }

        let ps = [
            Point3f::new(
                bounds.p_min.x() * self.nx as Float - 0.5,
                bounds.p_min.y() * self.ny as Float - 0.5,
                bounds.p_min.z() * self.nz as Float - 0.5,
            ),
            Point3f::new(
                bounds.p_max.x() * self.nx as Float - 0.5,
                bounds.p_max.y() * self.ny as Float - 0.5,
                bounds.p_max.z() * self.nz as Float - 0.5,
            ),
        ];

        let pi_min = Point3i::from(ps[0].floor()).max(Point3i::new(0, 0, 0));
        let pi_max = (Point3i::from(ps[1].floor()) + Vector3i::new(1, 1, 1)).min(Point3i::new(
            self.nx - 1,
            self.ny - 1,
            self.nz - 1,
        ));

        let mut max_value = self.lookup_int_convert(pi_min, &convert);

        for z in pi_min.z()..=pi_max.z() {
            for y in pi_min.y()..=pi_max.y() {
                for x in pi_min.x()..=pi_max.x() {
                    let val = self.lookup_int_convert(Point3i::new(x, y, z), &convert);
                    if val > max_value {
                        max_value = val;
                    }
                }
            }
        }

        max_value
    }

    pub fn max_value(&self, bounds: Bounds3f) -> T
    where
        T: PartialOrd + Default + Clone,
    {
        self.max_value_convert(bounds, |v| v.clone())
    }
}
