use super::*;
use crate::core::film::{CreateFilmBase, PixelSensor};
use crate::utils::containers::Array2D;
use crate::{Arena, FileLoc, ParameterDictionary};
use anyhow::{Result, anyhow};
use shared::Float;
use shared::core::camera::CameraTransform;
use shared::core::film::{FilmBase, SpectralFilm, SpectralPixel};
use shared::core::filter::FilterTrait;
use shared::spectra::{LAMBDA_MAX, LAMBDA_MIN, RGBColorSpace};
use shared::utils::AtomicFloat;
use shared::utils::containers::DeviceArray2D;
use shared::utils::math::SquareMatrix;
use std::path::Path;
use std::sync::Arc;

struct SpectralFilmStorage {
    pixels: DeviceArray2D<SpectralPixel>,
    bucket_sums: Vec<f64>,
    weight_sums: Vec<f64>,
    bucket_splats: Vec<AtomicFloat>,
}

pub struct SpectralFilmHost {
    pub device: SpectralFilm,
    storage: Arc<SpectralFilmStorage>,
}

impl SpectralFilmHost {
    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 = vec![0.0; total_buckets];
        let weight_sums = vec![0.0; total_buckets];

        let mut bucket_splats = Vec::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;
        }

        let storage = Arc::new(SpectralFilmStorage {
            pixels: pixels.device,
            bucket_sums,
            weight_sums,
            bucket_splats,
        });

        let device = 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: DeviceArray2D {
                values: pixels.values.as_mut_ptr(),
                extent: base.pixel_bounds,
                stride: base.pixel_bounds.p_max.x() - base.pixel_bounds.p_min.x(),
            },

            bucket_sums: storage.bucket_sums.as_ptr() as *mut f64,
            weight_sums: storage.weight_sums.as_ptr() as *mut f64,
            bucket_splats: storage.bucket_splats.as_ptr() as *mut AtomicFloat,
        };

        Self { device, storage }
    }
}

impl CreateFilm for SpectralFilm {
    fn create(
        params: &ParameterDictionary,
        exposure_time: Float,
        filter: Filter,
        _camera_transform: Option<CameraTransform>,
        loc: &FileLoc,
        _arena: &Arena,
    ) -> Result<Film> {
        let colorspace = params.color_space.as_ref().unwrap();
        let max_component_value = params.get_one_float("maxcomponentvalue", Float::INFINITY);
        let write_fp16 = params.get_one_bool("savefp16", true);
        let sensor = PixelSensor::create(params, colorspace.clone(), exposure_time, loc)?;
        let film_base = FilmBase::create(params, filter, Some(&sensor.device()), loc);

        let filename = params.get_one_string("filename", "pbrt.exr");
        if Path::new(&filename).extension() != Some("exr".as_ref()) {
            return Err(anyhow!("{}: EXR is the only format supported by GBufferFilm", loc).into());
        }

        let n_buckets = params.get_one_int("nbuckets", 16) as usize;
        let lambda_min = params.get_one_float("lambdamin", LAMBDA_MIN as Float);
        let lambda_max = params.get_one_float("lambdamin", LAMBDA_MAX as Float);
        if lambda_min < LAMBDA_MIN as Float && lambda_max > LAMBDA_MAX as Float {
            return Err(anyhow!(
                "{}: PBRT must be recompiled with different values of LAMBDA_MIN and LAMBDA_MAX",
                loc
            ));
        }

        let film = SpectralFilmHost::new(
            &film_base,
            lambda_min,
            lambda_max,
            n_buckets,
            colorspace,
            max_component_value,
            write_fp16,
        );

        Ok(Film::Spectral(film.device))
    }
}
