use crate::Arena;
use crate::core::image::{Image, ImageIO};
use crate::core::spectrum::spectrum_to_photometric;
use crate::spectra::{get_colorspace_context, get_spectra_context};
use crate::utils::sampling::{PiecewiseConstant2D, WindowedPiecewiseConstant2D};
use crate::utils::{FileLoc, ParameterDictionary, Upload, resolve_filename};
use anyhow::{Result, anyhow};

use rayon::iter::{IndexedParallelIterator, ParallelIterator};
use rayon::prelude::ParallelSliceMut;
use shared::core::camera::CameraTransform;
use shared::core::geometry::{Bounds2f, Frame, Point2f, Point2i, Point3f, VectorLike, cos_theta};
use shared::core::image::{PixelFormat, WrapMode};
use shared::core::light::{Light, LightBase, LightType};
use shared::core::medium::MediumInterface;
use shared::core::spectrum::Spectrum;
use shared::core::texture::SpectrumType;
use shared::lights::{ImageInfiniteLight, PortalInfiniteLight, UniformInfiniteLight};
use shared::spectra::RGBColorSpace;
use shared::utils::hash::hash_float;
use shared::utils::math::{equal_area_sphere_to_square, equal_area_square_to_sphere};
use shared::utils::{Ptr, Transform};
use shared::{Float, PI};
use std::path::Path;
use std::sync::Arc;

use crate::core::light::lookup_spectrum;

pub trait CreateImageInfiniteLight {
    fn new(
        render_from_light: Transform,
        medium_interface: MediumInterface,
        scale: Float,
        image: Arc<DeviceImage>,
        image_color_space: Arc<RGBColorSpace>,
    ) -> Self;
}

impl CreateImageInfiniteLight for ImageInfiniteLight {
    fn new(
        render_from_light: Transform,
        medium_interface: MediumInterface,
        scale: Float,
        image: Arc<Image>,
        image_color_space: Arc<RGBColorSpace>,
    ) -> Self {
        let base = LightBase::new(
            LightType::Infinite,
            render_from_light,
            MediumInterface::default(),
        );

        let desc = image
            .get_channel_desc(&["R", "G", "B"])
            .expect("Image used for DiffuseAreaLight doesn't have R, G, B channels");

        assert_eq!(3, desc.size());
        assert!(desc.is_identity());

        let res = image.resolution();
        assert_eq!(
            res.x(),
            res.y(),
            "Image resolution ({}, {}) is non-square. Unlikely to be an equal area environment map.",
            res.x(),
            res.y()
        );

        let n_u = res.x() as usize;
        let n_v = res.y() as usize;
        let mut data: Vec<Float> = (0..n_v)
            .flat_map(|v| {
                (0..n_u).map(move |u| {
                    image
                        .get_channels(Point2i::new(u as i32, v as i32))
                        .average()
                })
            })
            .collect();

        let distrib = PiecewiseConstant2D::new(&data, n_u, n_v);

        let slice = d.as_mut_slice();
        let average = slice.iter().sum::<Float>() / slice.len() as Float;

        let mut all_zero = true;
        for v in slice.iter_mut() {
            *v = (*v - average).max(0.0);
            all_zero &= *v == 0.0;
        }

        if all_zero {
            data.fill(1.0);
        }

        let compensated_distrib = PiecewiseConstant2D::new(&data, n_u, n_v);

        ImageInfiniteLight {
            base,
            image: Ptr::from(image.device_image()),
            image_color_space: Ptr::from(image_color_space.as_ref()),
            scene_center: Point3f::default(),
            scene_radius: 0.,
            scale,
            distrib: Ptr::from(&*distrib),
            compensated_distrib: Ptr::from(&*compensated_distrib),
        }
    }
}

#[derive(Debug)]
struct InfinitePortalLightStorage {
    image: Image,
    distribution: DeviceWindowedPiecewiseConstant2D,
    image_color_space: RGBColorSpace,
}

#[derive(Clone, Debug)]
pub struct PortalInfiniteLightHost {
    pub view: PortalInfiniteLight,
    pub filename: String,
    _storage: Arc<InfinitePortalLightStorage>,
}

pub trait CreatePortalInfiniteLight {
    fn new(
        render_from_light: Transform,
        scale: Float,
        image: Arc<Image>,
        image_color_space: Arc<RGBColorSpace>,
        points: Vec<Point3f>,
    ) -> Self;
}

impl CreatePortalInfiniteLight for PortalInfiniteLight {
    fn new(
        render_from_light: Transform,
        scale: Float,
        image: Arc<Image>,
        image_color_space: Arc<RGBColorSpace>,
        points: Vec<Point3f>,
    ) -> Self {
        let base = LightBase::new(
            LightType::Infinite,
            render_from_light,
            MediumInterface::default(),
        );

        let desc = image
            .get_channel_desc(&["R", "G", "B"])
            .unwrap_or_else(|_| {
                panic!("Image used for PortalImageInfiniteLight doesn't have R, G, B channels.",)
            });

        assert_eq!(3, desc.offset.len());
        let src_res = image.resolution();
        if src_res.x() != src_res.y() {
            panic!(
                "Image resolution ({}, {}) is non-square. It's unlikely this is an equal area environment map.",
                src_res.x(),
                src_res.y()
            );
        }

        if points.len() != 4 {
            panic!(
                "Expected 4 vertices for infinite light portal but given {}",
                points.len()
            );
        }

        let portal: [Point3f; 4] = [points[0], points[1], points[2], points[3]];

        let p01 = (portal[1] - portal[0]).normalize();
        let p12 = (portal[2] - portal[1]).normalize();
        let p32 = (portal[2] - portal[3]).normalize();
        let p03 = (portal[3] - portal[0]).normalize();

        if (p01.dot(p32) - 1.0).abs() > 0.001 || (p12.dot(p03) - 1.0).abs() > 0.001 {
            panic!("Infinite light portal isn't a planar quadrilateral (opposite edges)");
        }

        if p01.dot(p12).abs() > 0.001
            || p12.dot(p32).abs() > 0.001
            || p32.dot(p03).abs() > 0.001
            || p03.dot(p01).abs() > 0.001
        {
            panic!("Infinite light portal isn't a planar quadrilateral (perpendicular edges)");
        }

        let portal_frame = Frame::from_xy(p03, p01);

        let width = src_res.x();
        let height = src_res.y();

        let mut new_pixels = vec![0.0 as Float; (width * height * 3) as usize];

        new_pixels
            .par_chunks_mut((width * 3) as usize)
            .enumerate()
            .for_each(|(y, row_pixels)| {
                let y = y as i32;

                for x in 0..width {
                    let uv = Point2f::new(
                        (x as Float + 0.5) / width as Float,
                        (y as Float + 0.5) / height as Float,
                    );

                    let (w_world, _) = Self::render_from_image(portal_frame, uv);
                    let w_local = render_from_light.apply_inverse_vector(w_world).normalize();
                    let uv_equi = equal_area_sphere_to_square(w_local);

                    let pixel_idx = (x * 3) as usize;

                    for c in 0..3 {
                        let val = image.bilerp_channel_with_wrap(
                            uv_equi,
                            c,
                            WrapMode::OctahedralSphere.into(),
                        );
                        row_pixels[pixel_idx + c as usize] = val;
                    }
                }
            });

        let img = Image::new(
            PixelFormat::F32,
            src_res,
            &["R", "G", "B"],
            image.encoding().into(),
        );

        let duv_dw_closure = |p: Point2f| -> Float {
            let (_, jacobian) = Self::render_from_image(portal_frame, p);
            jacobian
        };

        let d = img.get_sampling_distribution(
            duv_dw_closure,
            Bounds2f::from_points(Point2f::new(0., 0.), Point2f::new(1., 1.)),
        );

        let distribution = DeviceWindowedPiecewiseConstant2D::new(d);

        PortalInfiniteLight {
            base,
            image: Ptr::from(img.device_image()),
            image_color_space: Ptr::from(&*image_color_space),
            scale,
            scene_center: Point3f::default(),
            scene_radius: 0.,
            portal,
            portal_frame,
            distribution,
        }
    }
}

pub trait CreateUniformInfiniteLight {
    fn new(render_from_light: Transform, le: Spectrum, scale: Float) -> Self;
}

impl CreateUniformInfiniteLight for UniformInfiniteLight {
    fn new(render_from_light: Transform, le: Spectrum, scale: Float) -> Self {
        let base = LightBase::new(
            LightType::Infinite,
            render_from_light,
            MediumInterface::default(),
        );
        let lemit = Ptr::from(&lookup_spectrum(&le));
        Self {
            base,
            lemit,
            scale,
            scene_center: Point3f::default(),
            scene_radius: 0.,
        }
    }
}

pub fn create(
    arena: &mut Arena,
    render_from_light: Transform,
    medium: MediumInterface,
    camera_transform: CameraTransform,
    parameters: &ParameterDictionary,
    colorspace: &RGBColorSpace,
    loc: &FileLoc,
) -> Result<Light> {
    let l = parameters.get_spectrum_array("L", SpectrumType::Illuminant);
    let mut scale = parameters.get_one_float("scale", 1.0);
    let portal = parameters.get_point3f_array("portal");
    let filename = resolve_filename(&parameters.get_one_string("filename", ""));
    let e_v = parameters.get_one_float("illuminance", -1.0);

    let has_spectrum = !l.is_empty();
    let has_file = !filename.is_empty();
    let has_portal = !portal.is_empty();

    if has_spectrum && has_file {
        return Err(anyhow!(loc, "cannot specify both \"L\" and \"filename\""));
    }

    if !has_file && !has_portal {
        let spectrum = if has_spectrum {
            scale /= spectrum_to_photometric(l[0]);
            l[0]
        } else {
            Spectrum::Dense(colorspace.illuminant)
        };

        if e_v > 0.0 {
            scale *= e_v / PI;
        }

        let light = UniformInfiniteLight::new(render_from_light, spectrum, scale);
        return Ok(Light::InfiniteUniform(light));
    }

    // Image based

    let (image, image_cs) = load_image_or_constant(&filename, &l, colorspace, loc)?;

    scale /= spectrum_to_photometric(Spectrum::Dense(image_cs.illuminant));

    if e_v > 0.0 {
        let k_e = compute_hemisphere_illuminance(&image, &image_cs);
        scale *= e_v / k_e;
    }

    // let image_ptr = image.upload(arena);
    // let cs_ptr = image_cs.upload(arena);

    if has_portal {
        let portal_render: Vec<Point3f> = portal
            .iter()
            .map(|p| camera_transform.camera_from_world(0.0).apply_to_point(*p))
            .collect();

        let (portal_ptr, portal_len) = arena.alloc_slice(&portal_render);
        let light =
            PortalInfiniteLight::new(render_from_light, scale, image.into(), cs, portal_render);
        Ok(Light::InfinitePortal(light))
    } else {
        let light = ImageInfiniteLight::new(render_from_light, medium, scale, image, cs);
        Ok(Light::InfiniteImage(light))
    }
}

fn load_image_or_constant(
    filename: &str,
    l: &[Spectrum],
    colorspace: &RGBColorSpace,
    loc: &FileLoc,
) -> Result<(Image, RGBColorSpace)> {
    if filename.is_empty() {
        let stdspec = get_spectra_context();
        let rgb = &l[0].to_rgb(colorspace, &stdspec);
        let rgb_values = [rgb.r, rgb.g, rgb.b];
        let image = Image::new_constant(Point2i::new(1, 1), &["R", "G", "B"], &rgb_values);
        Ok((image, colorspace.clone()))
    } else {
        let im = Image::read(Path::new(&filename), None)
            .map_err(|e| anyhow!(loc, "failed to load '{}': {}", filename, e))?;

        if im.image.has_any_infinite_pixels() || im.image.has_any_nan_pixels() {
            return Err(anyhow!(loc, "image '{}' has invalid pixels", filename));
        }

        im.image
            .get_channel_desc(&["R", "G", "B"])
            .map_err(|_| anyhow!(loc, "image '{}' must have R, G, B channels", filename))?;

        let cs = im.metadata.colorspace.unwrap_or_else(|| colorspace.clone());
        let image_desc = im.image.get_channel_desc(&["R", "G", "B"])?;
        let selected = im.image.select_channels(&image_desc);

        Ok((selected, cs))
    }
}

fn compute_hemisphere_illuminance(image: &Image, cs: &RGBColorSpace) -> Float {
    let lum = cs.luminance_vector();
    let res = image.resolution();
    let mut sum = 0.0;

    for y in 0..res.y() {
        let v = (y as Float + 0.5) / res.y() as Float;
        for x in 0..res.x() {
            let u = (x as Float + 0.5) / res.x() as Float;
            let w = equal_area_square_to_sphere(Point2f::new(u, v));

            if w.z() <= 0.0 {
                continue;
            }

            let r = image.get_channel(Point2i::new(x, y), 0);
            let g = image.get_channel(Point2i::new(x, y), 1);
            let b = image.get_channel(Point2i::new(x, y), 2);

            sum += (r * lum[0] + g * lum[1] + b * lum[2]) * cos_theta(w);
        }
    }

    sum * 2.0 * PI / (res.x() * res.y()) as Float
}