use crate::core::geometry::{Bounds2f, Bounds2i, Point2f, Point2i, Vector2f};
use crate::filters::*;
use crate::utils::math::{gaussian, gaussian_integral, lerp, sample_tent, windowed_sinc};
use crate::utils::sampling::PiecewiseConstant2D;
use crate::{Array2D, Float, Ptr};
use enum_dispatch::enum_dispatch;

pub struct FilterSample {
    pub p: Point2f,
    pub weight: Float,
}

#[repr(C)]
#[derive(Clone, Debug)]
pub struct FilterSampler {
    pub domain: Bounds2f,
    pub distrib: PiecewiseConstant2D,
    pub f: Array2D<Float>,
}

impl FilterSampler {
    pub fn new<F>(radius: Vector2f, func: F) -> Self
    where
        F: Fn(Point2f) -> Float,
    {
        let domain = Bounds2f::from_points(
            Point2f::new(-radius.x(), -radius.y()),
            Point2f::new(radius.x(), radius.y()),
        );
        let nx = (32.0 * radius.x()) as i32;
        let ny = (32.0 * radius.y()) as i32;
        let mut f = Array2D::new_dims(nx, ny);
        for y in 0..f.y_size() {
            for x in 0..f.x_size() {
                let p = domain.lerp(Point2f::new(
                    (x as Float + 0.5) / f.x_size() as Float,
                    (y as Float + 0.5) / f.y_size() as Float,
                ));
                f[(x as i32, y as i32)] = func(p);
            }
        }
        let distrib = PiecewiseConstant2D::new_with_bounds(&f, domain);
        Self { domain, distrib, f }
    }

    pub fn sample(&self, u: Point2f) -> FilterSample {
        let (p, pdf, pi) = self.distrib.sample(u);
        if pdf == 0.0 {
            return FilterSample { p, weight: 0.0 };
        }
        let idx = pi.x() as usize + pi.y() as usize * self.f.x_size();
        let weight = self.f.as_slice()[idx] / pdf;
        FilterSample { p, weight }
    }
}

#[enum_dispatch]
pub trait FilterTrait {
    fn radius(&self) -> Vector2f;
    fn evaluate(&self, p: Point2f) -> Float;
    fn integral(&self) -> Float;
    fn sample(&self, u: Point2f) -> FilterSample;
}

#[repr(C)]
#[enum_dispatch(FilterTrait)]
#[derive(Clone, Copy, Debug)]
pub enum Filter {
    Box(Ptr<BoxFilter>),
    Gaussian(Ptr<GaussianFilter>),
    Mitchell(Ptr<MitchellFilter>),
    LanczosSinc(Ptr<LanczosSincFilter>),
    Triangle(Ptr<TriangleFilter>),
}

impl<T: FilterTrait> FilterTrait for Ptr<T> {
    fn radius(&self) -> Vector2f {
        unsafe { self.as_ref().radius() }
    }
    fn integral(&self) -> Float {
        unsafe { self.as_ref().integral() }
    }
 
    fn evaluate(&self, p: Point2f) -> Float {
        unsafe { self.as_ref().evaluate(p) }
    }

    fn sample(&self, p: Point2f) -> FilterSample {
        unsafe { self.as_ref().sample(p) }
    }
}
