use super::{CameraBase, CameraRay, CameraTrait, ExitPupilSample, LensElementInterface};
use crate::core::film::FilmTrait;
use crate::core::pbrt::{Float, lerp};
use crate::core::sampler::CameraSample;
use crate::geometry::{
    Bounds2f, Normal3f, Point2f, Point2i, Point3f, Ray, Vector2i, Vector3f, VectorLike,
};
use crate::image::Image;
use crate::spectra::{SampledSpectrum, SampledWavelengths};
use crate::utils::math::{quadratic, square};
use crate::utils::scattering::refract;

pub struct RealisticCamera {
    base: CameraBase,
    focus_distance: Float,
    set_aperture_diameter: Float,
    aperture_image: Image,
    element_interface: Vec<LensElementInterface>,
    physical_extent: Bounds2f,
    exit_pupil_bounds: Vec<Bounds2f>,
}

impl RealisticCamera {
    pub fn new(
        &self,
        base: CameraBase,
        lens_params: Vec<Float>,
        focus_distance: Float,
        set_aperture_diameter: Float,
        aperture_image: Image,
    ) -> Self {
        let aspect =
            base.film.full_resolution().x() as Float / base.film.full_resolution().y() as Float;
        let diagonal = base.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 n_samples = 64;
        let half_diag = base.film.diagonal() / 2.0;
        let exit_pupil_bounds: Vec<_> = (0..n_samples)
            .map(|i| {
                let r0 = (i as Float / n_samples as Float) * half_diag;
                let r1 = ((i + 1) as Float / n_samples as Float) * half_diag;
                self.bound_exit_pupil(r0, r1)
            })
            .collect();

        Self {
            base,
            focus_distance,
            element_interface,
            physical_extent,
            set_aperture_diameter,
            aperture_image,
            exit_pupil_bounds,
        }
    }

    pub fn lens_rear_z(&self) -> Float {
        self.element_interface.last().unwrap().thickness
    }

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

    pub fn rear_element_radius(&self) -> Float {
        self.element_interface.last().unwrap().aperture_radius
    }

    pub fn bound_exit_pupil(&self, film_x_0: Float, film_x_1: Float) -> Bounds2f {
        let mut pupil_bounds = Bounds2f::default();
        let n_samples = 1024 * 1024;
        let rear_radius = self.rear_element_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, None), None)
            {
                pupil_bounds = pupil_bounds.union_point(Point2f::new(p_rear.x(), p_rear.y()));
            }
        }

        // Return degenerate bounds if no rays made it through the lens system
        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())
    }

    pub 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 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),
            None,
        );

        for i in (0..self.element_interface.len() - 1).rev() {
            let element: &LensElementInterface = &self.element_interface[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.evaluate(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 eta_t = if i > 0 && self.element_interface[i - 1].eta != 0. {
                    self.element_interface[i - 1].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),
            None,
        );

        Some((weight, r_out))
    }

    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 r_film = (square(p_film.x()) + square(p_film.y())).sqrt();
        let mut r_index =
            (r_film / (self.base.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 })
    }
}

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

    fn generate_ray(
        &self,
        sample: CameraSample,
        _lambda: &SampledWavelengths,
    ) -> Option<CameraRay> {
        // Find point on film, _pFilm_, corresponding to _sample.pFilm_
        let s = Point2f::new(
            sample.p_film.x() / self.base.film.full_resolution().x() as Float,
            sample.p_film.y() / self.base.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, None);
        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),
        })
    }
}