use crate::core::image::{HostImage, ImageChannelDesc, ImageChannelValues, ImageIO, ImageMetadata};
use crate::films::*;
use crate::spectra::data::get_named_spectrum;
use anyhow::{anyhow, Result};
use rayon::iter::ParallelIterator;
use rayon::prelude::IntoParallelIterator;
use shared::core::camera::CameraTransform;
use shared::core::color::{white_balance, RGB, SRGB, XYZ};
use shared::core::film::{Film, FilmBase, GBufferFilm, PixelSensor, RGBFilm, SpectralFilm};
use shared::core::filter::{Filter, FilterTrait};
use shared::core::geometry::{Bounds2f, Bounds2i, Point2f, Point2i};
use shared::core::image::PixelFormat;
use shared::core::spectrum::Spectrum;
use shared::spectra::{
    cie::SWATCHES_RAW, DenselySampledSpectrum, PiecewiseLinearSpectrum, RGBColorSpace,
};
use shared::utils::math::{linear_least_squares, SquareMatrix};
use shared::{Float, Ptr, leak};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, LazyLock};

use crate::spectra::{get_spectra_context, CIE_X_DATA, CIE_Y_DATA, CIE_Z_DATA};
use crate::{Arena, FileLoc, ParameterDictionary};

const N_SWATCH_REFLECTANCES: usize = 24;
const SWATCH_REFLECTANCES: LazyLock<[Spectrum; N_SWATCH_REFLECTANCES]> = LazyLock::new(|| {
    std::array::from_fn(|i| {
        let raw_data = SWATCHES_RAW[i];
        let pls = PiecewiseLinearSpectrum::from_interleaved(raw_data, false);
        Spectrum::Piecewise(leak(pls))
    })
});

pub fn get_swatches() -> Arc<[Spectrum; N_SWATCH_REFLECTANCES]> {
    Arc::new(*SWATCH_REFLECTANCES)
}

pub trait CreatePixelSensor: Sized {
    fn create(
        params: &ParameterDictionary,
        output_colorspace: Arc<RGBColorSpace>,
        exposure_time: Float,
        loc: &FileLoc,
        arena: &Arena,
    ) -> Result<Self>;

    fn new(
        r: &Spectrum,
        g: &Spectrum,
        b: &Spectrum,
        output_colorspace: Arc<RGBColorSpace>,
        sensor_illum: Option<&Spectrum>,
        imaging_ratio: Float,
        arena: &Arena,
    ) -> Self;

    fn new_with_white_balance(
        output_colorspace: &RGBColorSpace,
        sensor_illum: Option<&Spectrum>,
        imaging_ratio: Float,
        arena: &Arena,
    ) -> Self;
}

impl CreatePixelSensor for PixelSensor {
    fn create(
        params: &ParameterDictionary,
        output_colorspace: Arc<RGBColorSpace>,
        exposure_time: Float,
        loc: &FileLoc,
        arena: &Arena,
    ) -> Result<Self>
    where
        Self: Sized,
    {
        let iso = params.get_one_float("iso", 100.)?;
        let mut white_balance_temp = params.get_one_float("whitebalance", 0.)?;
        let sensor_name = params.get_one_string("sensor", "cie1931")?;
        if sensor_name != "cie1931" && white_balance_temp == 0. {
            white_balance_temp = 6500.;
        }
        let imaging_ratio = exposure_time * iso / 100.;

        let d_illum = if white_balance_temp == 0. {
            DenselySampledSpectrum::generate_cie_d(6500.)
        } else {
            DenselySampledSpectrum::generate_cie_d(white_balance_temp)
        };

        let d_ptr = arena.alloc(d_illum);

        let sensor_illum: Option<Arc<Spectrum>> = if white_balance_temp != 0. {
            Some(Spectrum::Dense(d_ptr).into())
        } else {
            None
        };

        if sensor_name == "cie1931" {
            Ok(Self::new_with_white_balance(
                output_colorspace.as_ref(),
                sensor_illum.as_deref(),
                imaging_ratio,
                arena
            ))
        } else {
            let r_opt = get_named_spectrum(&format!("{}_r", sensor_name));
            let g_opt = get_named_spectrum(&format!("{}_g", sensor_name));
            let b_opt = get_named_spectrum(&format!("{}_b", sensor_name));
            if r_opt.is_none() || g_opt.is_none() || b_opt.is_none() {
                return Err(anyhow!(
                    "{}: unknown sensor type '{}' (missing RGB spectral data)",
                    loc,
                    sensor_name
                ));
            }

            let r = r_opt.unwrap();
            let g = g_opt.unwrap();
            let b = b_opt.unwrap();

            Ok(Self::new(
                &r,
                &g,
                &b,
                output_colorspace.clone(),
                Some(
                    sensor_illum
                        .as_deref()
                        .expect("Sensor must have illuminant"),
                ),
                imaging_ratio,
                arena
            ))
        }
    }

    fn new(
        r: &Spectrum,
        g: &Spectrum,
        b: &Spectrum,
        output_colorspace: Arc<RGBColorSpace>,
        sensor_illum: Option<&Spectrum>,
        imaging_ratio: Float,
        arena: &Arena,
    ) -> Self {
        let illum: &Spectrum = match sensor_illum {
            Some(arc_illum) => arc_illum,
            None => &Spectrum::Dense(output_colorspace.as_ref().illuminant),
        };

        let r_bar = DenselySampledSpectrum::from_spectrum(r);
        let g_bar = DenselySampledSpectrum::from_spectrum(g);
        let b_bar = DenselySampledSpectrum::from_spectrum(b);
        let r_ptr = arena.alloc(r_bar);
        let g_ptr = arena.alloc(g_bar);
        let b_ptr = arena.alloc(b_bar);
        let mut rgb_camera = [[0.; 3]; N_SWATCH_REFLECTANCES];

        let swatches = get_swatches();

        for i in 0..N_SWATCH_REFLECTANCES {
            let rgb = PixelSensor::project_reflectance::<RGB>(
                &swatches[i],
                illum,
                &Spectrum::Dense(r_ptr),
                &Spectrum::Dense(g_ptr),
                &Spectrum::Dense(b_ptr),
            );
            for c in 0..3 {
                rgb_camera[i][c] = rgb[c];
            }
        }

        let mut xyz_output = [[0.; 3]; N_SWATCH_REFLECTANCES];
        let spectra = get_spectra_context();
        let sensor_white_g = illum.inner_product(&Spectrum::Dense(g_ptr));
        let sensor_white_y = illum.inner_product(&Spectrum::Dense(spectra.y));
        for i in 0..N_SWATCH_REFLECTANCES {
            let s = swatches[i];
            let xyz = PixelSensor::project_reflectance::<XYZ>(
                &s,
                illum,
                &Spectrum::Dense(spectra.x),
                &Spectrum::Dense(spectra.y),
                &Spectrum::Dense(spectra.z),
            ) * (sensor_white_y / sensor_white_g);
            for c in 0..3_u32 {
                xyz_output[i][c as usize] = xyz[c].try_into().unwrap();
            }
        }

        let xyz_from_sensor_rgb = linear_least_squares(rgb_camera, xyz_output)
            .expect("Could not convert sensor illuminance to XYZ space");

        PixelSensor {
            r_bar: r_ptr,
            g_bar: g_ptr,
            b_bar: b_ptr,
            imaging_ratio,
            xyz_from_sensor_rgb,
        }
    }

    fn new_with_white_balance(
        output_colorspace: &RGBColorSpace,
        sensor_illum: Option<&Spectrum>,
        imaging_ratio: Float,
        arena: &Arena
    ) -> Self {
        let spectra = get_spectra_context();
        let r_bar = CIE_X_DATA.clone();
        let g_bar = CIE_Y_DATA.clone();
        let b_bar = CIE_Z_DATA.clone();
        let xyz_from_sensor_rgb: SquareMatrix<Float, 3>;

        if let Some(illum) = sensor_illum {
            let source_white = illum.to_xyz(&spectra).xy();
            let target_white = output_colorspace.w;
            xyz_from_sensor_rgb = white_balance(source_white, target_white);
        } else {
            xyz_from_sensor_rgb = SquareMatrix::<Float, 3>::identity();
        }

        PixelSensor {
            r_bar: arena.alloc(r_bar),
            g_bar: arena.alloc(g_bar),
            b_bar: arena.alloc(b_bar),
            xyz_from_sensor_rgb,
            imaging_ratio,
        }
    }

}

pub trait CreateFilmBase {
    fn create(
        params: &ParameterDictionary,
        filter: Filter,
        sensor: Ptr<PixelSensor>,
        loc: &FileLoc,
    ) -> Result<Self>
    where
        Self: Sized;
}

impl CreateFilmBase for FilmBase {
    fn create(
        params: &ParameterDictionary,
        filter: Filter,
        sensor: Ptr<PixelSensor>,
        loc: &FileLoc,
    ) -> Result<Self>
    where
        Self: Sized,
    {
        let x_res = params.get_one_int("xresolution", 1280)?;
        let y_res = params.get_one_int("yresolution", 720)?;

        if x_res <= 0 || y_res <= 0 {
            eprintln!(
                "{}: Film resolution must be > 0. Defaulting to 1280x720.",
                loc
            );
        }
        let full_resolution = Point2i::new(x_res.max(1), y_res.max(1));

        let crop_data = params.get_float_array("cropwindow")?;
        let crop = if crop_data.len() == 4 {
            Bounds2f::from_points(
                Point2f::new(crop_data[0], crop_data[2]),
                Point2f::new(crop_data[1], crop_data[3]),
            )
        } else {
            Bounds2f::from_points(Point2f::zero(), Point2f::new(1.0, 1.0))
        };

        let p_min = Point2i::new(
            (full_resolution.x() as Float * crop.p_min.x()).ceil() as i32,
            (full_resolution.y() as Float * crop.p_min.y()).ceil() as i32,
        );
        let p_max = Point2i::new(
            (full_resolution.x() as Float * crop.p_max.x()).ceil() as i32,
            (full_resolution.y() as Float * crop.p_max.y()).ceil() as i32,
        );

        let mut pixel_bounds = Bounds2i::from_points(p_min, p_max);

        if pixel_bounds.is_empty() {
            eprintln!("{}: Film crop window results in empty pixel bounds.", loc);
        }

        let rad = filter.radius();
        let expansion = Point2i::new(rad.x().ceil() as i32, rad.y().ceil() as i32);
        pixel_bounds = pixel_bounds.expand(expansion);

        let diagonal_mm = params.get_one_float("diagonal", 35.0)?;
        // let filename = params.get_one_string("filename", "pbrt.exr");

        Ok(Self {
            full_resolution,
            pixel_bounds,
            filter,
            diagonal: diagonal_mm * 0.001,
            sensor,
        })
    }
}

pub trait FilmTrait: Sync {
    fn base(&self) -> &FilmBase;
    fn get_pixel_rgb(&self, p: Point2i, splat_scale: Option<Float>) -> RGB;
    fn write_image(&self, metadata: &ImageMetadata, splat_scale: Float, filename: &str) {
        let image = self.get_image(metadata, splat_scale);
        image.write(filename, metadata).expect("Something")
    }

    fn get_image(&self, _metadata: &ImageMetadata, splat_scale: Float) -> HostImage {
        let write_fp16 = true;
        let format = if write_fp16 {
            PixelFormat::F16
        } else {
            PixelFormat::F32
        };

        let channel_names = &["R", "G", "B"];

        let pixel_bounds = self.base().pixel_bounds;
        let resolution = Point2i::from(pixel_bounds.diagonal());

        let n_clamped = Arc::new(AtomicUsize::new(0));
        let processed_rows: Vec<Vec<Float>> = (pixel_bounds.p_min.y()..pixel_bounds.p_max.y())
            .into_par_iter()
            .map(|y| {
                let n_clamped = Arc::clone(&n_clamped);
                let mut row_data = Vec::with_capacity(resolution.x() as usize * 3);
                for x in pixel_bounds.p_min.x()..pixel_bounds.p_max.x() {
                    let p = Point2i::new(x, y);
                    let mut rgb = self.get_pixel_rgb(p, Some(splat_scale));
                    let mut was_clamped = false;
                    if write_fp16 {
                        if rgb.r > 65504.0 {
                            rgb.r = 65504.0;
                            was_clamped = true;
                        }
                        if rgb.g > 65504.0 {
                            rgb.g = 65504.0;
                            was_clamped = true;
                        }
                        if rgb.b > 65504.0 {
                            rgb.b = 65504.0;
                            was_clamped = true;
                        }
                    }
                    if was_clamped {
                        n_clamped.fetch_add(1, Ordering::SeqCst);
                    }
                    row_data.push(rgb.r);
                    row_data.push(rgb.g);
                    row_data.push(rgb.b);
                }
                row_data
            })
            .collect();

        let mut image = HostImage::new(format, resolution, channel_names, SRGB);
        let _rgb_desc = ImageChannelDesc::new(&[0, 1, 2]);

        for (iy, row_data) in processed_rows.into_iter().enumerate() {
            for (ix, rgb_chunk) in row_data.chunks_exact(3).enumerate() {
                let p_offset = Point2i::new(ix as i32, iy as i32);
                let values = ImageChannelValues::from(rgb_chunk);
                image.set_channels(p_offset, &values);
            }
        }

        let clamped_count = n_clamped.load(Ordering::SeqCst);
        if clamped_count > 0 {
            println!(
                "{} pixel values clamped to maximum fp16 value.",
                clamped_count
            );
        }

        image
    }
}

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

    fn get_pixel_rgb(&self, p: Point2i, splat_scale: Option<Float>) -> RGB {
        match self {
            Film::RGB(f) => f.get_pixel_rgb(p, splat_scale),
            Film::GBuffer(f) => f.get_pixel_rgb(p, splat_scale),
            Film::Spectral(f) => f.get_pixel_rgb(p, splat_scale),
        }
    }
}

pub trait FilmFactory {
    fn create(
        name: &str,
        params: &ParameterDictionary,
        exposure_time: Float,
        filter: Filter,
        _camera_transform: Option<CameraTransform>,
        loc: &FileLoc,
        arena: &Arena,
    ) -> Result<Self>
    where
        Self: Sized;
}

impl FilmFactory for Film {
    fn create(
        name: &str,
        params: &ParameterDictionary,
        exposure_time: Float,
        filter: Filter,
        camera_transform: Option<CameraTransform>,
        loc: &FileLoc,
        arena: &Arena,
    ) -> Result<Self>
    where
        Self: Sized,
    {
        match name {
            "gbuffer" => {
                GBufferFilm::create(params, exposure_time, filter, camera_transform, loc, arena)
            }
            "rgb" => RGBFilm::create(params, exposure_time, filter, camera_transform, loc, arena),
            "spectral" => {
                SpectralFilm::create(params, exposure_time, filter, camera_transform, loc, arena)
            }
            _ => Err(anyhow!("Film type '{}' unknown at {}", name, loc)),
        }
    }
}