use crate::spectra::dense::DenselySampledSpectrumBuffer;
use shared::core::geometry::{Bounds3f, Point3i};
use shared::core::medium::{GridMedium, HGPhaseFunction, HomogeneousMedium, RGBGridMedium};
use shared::core::spectrum::Spectrum;
use shared::spectra::{DenselySampledSpectrum, RGBIlluminantSpectrum, RGBUnboundedSpectrum};
use shared::utils::Transform;
use shared::utils::containers::SampledGrid;
use shared::{Float, core::medium::MajorantGrid};

pub struct MajorantGridHost {
    pub device: MajorantGrid,
    voxels: Vec<Float>,
}

impl MajorantGridHost {
    pub fn new(bounds: Bounds3f, res: Point3i) -> Self {
        let n = (res.x() * res.y() * res.z()) as usize;
        let voxels = vec![0.0; n];

        let device = MajorantGrid {
            bounds,
            res,
            voxels: std::ptr::null_mut(),
            n_voxels: n as u32,
        };

        Self { device, voxels }
    }
}

pub trait RGBGridMediumCreator {
    fn new(
        bounds: &Bounds3f,
        render_from_medium: &Transform,
        g: Float,
        sigma_a_grid: SampledGrid<RGBUnboundedSpectrum>,
        sigma_s_grid: SampledGrid<RGBUnboundedSpectrum>,
        sigma_scale: Float,
        le_grid: SampledGrid<RGBIlluminantSpectrum>,
        le_scale: Float,
    ) -> Self;
}

impl RGBGridMediumCreator for RGBGridMedium {
    #[allow(clippy::too_many_arguments)]
    fn new(
        bounds: &Bounds3f,
        render_from_medium: &Transform,
        g: Float,
        sigma_a_grid: SampledGrid<RGBUnboundedSpectrum>,
        sigma_s_grid: SampledGrid<RGBUnboundedSpectrum>,
        sigma_scale: Float,
        le_grid: SampledGrid<RGBIlluminantSpectrum>,
        le_scale: Float,
    ) -> Self {
        let majorant_grid = MajorantGrid::new(*bounds, Point3i::new(16, 16, 16));
        for z in 0..majorant_grid.res.x() {
            for y in 0..majorant_grid.res.y() {
                for x in 0..majorant_grid.res.x() {
                    let bounds = majorant_grid.voxel_bounds(x, y, z);
                    let convert = |s: &RGBUnboundedSpectrum| s.max_value();
                    let max_sigma_t = sigma_a_grid.max_value_convert(bounds, convert)
                        + sigma_s_grid.max_value_convert(bounds, convert);
                    majorant_grid.set(x, y, z, sigma_scale * max_sigma_t);
                }
            }
        }

        Self {
            bounds: *bounds,
            render_from_medium: *render_from_medium,
            le_grid,
            le_scale,
            phase: HGPhaseFunction::new(g),
            sigma_a_grid,
            sigma_s_grid,
            sigma_scale,
            majorant_grid,
        }
    }
}

pub trait GridMediumCreator {
    fn new(
        bounds: &Bounds3f,
        render_from_medium: &Transform,
        sigma_a: &Spectrum,
        sigma_s: &Spectrum,
        sigma_scale: Float,
        g: Float,
        density_grid: SampledGrid<Float>,
        temperature_grid: SampledGrid<Float>,
        le: &Spectrum,
        le_scale: SampledGrid<Float>,
    ) -> Self;
}

impl GridMediumCreator for GridMedium {
    #[allow(clippy::too_many_arguments)]
    fn new(
        bounds: &Bounds3f,
        render_from_medium: &Transform,
        sigma_a: &Spectrum,
        sigma_s: &Spectrum,
        sigma_scale: Float,
        g: Float,
        density_grid: SampledGrid<Float>,
        temperature_grid: SampledGrid<Float>,
        le: &Spectrum,
        le_scale: SampledGrid<Float>,
    ) -> Self {
        let mut sigma_a_spec = DenselySampledSpectrum::from_spectrum(sigma_a);
        let mut sigma_s_spec = DenselySampledSpectrum::from_spectrum(sigma_s);
        let le_spec = DenselySampledSpectrum::from_spectrum(le);
        sigma_a_spec.scale(sigma_scale);
        sigma_s_spec.scale(sigma_scale);

        let mut majorant_grid = MajorantGrid::new(*bounds, Point3i::new(16, 16, 16));
        let is_emissive = if temperature_grid.is_some() {
            true
        } else {
            le_spec.max_value() > 0.
        };

        for z in 0..majorant_grid.res.z() {
            for y in 0..majorant_grid.res.y() {
                for x in 0..majorant_grid.res.x() {
                    let bounds = majorant_grid.voxel_bounds(x, y, z);
                    majorant_grid.set(x, y, z, density_grid.max_value(bounds));
                }
            }
        }

        Self {
            bounds: *bounds,
            render_from_medium: *render_from_medium,
            sigma_a_spec,
            sigma_s_spec,
            density_grid,
            phase: HGPhaseFunction::new(g),
            temperature_grid,
            le_spec,
            le_scale,
            is_emissive,
            majorant_grid,
        }
    }
}

pub trait HomogeneousMediumCreator {
    fn new(
        sigma_a: Spectrum,
        sigma_s: Spectrum,
        sigma_scale: Float,
        le: Spectrum,
        le_scale: Float,
        g: Float,
    ) -> Self;
}

impl HomogeneousMediumCreator for HomogeneousMedium {
    fn new(
        sigma_a: Spectrum,
        sigma_s: Spectrum,
        sigma_scale: Float,
        le: Spectrum,
        le_scale: Float,
        g: Float,
    ) -> Self {
        let mut sigma_a_spec = DenselySampledSpectrumBuffer::from_spectrum(&sigma_a);
        let mut sigma_s_spec = DenselySampledSpectrumBuffer::from_spectrum(&sigma_s);
        let mut le_spec = DenselySampledSpectrum::from_spectrum(&le);

        sigma_a_spec.scale(sigma_scale);
        sigma_s_spec.scale(sigma_scale);
        le_spec.scale(le_scale);

        Self {
            sigma_a_spec,
            sigma_s_spec,
            le_spec,
            phase: HGPhaseFunction::new(g),
        }
    }
}