use crate::Arena;
use crate::core::image::{Image, ImageIO};
use crate::core::light::lookup_spectrum;
use crate::core::spectrum::spectrum_to_photometric;
use crate::spectra::get_spectra_context;
use crate::utils::sampling::{PiecewiseConstant2D, WindowedPiecewiseConstant2D};
use crate::utils::{FileLoc, ParameterDictionary, Upload, resolve_filename};
use anyhow::{Result, anyhow};
use rayon::prelude::*;
use shared::core::camera::CameraTransform;
use shared::core::geometry::{Bounds2f, Frame, Point2f, Point2i, Point3f, VectorLike, cos_theta};
use shared::core::image::{DeviceImage, 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::{DenselySampledSpectrum, RGBColorSpace};
use shared::utils::math::{equal_area_sphere_to_square, equal_area_square_to_sphere};
use shared::utils::sampling::{DevicePiecewiseConstant2D, DeviceWindowedPiecewiseConstant2D};
use shared::utils::{Ptr, Transform};
use shared::{Float, PI};
use std::path::Path;

pub trait CreateImageInfiniteLight {
    fn new(
        render_from_light: Transform,
        scale: Float,
        image: Ptr<DeviceImage>,
        image_color_space: Ptr<RGBColorSpace>,
        distrib: Ptr<DevicePiecewiseConstant2D>,
        compensated_distrib: Ptr<DevicePiecewiseConstant2D>,
    ) -> Self;
}

impl CreateImageInfiniteLight for ImageInfiniteLight {
    fn new(
        render_from_light: Transform,
        scale: Float,
        image: Ptr<DeviceImage>,
        image_color_space: Ptr<RGBColorSpace>,
        distrib: Ptr<DevicePiecewiseConstant2D>,
        compensated_distrib: Ptr<DevicePiecewiseConstant2D>,
    ) -> Self {
        let base = LightBase::new(
            LightType::Infinite,
            render_from_light,
            MediumInterface::default(),
        );
        Self {
            base,
            image,
            image_color_space,
            scale,
            distrib,
            compensated_distrib,
            scene_center: Point3f::default(),
            scene_radius: 0.0,
        }
    }
}

pub trait CreatePortalInfiniteLight {
    fn new(
        render_from_light: Transform,
        scale: Float,
        image: Ptr<DeviceImage>,
        image_color_space: Ptr<RGBColorSpace>,
        portal: [Point3f; 4],
        portal_frame: Frame,
        distribution: Ptr<DeviceWindowedPiecewiseConstant2D>,
    ) -> Self;
}

impl CreatePortalInfiniteLight for PortalInfiniteLight {
    fn new(
        render_from_light: Transform,
        scale: Float,
        image: Ptr<DeviceImage>,
        image_color_space: Ptr<RGBColorSpace>,
        portal: [Point3f; 4],
        portal_frame: Frame,
        distribution: Ptr<DeviceWindowedPiecewiseConstant2D>,
    ) -> Self {
        let base = LightBase::new(
            LightType::Infinite,
            render_from_light,
            MediumInterface::default(),
        );
        Self {
            base,
            image,
            image_color_space,
            scale,
            portal,
            portal_frame,
            distribution: *distribution,
            scene_center: Point3f::default(),
            scene_radius: 0.0,
        }
    }
}

pub trait CreateUniformInfiniteLight {
    fn new(render_from_light: Transform, scale: Float, lemit: Ptr<DenselySampledSpectrum>) -> Self;
}

impl CreateUniformInfiniteLight for UniformInfiniteLight {
    fn new(render_from_light: Transform, scale: Float, lemit: Ptr<DenselySampledSpectrum>) -> Self {
        let base = LightBase::new(
            LightType::Infinite,
            render_from_light,
            MediumInterface::default(),
        );
        Self {
            base,
            lemit,
            scale,
            scene_center: Point3f::default(),
            scene_radius: 0.0,
        }
    }
}

pub fn create(
    arena: &mut Arena,
    render_from_light: Transform,
    _medium: MediumInterface,
    camera_transform: CameraTransform,
    parameters: &ParameterDictionary,
    colorspace: Option<&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\""));
    }

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

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

        let lemit = lookup_spectrum(&spectrum);
        let light = UniformInfiniteLight::new(render_from_light, scale, lemit.upload(arena));
        return Ok(Light::InfiniteUniform(light));
    }

    // Image-based lights
    let (image, image_cs) = load_image(&filename, &l, colorspace.unwrap(), 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;
    }

    if has_portal {
        create_portal_light(
            arena,
            render_from_light,
            scale,
            image,
            image_cs,
            &portal,
            camera_transform,
            loc,
        )
    } else {
        create_image_light(arena, render_from_light, scale, image, image_cs)
    }
}

fn create_image_light(
    arena: &mut Arena,
    render_from_light: Transform,
    scale: Float,
    image: Image,
    image_cs: RGBColorSpace,
) -> Result<Light> {
    let res = image.resolution();
    assert_eq!(
        res.x(),
        res.y(),
        "Image must be square for equal-area mapping"
    );

    let (n_u, n_v) = (res.x() as usize, res.y() as usize);

    // Extract luminance data
    let image_ptr = image.upload(arena);
    let value = &image;
    let mut data: Vec<Float> = (0..n_v)
        .flat_map(|v| {
            (0..n_u).map(move |u| {
                value
                    .get_channels(Point2i::new(u as i32, v as i32))
                    .average()
            })
        })
        .collect();

    let distrib = PiecewiseConstant2D::from_slice(&data_u, n_u, n_v, Bounds2f::unit());

    // Build compensated distribution
    let average = data.iter().sum::<Float>() / data.len() as Float;
    let mut all_zero = true;
    for v in &mut data {
        *v = (*v - average).max(0.0);
        all_zero &= *v == 0.0;
    }
    if all_zero {
        data.fill(1.0);
    }
    let compensated_distrib = PiecewiseConstant2D::from_slice(&data, n_u, n_v, Bounds2f::unit());

    let light = ImageInfiniteLight::new(
        render_from_light,
        scale,
        image_ptr,
        image_cs.upload(arena),
        distrib.upload(arena),
        compensated_distrib.upload(arena),
    );

    Ok(Light::InfiniteImage(light))
}

fn create_portal_light(
    arena: &mut Arena,
    render_from_light: Transform,
    scale: Float,
    image: Image,
    image_cs: RGBColorSpace,
    portal_points: &[Point3f],
    camera_transform: CameraTransform,
    loc: &FileLoc,
) -> Result<Light> {
    let res = image.resolution();
    if res.x() != res.y() {
        return Err(anyhow!(loc, "Portal light image must be square"));
    }

    // Validate portal
    if portal_points.len() != 4 {
        return Err(anyhow!(
            loc,
            "Portal requires exactly 4 vertices, got {}",
            portal_points.len()
        ));
    }

    let portal: [Point3f; 4] = portal_points
        .iter()
        .map(|p| camera_transform.camera_from_world(0.0).apply_to_point(*p))
        .collect::<Vec<_>>()
        .try_into()
        .unwrap();

    let portal_frame = validate_and_build_portal_frame(&portal, loc)?;

    // Remap image through portal
    let remapped = remap_image_through_portal(&image, &render_from_light, &portal_frame);

    // Build distribution
    let duv_dw = |p: Point2f| -> Float {
        let (_, jacobian) = PortalInfiniteLight::render_from_image(portal_frame, p);
        jacobian
    };
    let d = remapped.get_sampling_distribution(
        duv_dw,
        Bounds2f::from_points(Point2f::zero(), Point2f::fill(1.)),
    );
    let distribution = WindowedPiecewiseConstant2D::new(d);

    let light = PortalInfiniteLight::new(
        render_from_light,
        scale,
        remapped.upload(arena),
        image_cs.upload(arena),
        portal,
        portal_frame,
        distribution.upload(arena),
    );

    Ok(Light::InfinitePortal(light))
}

fn validate_and_build_portal_frame(portal: &[Point3f; 4], loc: &FileLoc) -> Result<Frame> {
    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 {
        return Err(anyhow!(loc, "Portal edges not parallel"));
    }

    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
    {
        return Err(anyhow!(loc, "Portal edges not perpendicular"));
    }

    Ok(Frame::from_xy(p03, p01))
}

fn remap_image_through_portal(
    image: &Image,
    render_from_light: &Transform,
    portal_frame: &Frame,
) -> Image {
    let res = image.resolution();
    let (width, height) = (res.x() as usize, res.y() as usize);

    let mut pixels = vec![0.0f32; width * height * 3];

    pixels
        .par_chunks_mut(width * 3)
        .enumerate()
        .for_each(|(y, row)| {
            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, _) = PortalInfiniteLight::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);

                for c in 0..3 {
                    row[x * 3 + c] = image.bilerp_channel_with_wrap(
                        uv_equi,
                        c as i32,
                        WrapMode::OctahedralSphere.into(),
                    );
                }
            }
        });

    Image::from_f32(pixels, res, &["R", "G", "B"])
}

fn load_image(
    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 image =
            Image::new_constant(Point2i::new(1, 1), &["R", "G", "B"], &[rgb.r, rgb.g, rgb.b]);
        return Ok((image, colorspace.clone()));
    }

    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));
    }

    let desc = 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());
    Ok((im.image.select_channels(&desc), cs))
}

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

    let sum: Float = (0..res.y())
        .flat_map(|y| (0..res.x()).map(move |x| (x, y)))
        .filter_map(|(x, y)| {
            let u = (x as Float + 0.5) / res.x() as Float;
            let v = (y as Float + 0.5) / res.y() as Float;
            let w = equal_area_square_to_sphere(Point2f::new(u, v));

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

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

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

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