use crate::Arena;
use crate::core::image::Image;
use crate::core::spectrum::spectrum_to_photometric;
use crate::utils::{FileLoc, ParameterDictionary, resolve_filename};
use log::error;
use shared::core::camera::CameraTransform;
use shared::core::geometry::{Bounds2f, Frame, Point2f, Point2i, Point3f, 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::sampling::{DevicePiecewiseConstant2D, WindowedPiecewiseConstant2D};
use shared::utils::{Ptr, Transform};
use shared::{Float, PI};
use std::fmt::Error;
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: Ptr<Image>,
        image_color_space: Ptr<RGBColorSpace>,
    ) -> Self;
}

impl CreateImageInfiniteLight for ImageInfiniteLight {
    fn new(
        render_from_light: Transform,
        medium_interface: MediumInterface,
        scale: Float,
        image: Ptr<Image>,
        image_color_space: Ptr<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());
        if image.resolution().x() != image.resolution().y() {
            let into = hash_float(hashee);
            panic!(
                "Image resolution ({}, {}) is non-square. It's unlikely this is an equal area environment map.",
                image.resolution.x(),
                image.resolution.y()
            );
        }
        let mut d = image.get_sampling_distribution_uniform();
        let domain = Bounds2f::from_points(Point2f::new(0., 0.), Point2f::new(1., 1.));
        let distrib = DevicePiecewiseConstant2D::new_with_bounds(&d, domain);
        let slice = &mut d.values; // or d.as_slice_mut()
        let count = slice.len() as Float;
        let sum: Float = slice.iter().sum();
        let average = sum / count;

        for v in slice.iter_mut() {
            *v = (*v - average).max(0.0);
        }

        let all_zero = slice.iter().all(|&v| v == 0.0);
        if all_zero {
            for v in slice.iter_mut() {
                *v = 1.0;
            }
        }

        let compensated_distrib = DevicePiecewiseConstant2D::new_with_bounds(&d, domain);

        ImageInfiniteLight {
            base,
            image: &image,
            image_color_space,
            scene_center: Point3f::default(),
            scene_radius: 0.,
            scale,
            distrib: &distrib,
            compensated_distrib: &compensated_distrib,
        };
    }
}

#[derive(Debug)]
struct InfinitePortalLightStorage {
    image: Image,
    distribution: WindowedPiecewiseConstant2D,
    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: Ptr<Image>,
        image_color_space: Ptr<RGBColorSpace>,
        points: Ptr<Point3f>,
    ) -> Self;
}

impl CreatePortalInfiniteLight for PortalInfiniteLight {
    fn new(
        render_from_light: Transform,
        scale: Float,
        image: Ptr<Image>,
        image_color_space: Ptr<RGBColorSpace>,
        points: Ptr<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] = val;
                    }
                }
            });

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

        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 = WindowedPiecewiseConstant2D::new(d);

        PortalInfiniteLight {
            base,
            image: img,
            image_color_space: &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 = 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, Error> {
    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(error!(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 {
            &colorspace.illuminant
        };

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

        let light = UniformInfiniteLight::new(render_from_light, lookup_spectrum(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(&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.render_from_world(*p))
            .collect();

        let light = PortalInfiniteLight::new(
            render_from_light,
            scale,
            image_ptr,
            cs_ptr,
            arena.alloc_slice(&portal_render),
        );
        Ok(Light::InfinitePortal(light))
    } else {
        let light = ImageInfiniteLight::new(render_from_light, medium, scale, image_ptr, cs_ptr);
        Ok(Light::InfiniteImage(light))
    }
}

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

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

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

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

        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
}
