use crate::PI;
use crate::core::camera::{CameraBase, CameraRay, CameraTrait, CameraTransform};
use crate::core::color::SRGB;
use crate::core::film::Film;
use crate::core::geometry::{
    Bounds2f, Normal3f, Point2f, Point2i, Point3f, Ray, Vector2f, Vector2i, Vector3f, VectorLike,
};
use crate::core::image::{DeviceImage, PixelFormat};
use crate::core::medium::Medium;
use crate::core::pbrt::Float;
use crate::core::sampler::CameraSample;
use crate::core::scattering::refract;
use crate::spectra::{SampledSpectrum, SampledWavelengths};
use crate::utils::Ptr;
use crate::utils::math::{lerp, quadratic, square};

#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct LensElementInterface {
    pub curvature_radius: Float,
    pub thickness: Float,
    pub eta: Float,
    pub aperture_radius: Float,
}

#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct ExitPupilSample {
    pub p_pupil: Point3f,
    pub pdf: Float,
}

const EXIT_PUPIL_SAMPLES: usize = 64;

#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct RealisticCamera {
    base: CameraBase,
    focus_distance: Float,
    set_aperture_diameter: Float,
    aperture_image: Ptr<DeviceImage>,
    element_interfaces: Ptr<LensElementInterface>,
    n_elements: usize,
    physical_extent: Bounds2f,
    exit_pupil_bounds: [Bounds2f; EXIT_PUPIL_SAMPLES],
}

#[cfg(not(target_os = "cuda"))]
impl RealisticCamera {
    pub fn new(
        base: CameraBase,
        lens_params: &[Float],
        focus_distance: Float,
        set_aperture_diameter: Float,
        aperture_image: Ptr<DeviceImage>,
    ) -> Self {
        let film_ptr = base.film;
        if film_ptr.is_null() {
            panic!("Camera must have a film");
        }
        let film = &*film_ptr;

        let aspect = film.full_resolution().x() as Float / film.full_resolution().y() as Float;
        let diagonal = film.diagonal();
        let x = (square(diagonal) / (1.0 + square(diagonal))).sqrt();
        let y = x * aspect;
        let physical_extent =
            Bounds2f::from_points(Point2f::new(-x / 2., -y / 2.), Point2f::new(x / 2., y / 2.));
        let mut element_interface: Vec<LensElementInterface> = Vec::new();

        for i in (0..lens_params.len()).step_by(4) {
            let curvature_radius = lens_params[i] / 1000.0;
            let thickness = lens_params[i + 1] / 1000.0;
            let eta = lens_params[i + 2];
            let mut aperture_diameter = lens_params[i + 3] / 1000.0;

            if curvature_radius == 0.0 {
                aperture_diameter /= 1000.0;
                if set_aperture_diameter > aperture_diameter {
                    println!("Aperture is larger than possible")
                } else {
                    aperture_diameter = set_aperture_diameter;
                }
            }
            let el_int = LensElementInterface {
                curvature_radius,
                thickness,
                eta,
                aperture_radius: aperture_diameter / 2.0,
            };
            element_interface.push(el_int);
        }

        let half_diag = film.diagonal() / 2.0;
        let mut exit_pupil_bounds = [Bounds2f::default(); EXIT_PUPIL_SAMPLES];

        for i in 0..EXIT_PUPIL_SAMPLES {
            let r0 = (i as Float / EXIT_PUPIL_SAMPLES as Float) * half_diag;
            let r1 = ((i + 1) as Float / EXIT_PUPIL_SAMPLES as Float) * half_diag;
            exit_pupil_bounds[i] = Self::compute_exit_pupil_bounds(&element_interface, r0, r1);
        }

        let n_elements = element_interface.len();
        let element_interfaces = element_interface.as_ptr();
        std::mem::forget(element_interface);

        Self {
            base,
            focus_distance,
            element_interfaces: Ptr::from(element_interfaces),
            n_elements,
            physical_extent,
            set_aperture_diameter,
            aperture_image,
            exit_pupil_bounds,
        }
    }

    pub fn compute_cardinal_points(r_in: Ray, r_out: Ray) -> (Float, Float) {
        let tf = -r_out.o.x() / r_out.d.x();
        let tp = (r_in.o.x() - r_out.o.x()) / r_out.d.x();
        (-r_out.at(tf).z(), -r_out.at(tp).z())
    }

    pub fn compute_thick_lens_approximation(&self) -> ([Float; 2], [Float; 2]) {
        use crate::utils::Ptr;

        let x = 0.001 * self.get_film().diagonal();
        let r_scene = Ray::new(
            Point3f::new(0., x, self.lens_front_z() + 1.),
            Vector3f::new(0., 0., -1.),
            None,
            &Ptr::null(),
        );
        let Some((_, r_film)) = self.trace_lenses_from_film(&r_scene) else {
            panic!(
                "Unable to trace ray from scene to film for thick lens approx. Is aperture very small?"
            )
        };
        let (pz0, fz0) = Self::compute_cardinal_points(r_scene, r_film);
        let r_film = Ray::new(
            Point3f::new(x, 0., self.lens_rear_z() - 1.),
            Vector3f::new(0., 0., 1.),
            None,
            &Ptr::null(),
        );
        let Some((_, r_scene)) = self.trace_lenses_from_film(&r_film) else {
            panic!(
                "Unable to trace ray from scene to film for thick lens approx. Is aperture very small?"
            )
        };
        let (pz1, fz1) = Self::compute_cardinal_points(r_film, r_scene);
        ([pz0, pz1], [fz0, fz1])
    }

    pub fn focus_thick_lens(&self, focus_distance: Float) -> Float {
        let (pz, fz) = self.compute_thick_lens_approximation();
        let f = fz[0] - pz[0];
        let z = -focus_distance;
        let c = (pz[1] - z - pz[0]) * (pz[1] - z - 4. * f - pz[0]);
        if c <= 0. {
            panic!(
                "Coefficient must be positive. It looks focusDistance {}  is too short for a given lenses configuration",
                focus_distance
            );
        }
        let delta = (pz[1] - z + pz[0] - c.sqrt()) / 2.;
        let last_interface = unsafe { self.element_interfaces.add(self.n_elements) };
        last_interface.thickness + delta
    }

    pub fn bound_exit_pupil(&self, film_x_0: Float, film_x_1: Float) -> Bounds2f {
        let interface_array = unsafe {
            core::slice::from_raw_parts(self.element_interfaces.as_raw(), self.n_elements as usize)
        };
        Self::compute_exit_pupil_bounds(interface_array, film_x_0, film_x_1)
    }

    fn compute_exit_pupil_bounds(
        elements: &[LensElementInterface],
        film_x_0: Float,
        film_x_1: Float,
    ) -> Bounds2f {
        let mut pupil_bounds = Bounds2f::default();
        let n_samples = 1024 * 1024;
        let rear_radius = elements.last().unwrap().aperture_radius;
        let proj_rear_bounds = Bounds2f::from_points(
            Point2f::new(-1.5 * rear_radius, -1.5 * rear_radius),
            Point2f::new(1.5 * rear_radius, 1.5 * rear_radius),
        );

        let radical_inverse = |x: i32, _y: i64| x as Float;
        let lens_rear_z = || 1.;
        let trace_lenses_from_film = |_ray: Ray, _place: Option<Ray>| true;
        for i in 0..n_samples {
            // Find location of sample points on $x$ segment and rear lens element
            //
            let p_film = Point3f::new(
                lerp((i as Float + 0.5) / n_samples as Float, film_x_0, film_x_1),
                0.,
                0.,
            );
            let u: [Float; 2] = [radical_inverse(0, i), radical_inverse(1, i)];
            let p_rear = Point3f::new(
                lerp(u[0], proj_rear_bounds.p_min.x(), proj_rear_bounds.p_max.x()),
                lerp(u[1], proj_rear_bounds.p_min.y(), proj_rear_bounds.p_max.y()),
                lens_rear_z(),
            );

            // Expand pupil bounds if ray makes it through the lens system
            if !pupil_bounds.contains(Point2f::new(p_rear.x(), p_rear.y()))
                && trace_lenses_from_film(
                    Ray::new(p_film, p_rear - p_film, None, &Ptr::null()),
                    None,
                )
            {
                pupil_bounds = pupil_bounds.union_point(Point2f::new(p_rear.x(), p_rear.y()));
            }
        }

        if pupil_bounds.is_degenerate() {
            print!(
                "Unable to find exit pupil in x = {},{} on film.",
                film_x_0, film_x_1
            );
            return pupil_bounds;
        }

        pupil_bounds.expand(2. * proj_rear_bounds.diagonal().norm() / (n_samples as Float).sqrt());
        pupil_bounds
    }

    pub fn sample_exit_pupil(&self, p_film: Point2f, u_lens: Point2f) -> Option<ExitPupilSample> {
        // Find exit pupil bound for sample distance from film center
        let film = self.get_film();
        let r_film = (square(p_film.x()) + square(p_film.y())).sqrt();
        let mut r_index = (r_film / (film.diagonal() / 2.)) as usize * self.exit_pupil_bounds.len();
        r_index = (self.exit_pupil_bounds.len() - 1).min(r_index);

        let pupil_bounds = self.exit_pupil_bounds[r_index];
        if pupil_bounds.is_degenerate() {
            return None;
        }

        // Generate sample point inside exit pupil bound
        let p_lens = pupil_bounds.lerp(u_lens);
        let pdf = 1. / pupil_bounds.area();

        // Return sample point rotated by angle of _pFilm_ with $+x$ axis
        let sin_theta = if r_film != 0. {
            p_film.y() / r_film
        } else {
            0.
        };
        let cos_theta = if r_film != 0. {
            p_film.x() / r_film
        } else {
            1.
        };
        let p_pupil = Point3f::new(
            cos_theta * p_lens.x() - sin_theta * p_lens.y(),
            sin_theta * p_lens.x() + cos_theta * p_lens.y(),
            self.lens_rear_z(),
        );

        Some(ExitPupilSample { p_pupil, pdf })
    }

    pub fn trace_lenses_from_film(&self, r_camera: &Ray) -> Option<(Float, Ray)> {
        let mut element_z = 0.;
        let weight = 1.;
        // Transform _r_camera_ from camera to lens system space
        let mut r_lens = Ray::new(
            Point3f::new(r_camera.o.x(), r_camera.o.y(), -r_camera.o.z()),
            Vector3f::new(r_camera.d.x(), r_camera.d.y(), -r_camera.d.z()),
            Some(r_camera.time),
            &Ptr::null(),
        );

        for i in (0..self.n_elements - 1).rev() {
            let element: &LensElementInterface = unsafe { &self.element_interfaces.add(i) };
            // Update ray from film accounting for interaction with _element_
            element_z -= element.thickness;

            let is_stop = element.curvature_radius == 0.;
            let t: Float;
            let mut n = Normal3f::default();
            if is_stop {
                // Compute _t_ at plane of aperture stop
                t = (element_z - r_lens.o.z()) / r_lens.d.z();
            } else {
                // Intersect ray with element to compute _t_ and _n_
                let radius = element.curvature_radius;
                let z_center = element_z + element.curvature_radius;
                if let Some((intersect_t, intersect_n)) =
                    RealisticCamera::intersect_spherical_element(radius, z_center, &r_lens)
                {
                    t = intersect_t;
                    n = intersect_n;
                } else {
                    return None; // Ray missed the element
                }
            }

            if t < 0. {
                return None;
            }

            // Test intersection point against element aperture
            let p_hit = r_lens.at(t);
            if square(p_hit.x()) + square(p_hit.y()) > square(element.aperture_radius) {
                return None;
            }

            r_lens.o = p_hit;

            // Update ray path for element interface interaction
            if !is_stop {
                let eta_i = element.eta;
                let interface_i = unsafe { self.element_interfaces.add(i - 1) };
                let eta_t = if i > 0 && interface_i.eta != 0. {
                    interface_i.eta
                } else {
                    1.
                };
                let wi = -r_lens.d.normalize();
                let eta_ratio = eta_t / eta_i;

                // Handle refraction idiomatically
                if let Some((wt, _final_eta)) = refract(wi, n, eta_ratio) {
                    r_lens.d = wt;
                } else {
                    // Total internal reflection occurred
                    return None;
                }
            }
        }

        // Transform lens system space ray back to camera space
        let r_out = Ray::new(
            Point3f::new(r_lens.o.x(), r_lens.o.y(), -r_lens.o.z()),
            Vector3f::new(r_lens.d.x(), r_lens.d.y(), -r_lens.d.z()),
            Some(r_lens.time),
            &Ptr::null(),
        );

        Some((weight, r_out))
    }

    fn intersect_spherical_element(
        radius: Float,
        z_center: Float,
        ray: &Ray,
    ) -> Option<(Float, Normal3f)> {
        let o = ray.o - Vector3f::new(0.0, 0.0, z_center);

        let a = ray.d.norm_squared();
        let b = 2.0 * ray.d.dot(o.into());
        let c = Vector3f::from(o).norm_squared() - radius * radius;

        let (t0, t1) = quadratic(a, b, c)?;

        let use_closer_t = (ray.d.z() > 0.0) ^ (radius < 0.0);
        let t = if use_closer_t { t0.min(t1) } else { t0.max(t1) };

        // Intersection is behind the ray's origin.
        if t < 0.0 {
            return None;
        }

        let p_hit_relative = o + ray.d * t;
        // Ensures the normal points towards the incident ray.
        let n = Normal3f::from(Vector3f::from(p_hit_relative))
            .normalize()
            .face_forward(-ray.d);

        Some((t, n))
    }

    pub fn lens_rear_z(&self) -> Float {
        let last_interface = unsafe { self.element_interfaces.add(self.n_elements - 1) };
        last_interface.thickness
    }

    pub fn lens_front_z(&self) -> Float {
        let mut z_sum = 0.;
        for i in 0..self.n_elements {
            let element = unsafe { self.element_interfaces.add(i) };
            z_sum += element.thickness;
        }
        z_sum
    }

    pub fn rear_element_radius(&self) -> Float {
        let last_interface = unsafe { self.element_interfaces.add(self.n_elements - 1) };
        last_interface.aperture_radius
    }
}

impl CameraTrait for RealisticCamera {
    fn base(&self) -> &CameraBase {
        &self.base
    }

    fn generate_ray(
        &self,
        sample: CameraSample,
        _lambda: &SampledWavelengths,
    ) -> Option<CameraRay> {
        let film = self.get_film();
        let s = Point2f::new(
            sample.p_film.x() / film.full_resolution().x() as Float,
            sample.p_film.y() / film.full_resolution().y() as Float,
        );
        let p_film2 = self.physical_extent.lerp(s);
        let p_film = Point3f::new(-p_film2.x(), p_film2.y(), 0.);

        // Trace ray from _pFilm_ through lens system
        let eps = self.sample_exit_pupil(Point2f::new(p_film.x(), p_film.y()), sample.p_lens)?;

        let p_pupil = Point3f::new(0., 0., 0.);
        let r_film = Ray::new(p_film, p_pupil - p_film, None, &Ptr::null());
        let (weight, mut ray) = self.trace_lenses_from_film(&r_film)?;
        if weight == 0. {
            return None;
        }

        // Finish initialization of _RealisticCamera_ ray
        ray.time = self.sample_time(sample.time);
        ray.medium = self.base.medium.clone();
        ray = self.render_from_camera(&ray, &mut None);
        ray.d = ray.d.normalize();

        // Compute weighting for _RealisticCamera_ ray
        let cos_theta = r_film.d.normalize().z();
        let final_weight =
            weight * cos_theta.powf(4.) / (eps.pdf as Float * square(self.lens_rear_z()));

        Some(CameraRay {
            ray,
            weight: SampledSpectrum::new(final_weight),
        })
    }
}