use crate::core::camera::CameraTransform;
use crate::core::color::{white_balance, MatrixMulColor, RGB, SRGB, XYZ};
use crate::core::filter::{Filter, FilterTrait};
use crate::core::geometry::{
    Bounds2f, Bounds2fi, Bounds2i, Normal3f, Point2f, Point2i, Point3f, Tuple, Vector2f, Vector2fi,
    Vector2i, Vector3f,
};
use crate::core::image::{Image, PixelFormat};
use crate::core::interaction::SurfaceInteraction;
use crate::core::spectrum::{Spectrum, SpectrumTrait, StandardSpectra};
use crate::spectra::{
    colorspace, ConstantSpectrum, DenselySampledSpectrum, PiecewiseLinearSpectrum, RGBColorSpace,
    SampledSpectrum, SampledWavelengths, LAMBDA_MAX, LAMBDA_MIN, N_SPECTRUM_SAMPLES,
};
use crate::utils::math::linear_least_squares;
use crate::utils::math::{wrap_equal_area_square, SquareMatrix};
use crate::utils::sampling::VarianceEstimator;
use crate::utils::transform::AnimatedTransform;
use crate::utils::{gpu_array_from_fn, AtomicFloat};
use crate::{gvec_from_slice, gvec_with_capacity, Array2D, Float, GVec, Ptr};
use num_traits::Float as NumFloat;

#[repr(C)]
#[derive(Debug, Clone)]
pub struct RGBFilm {
    pub base: FilmBase,
    pub max_component_value: Float,
    pub write_fp16: bool,
    pub filter_integral: Float,
    pub output_rgbf_from_sensor_rgb: SquareMatrix<Float, 3>,
    pub pixels: Array2D<RGBPixel>,
}

#[repr(C)]
#[derive(Debug, Clone)]
pub struct RGBPixel {
    rgb_sum: [AtomicFloat; 3],
    weight_sum: AtomicFloat,
    rgb_splat: [AtomicFloat; 3],
}

impl Default for RGBPixel {
    fn default() -> Self {
        Self {
            rgb_sum: gpu_array_from_fn(|_| AtomicFloat::default()),
            weight_sum: AtomicFloat::default(),
            rgb_splat: gpu_array_from_fn(|_| AtomicFloat::default()),
        }
    }
}

impl RGBFilm {
    pub fn new(
        base: FilmBase,
        colorspace: &RGBColorSpace,
        max_component_value: Float,
        write_fp16: bool,
    ) -> Self {
        let sensor_ptr = base.sensor;
        // TODO: This wont work on gpu, need to add check on host side
        if sensor_ptr.is_null() {
            panic!("Film must have a sensor");
        }
        let sensor = &*sensor_ptr;
        let filter_integral = base.filter.integral();
        let sensor_matrix = sensor.xyz_from_sensor_rgb;
        let output_rgbf_from_sensor_rgb = colorspace.rgb_from_xyz * sensor_matrix;

        let width = base.pixel_bounds.p_max.x() - base.pixel_bounds.p_min.x();
        let height = base.pixel_bounds.p_max.y() - base.pixel_bounds.p_min.y();
        let count = (width * height) as usize;

        let mut pixel_vec = gvec_with_capacity(count);
        for _ in 0..count {
            pixel_vec.push(RGBPixel::default());
        }

        let pixels: Array2D<RGBPixel> = Array2D::new(base.pixel_bounds);

        RGBFilm {
            base,
            max_component_value,
            write_fp16,
            filter_integral,
            output_rgbf_from_sensor_rgb,
            pixels,
        }
    }

    pub fn base(&self) -> &FilmBase {
        &self.base
    }

    pub fn base_mut(&mut self) -> &mut FilmBase {
        &mut self.base
    }

    pub fn get_sensor(&self) -> &PixelSensor {
        #[cfg(not(target_os = "cuda"))]
        {
            if self.base.sensor.is_null() {
                panic!(
                    "FilmBase error: PixelSensor pointer is null. This should have been checked during construction."
                );
            }
        }
        &self.base.sensor
    }

    pub fn add_sample(
        &self,
        p_film: Point2i,
        l: SampledSpectrum,
        lambda: &SampledWavelengths,
        _vi: Option<&VisibleSurface>,
        weight: Float,
    ) {
        if !self.base.pixel_bounds.contains_exclusive(p_film) {
            return;
        }
        let sensor = self.get_sensor();
        let mut rgb = sensor.to_sensor_rgb(l, lambda);
        let m = rgb.into_iter().copied().fold(f32::NEG_INFINITY, f32::max);
        if m > self.max_component_value {
            rgb *= self.max_component_value / m;
        }

        let pixel = &self.pixels[p_film];
        for c in 0..3 {
            pixel.rgb_sum[c].add(weight * rgb[c as u32]);
        }
        pixel.weight_sum.add(weight);
    }

    pub fn add_splat(&mut self, p: Point2f, l: SampledSpectrum, lambda: &SampledWavelengths) {
        let sensor = self.get_sensor();
        let mut rgb = sensor.to_sensor_rgb(l, lambda);
        let m = rgb.into_iter().copied().fold(f32::NEG_INFINITY, f32::max);
        if m > self.max_component_value {
            rgb *= self.max_component_value / m;
        }

        let p_discrete = p + Vector2f::new(0.5, 0.5);
        let radius = self.base.filter.radius();

        let splat_bounds = Bounds2i::from_points(
            (p_discrete - radius).floor(),
            (p_discrete + radius).floor() + Vector2i::new(1, 1),
        );

        let splat_intersect = splat_bounds.union(self.base().pixel_bounds);
        for pi in &splat_intersect {
            let pi_f: Point2f = (*pi).into();
            let wt = self
                .base()
                .filter
                .evaluate((p - pi_f - Vector2f::new(0.5, 0.5)).into());
            if wt != 0. {
                let pixel = &self.pixels[*pi];
                for i in 0..3 {
                    pixel.rgb_splat[i].add((wt * rgb[i as u32]) as f32);
                }
            }
        }
    }

    pub fn get_pixel_rgb(&self, p: Point2i, splat_scale: Option<Float>) -> RGB {
        let pixel = &self.pixels.get(p);
        let mut rgb = RGB::new(
            pixel.rgb_sum[0].get() as Float,
            pixel.rgb_sum[1].get() as Float,
            pixel.rgb_sum[2].get() as Float,
        );
        let weight_sum = pixel.weight_sum.get();
        if weight_sum != 0. {
            rgb /= weight_sum as Float
        }

        if let Some(splat) = splat_scale {
            for c in 0..3 {
                let splat_val = pixel.rgb_splat[c].get();
                rgb[c] += splat * splat_val as Float / self.filter_integral;
            }
        } else {
            for c in 0..3 {
                let splat_val = pixel.rgb_splat[c].get();
                rgb[c] += splat_val as Float / self.filter_integral;
            }
        }
        self.output_rgbf_from_sensor_rgb.mul_rgb(rgb)
    }

    pub fn to_output_rgb(&self, l: SampledSpectrum, lambda: &SampledWavelengths) -> RGB {
        let sensor = self.get_sensor();
        let sensor_rgb = sensor.to_sensor_rgb(l, lambda);
        self.output_rgbf_from_sensor_rgb.mul_rgb(sensor_rgb)
    }

    fn uses_visible_surface(&self) -> bool {
        false
    }
}

#[repr(C)]
#[derive(Debug, Clone)]
#[cfg_attr(target_os = "cuda", derive(Copy))]
pub struct GBufferPixel {
    pub rgb_sum: [AtomicFloat; 3],
    pub weight_sum: AtomicFloat,
    pub g_buffer_weight_sum: AtomicFloat,
    pub rgb_splat: [AtomicFloat; 3],
    pub p_sum: Point3f,
    pub dz_dx_sum: AtomicFloat,
    pub dz_dy_sum: AtomicFloat,
    pub n_sum: Normal3f,
    pub ns_sum: Normal3f,
    pub uv_sum: Point2f,
    pub rgb_albedo_sum: [AtomicFloat; 3],
    pub rgb_variance: VarianceEstimator,
}

impl Default for GBufferPixel {
    fn default() -> Self {
        Self {
            rgb_sum: gpu_array_from_fn(|_| AtomicFloat::default()),
            weight_sum: AtomicFloat::default(),
            rgb_splat: gpu_array_from_fn(|_| AtomicFloat::default()),
            g_buffer_weight_sum: AtomicFloat::default(),
            p_sum: Point3f::default(),
            dz_dx_sum: AtomicFloat::default(),
            dz_dy_sum: AtomicFloat::default(),
            n_sum: Normal3f::default(),
            ns_sum: Normal3f::default(),
            uv_sum: Point2f::default(),
            rgb_albedo_sum: gpu_array_from_fn(|_| AtomicFloat::default()),
            rgb_variance: VarianceEstimator::default(),
        }
    }
}

#[repr(C)]
#[derive(Debug, Clone)]
#[cfg_attr(target_os = "cuda", derive(Copy))]
pub struct GBufferFilm {
    pub base: FilmBase,
    pub output_from_render: AnimatedTransform,
    pub apply_inverse: bool,
    pub pixels: Array2D<GBufferPixel>,
    pub colorspace: RGBColorSpace,
    pub max_component_value: Float,
    pub write_fp16: bool,
    pub filter_integral: Float,
    pub output_rgbf_from_sensor_rgb: SquareMatrix<Float, 3>,
}

impl GBufferFilm {
    pub fn new(
        base: &FilmBase,
        output_from_render: &AnimatedTransform,
        apply_inverse: bool,
        colorspace: &RGBColorSpace,
        max_component_value: Float,
        write_fp16: bool,
    ) -> Self {
        assert!(!base.pixel_bounds.is_empty());
        let sensor_ptr = base.sensor;
        if sensor_ptr.is_null() {
            panic!("Film must have a sensor");
        }
        let sensor = &*sensor_ptr;
        let output_rgbf_from_sensor_rgb = colorspace.rgb_from_xyz * sensor.xyz_from_sensor_rgb;
        let filter_integral = base.filter.integral();
        let pixels = Array2D::new(base.pixel_bounds);

        GBufferFilm {
            base: base.clone(),
            output_from_render: *output_from_render,
            apply_inverse,
            pixels,
            colorspace: colorspace.clone(),
            max_component_value,
            write_fp16,
            filter_integral,
            output_rgbf_from_sensor_rgb,
        }
    }

    pub fn base(&self) -> &FilmBase {
        &self.base
    }

    pub fn base_mut(&mut self) -> &mut FilmBase {
        &mut self.base
    }

    pub fn get_sensor(&self) -> &PixelSensor {
        #[cfg(not(target_os = "cuda"))]
        {
            if self.base.sensor.is_null() {
                panic!(
                    "FilmBase error: PixelSensor pointer is null. This should have been checked during construction."
                );
            }
        }
        &self.base.sensor
    }

    pub fn add_sample(
        &self,
        _p_film: Point2i,
        _l: SampledSpectrum,
        _lambda: &SampledWavelengths,
        _visible_surface: Option<&VisibleSurface>,
        _weight: Float,
    ) {
        todo!()
    }

    pub fn add_splat(&mut self, p: Point2f, l: SampledSpectrum, lambda: &SampledWavelengths) {
        let sensor = self.get_sensor();
        let mut rgb = sensor.to_sensor_rgb(l, lambda);
        let m = rgb.into_iter().copied().fold(f32::NEG_INFINITY, f32::max);
        if m > self.max_component_value {
            rgb *= self.max_component_value / m;
        }

        let p_discrete = p + Vector2f::new(0.5, 0.5);
        let radius = self.base().filter.radius();

        let splat_bounds = Bounds2i::from_points(
            (p_discrete - radius).floor(),
            (p_discrete + radius).floor() + Vector2i::new(1, 1),
        );

        let splat_intersect = splat_bounds.union(self.base.pixel_bounds);
        for pi in &splat_intersect {
            let pi_f: Point2f = (*pi).into();
            let wt = self
                .base
                .filter
                .evaluate((p - pi_f - Vector2f::new(0.5, 0.5)).into());
            if wt != 0. {
                let pixel = &self.pixels[*pi];
                for i in 0..3 {
                    pixel.rgb_splat[i].add((wt * rgb[i]) as f32);
                }
            }
        }
    }

    pub fn to_output_rgb(&self, l: SampledSpectrum, lambda: &SampledWavelengths) -> RGB {
        let sensor = self.get_sensor();
        let sensor_rgb = sensor.to_sensor_rgb(l, lambda);
        self.output_rgbf_from_sensor_rgb.mul_rgb(sensor_rgb)
    }

    pub fn get_pixel_rgb(&self, p: Point2i, splat_scale: Option<Float>) -> RGB {
        let pixel = &self.pixels.get(p);
        let mut rgb = RGB::new(
            pixel.rgb_sum[0].get() as Float,
            pixel.rgb_sum[1].get() as Float,
            pixel.rgb_sum[2].get() as Float,
        );
        let weight_sum = pixel.weight_sum.get();
        if weight_sum != 0. {
            rgb /= weight_sum as Float
        }

        if let Some(splat) = splat_scale {
            for c in 0..3 {
                let splat_val = pixel.rgb_splat[c].get();
                rgb[c] += splat * splat_val as Float / self.filter_integral;
            }
        } else {
            for c in 0..3 {
                let splat_val = pixel.rgb_splat[c].get();
                rgb[c] += splat_val as Float / self.filter_integral;
            }
        }
        self.output_rgbf_from_sensor_rgb.mul_rgb(rgb)
    }

    pub fn uses_visible_surface(&self) -> bool {
        true
    }
}

#[repr(C)]
#[derive(Debug)]
#[cfg_attr(target_os = "cuda", derive(Copy))]
pub struct SpectralPixel {
    pub rgb_sum: [AtomicFloat; 3],
    pub rgb_weight_sum: AtomicFloat,
    pub rgb_splat: [AtomicFloat; 3],
    pub bucket_offset: usize,
}

impl Clone for SpectralPixel {
    fn clone(&self) -> Self {
        Self {
            rgb_sum: gpu_array_from_fn(|i| AtomicFloat::new(self.rgb_sum[i].get())),
            rgb_weight_sum: AtomicFloat::new(self.rgb_weight_sum.get()),
            rgb_splat: gpu_array_from_fn(|i| AtomicFloat::new(self.rgb_splat[i].get())),
            bucket_offset: self.bucket_offset,
        }
    }
}

impl Default for SpectralPixel {
    fn default() -> Self {
        Self {
            rgb_sum: gpu_array_from_fn(|_| AtomicFloat::new(0.0)),
            rgb_weight_sum: AtomicFloat::new(0.0),
            rgb_splat: gpu_array_from_fn(|_| AtomicFloat::new(0.0)),
            bucket_offset: 0,
        }
    }
}

#[repr(C)]
#[derive(Debug)]
#[cfg_attr(target_os = "cuda", derive(Copy, Clone))]
pub struct SpectralFilm {
    pub base: FilmBase,
    pub lambda_min: Float,
    pub lambda_max: Float,
    pub n_buckets: usize,
    pub max_component_value: Float,
    pub write_fp16: bool,
    pub filter_integral: Float,
    pub colorspace: RGBColorSpace,
    pub pixels: Array2D<SpectralPixel>,
    pub output_rgbf_from_sensor_rgb: SquareMatrix<Float, 3>,
    pub bucket_sums: GVec<f64>,
    pub weight_sums: GVec<f64>,
    pub bucket_splats: GVec<AtomicFloat>,
}

unsafe impl Send for SpectralFilm {}
unsafe impl Sync for SpectralFilm {}

impl SpectralFilm {
    pub fn new(
        base: &FilmBase,
        lambda_min: Float,
        lambda_max: Float,
        n_buckets: usize,
        colorspace: &RGBColorSpace,
        max_component_value: Float,
        write_fp16: bool,
    ) -> Self {
        let n_pixels = base.pixel_bounds.area() as usize;
        let total_buckets = n_pixels * n_buckets;

        let bucket_sums = gvec_with_capacity(total_buckets);
        let weight_sums = gvec_with_capacity(total_buckets);

        let mut bucket_splats = gvec_with_capacity(total_buckets);
        for _ in 0..total_buckets {
            bucket_splats.push(AtomicFloat::new(0.0));
        }

        let mut pixels = Array2D::<SpectralPixel>::new(base.pixel_bounds);

        for i in 0..n_pixels {
            let pixel = pixels.get_linear_mut(i);
            pixel.bucket_offset = i * n_buckets;
        }

        SpectralFilm {
            base: *base,
            colorspace: colorspace.clone(),
            lambda_min,
            lambda_max,
            n_buckets,
            max_component_value,
            write_fp16,
            filter_integral: base.filter.integral(),
            output_rgbf_from_sensor_rgb: SquareMatrix::identity(),
            pixels: Array2D::from_slice(base.pixel_bounds, pixels.as_slice()),
            bucket_sums,
            weight_sums,
            bucket_splats,
        }
    }

    pub fn base(&self) -> &FilmBase {
        &self.base
    }

    pub fn base_mut(&mut self) -> &mut FilmBase {
        &mut self.base
    }

    fn uses_visible_surface(&self) -> bool {
        true
    }

    pub fn add_sample(
        &self,
        _p_film: Point2i,
        _l: SampledSpectrum,
        _lambda: &SampledWavelengths,
        _visible_surface: Option<&VisibleSurface>,
        _weight: Float,
    ) {
        todo!()
    }

    pub fn add_splat(&mut self, _p: Point2f, _v: SampledSpectrum, _lambda: &SampledWavelengths) {
        todo!()
    }

    pub fn get_pixel_rgb(&self, _p: Point2i, _splat_scale: Option<Float>) -> RGB {
        todo!()
    }
}

#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct PixelSensor {
    pub xyz_from_sensor_rgb: SquareMatrix<Float, 3>,
    pub r_bar: Ptr<DenselySampledSpectrum>,
    pub g_bar: Ptr<DenselySampledSpectrum>,
    pub b_bar: Ptr<DenselySampledSpectrum>,
    pub imaging_ratio: Float,
}

impl PixelSensor {
    pub fn project_reflectance<T>(
        refl: &Spectrum,
        illum: &Spectrum,
        b1: &Spectrum,
        b2: &Spectrum,
        b3: &Spectrum,
    ) -> T
    where
        T: From<[Float; 3]>,
    {
        let mut result = [0.; 3];
        let mut g_integral = 0.;

        for lambda_ind in LAMBDA_MIN..=LAMBDA_MAX {
            let lambda = lambda_ind as Float;
            let illum_val = illum.evaluate(lambda);

            g_integral += b2.evaluate(lambda) * illum_val;
            let refl_illum = refl.evaluate(lambda) * illum_val;
            result[0] += b1.evaluate(lambda) * refl_illum;
            result[1] += b2.evaluate(lambda) * refl_illum;
            result[2] += b3.evaluate(lambda) * refl_illum;
        }

        if g_integral > 0. {
            let inv_g = 1. / g_integral;
            result[0] *= inv_g;
            result[1] *= inv_g;
            result[2] *= inv_g;
        }

        T::from([result[0], result[1], result[2]])
    }

    pub fn to_sensor_rgb(&self, l: SampledSpectrum, lambda: &SampledWavelengths) -> RGB {
        let l_norm = SampledSpectrum::safe_div(&l, &lambda.pdf());
        self.imaging_ratio
            * RGB::new(
                (self.r_bar.sample(lambda) * l_norm).average(),
                (self.g_bar.sample(lambda) * l_norm).average(),
                (self.b_bar.sample(lambda) * l_norm).average(),
            )
    }
}

#[repr(C)]
#[derive(Default, Copy, Clone)]
pub struct VisibleSurface {
    pub p: Point3f,
    pub n: Normal3f,
    pub ns: Normal3f,
    pub uv: Point2f,
    pub time: Float,
    pub dpdx: Vector3f,
    pub dpdy: Vector3f,
    pub albedo: SampledSpectrum,
    pub set: bool,
}

impl VisibleSurface {
    pub fn new(
        _si: &SurfaceInteraction,
        albedo: &SampledSpectrum,
        _lambda: &SampledWavelengths,
    ) -> Self {
        VisibleSurface {
            albedo: *albedo,
            ..Default::default()
        }
    }
}

#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct FilmBase {
    pub full_resolution: Point2i,
    pub pixel_bounds: Bounds2i,
    pub filter: Filter,
    pub diagonal: Float,
    pub sensor: Ptr<PixelSensor>,
}

#[repr(C)]
#[derive(Debug)]
#[cfg_attr(target_os = "cuda", derive(Copy, Clone))]
pub enum Film {
    RGB(RGBFilm),
    GBuffer(GBufferFilm),
    Spectral(SpectralFilm),
}

unsafe impl Send for Film {}
unsafe impl Sync for Film {}

impl Film {
    pub fn base(&self) -> &FilmBase {
        match self {
            Film::RGB(f) => f.base(),
            Film::GBuffer(f) => f.base(),
            Film::Spectral(f) => f.base(),
        }
    }

    pub fn base_mut(&mut self) -> &mut FilmBase {
        match self {
            Film::RGB(f) => f.base_mut(),
            Film::GBuffer(f) => f.base_mut(),
            Film::Spectral(f) => f.base_mut(),
        }
    }

    pub fn add_sample(
        &self,
        p_film: Point2i,
        l: SampledSpectrum,
        lambda: &SampledWavelengths,
        visible_surface: Option<&VisibleSurface>,
        weight: Float,
    ) {
        match self {
            Film::RGB(f) => f.add_sample(p_film, l, lambda, visible_surface, weight),
            Film::GBuffer(f) => f.add_sample(p_film, l, lambda, visible_surface, weight),
            Film::Spectral(f) => f.add_sample(p_film, l, lambda, visible_surface, weight),
        }
    }

    pub fn add_splat(&mut self, p: Point2f, v: SampledSpectrum, lambda: &SampledWavelengths) {
        match self {
            Film::RGB(f) => f.add_splat(p, v, lambda),
            Film::GBuffer(f) => f.add_splat(p, v, lambda),
            Film::Spectral(f) => f.add_splat(p, v, lambda),
        }
    }

    pub fn sample_wavelengths(&self, u: Float) -> SampledWavelengths {
        SampledWavelengths::sample_visible(u)
    }

    pub fn uses_visible_surface(&self) -> bool {
        match self {
            Film::RGB(f) => f.uses_visible_surface(),
            Film::GBuffer(f) => f.uses_visible_surface(),
            Film::Spectral(f) => f.uses_visible_surface(),
        }
    }

    pub fn full_resolution(&self) -> Point2i {
        self.base().full_resolution
    }

    pub fn pixel_bounds(&self) -> Bounds2i {
        self.base().pixel_bounds
    }

    pub fn sample_bounds(&self) -> Bounds2f {
        let pixel_bounds = self.pixel_bounds();
        let radius = self.get_filter().radius();
        Bounds2f::from_points(
            Point2f::from(pixel_bounds.p_min) - radius + Vector2f::new(0.5, 0.5),
            Point2f::from(pixel_bounds.p_max) + radius - Vector2f::new(0.5, 0.5),
        )
    }

    pub fn diagonal(&self) -> Float {
        self.base().diagonal
    }

    pub fn get_filter(&self) -> &Filter {
        &self.base().filter
    }

    pub fn get_pixel_rgb(&self, _p: Point2i, _splat_scale: Option<Float>) -> RGB {
        todo!()
    }

    pub fn reset_pixel(&mut self, _p: Point2i) {
        todo!()
    }

    pub fn to_output_rgb(&self, _v: SampledSpectrum, _lambda: &SampledWavelengths) -> RGB {
        todo!()
    }
}