// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at https://mozilla.org/MPL/2.0/.

//! Scene Definition and Rendering Control.

mod light;
mod object;

pub use self::light::Light;
pub(crate) use self::object::Object;
pub use self::object::Segment;

use crate::geom::{Point, Rect};
use crate::image::RenderImage;
use crate::ray::{Ray, RayResult};

use rand::prelude::*;
use rand_pcg::Pcg64Mcg;

use std::sync::Arc;
use std::thread;
use std::time;

/// Render termination predicate.
#[derive(Clone, Copy, Debug)]
pub enum RenderConstraint {
    /// Render Will terminate after a fixed number of rays have been generated by the light sources.
    /// This value does not count rays generated by bounces.
    Count(usize),
    /// Render will terminate after a fixed number of miliseconds.
    TimeMS(usize),
}

impl From<usize> for RenderConstraint {
    /// `usize` is interpreted as a fixed number of rays, mostly for backwards compatiblity.
    fn from(value: usize) -> Self {
        Self::Count(value)
    }
}

const BATCH_SIZE: usize = 1000;

type R = Pcg64Mcg;

/// Holds scene Configuration and logic
#[derive(Clone)]
pub struct Scene {
    lights: Vec<Light<R>>,
    objects: Vec<Object<R>>,
    total_light_power: f32,
    viewport: Rect,
}

impl Scene {
    /// Creates new Renderer ready for defining a scene.
    pub fn new(resolution_x: usize, resolution_y: usize) -> Self {
        Self {
            lights: vec![],
            objects: vec![],
            viewport: Rect::from_points(
                &Point { x: 0.0, y: 0.0 },
                &Point {
                    x: resolution_x as f64,
                    y: resolution_y as f64,
                },
            ),
            total_light_power: 0.0,
        }
    }

    /// Adds a light to the scene
    pub fn add_light(&mut self, light: Light<R>) {
        self.total_light_power += light.power.bounds().1 as f32;
        self.lights.push(light);
    }

    /// Adds a light to the scene - Chainable variant
    pub fn with_light(mut self, light: Light<R>) -> Self {
        self.add_light(light);
        self
    }

    /// Adds an object to the scene.
    /// `Object` is an internal representation used by the raytracer, `A`
    /// should be anything that implements `Into<Object>` currently only
    /// `Segment` provides this.
    ///
    /// # Example Usage:
    /// ```
    /// use rustic_zen::prelude::*;
    /// let a = (0.0, 0.0);
    /// let b = (10.0, 10.0);
    /// // Or other math to construct your scene here,
    /// let line = Segment::line_from_points(a, b, material::hqz_legacy_default());
    ///
    /// let mut s = Scene::new(1920,1080);
    /// s.add_object(line); //Segment automatically converted to hidden object type here.
    /// ```
    #[allow(private_bounds)]
    pub fn add_object<A>(&mut self, object: A)
    where
        A: Into<Object<R>>,
    {
        self.objects.push(object.into());
    }

    /// Adds an object to the scene - Chainable Variant.
    /// `Object` is an internal representation used by the raytracer, `A`
    /// should be anything that implements `Into<Object>` currently only
    /// `Segment` provides this.
    ///
    /// # Example Usage:
    /// ```
    /// use rustic_zen::prelude::*;
    /// let a = Point::new(0.0, 0.0);
    /// let b = Point::new(10.0, 10.0);
    /// // Or other math to construct your scene here,
    /// let line = Segment::line_from_points(a, b, material::hqz_legacy_default());
    ///
    /// //Segment automatically converted to hidden object type here.
    /// let s = Scene::new(1920,1080).with_object(line);
    /// ```
    #[allow(private_bounds)]
    pub fn with_object<A>(mut self, object: A) -> Self
    where
        A: Into<Object<R>>,
    {
        self.add_object(object);
        self
    }

    #[inline(always)]
    fn choose_light(&self, rng: &mut R) -> &Light<R> {
        let threshold = rng.gen_range(0.0..self.total_light_power) as f64;
        let mut sum: f64 = 0.0;
        for light in &self.lights {
            sum += light.power.sample(rng);
            if threshold <= sum {
                return light;
            }
        }
        return self.lights.last().expect("Scene has no lights");
    }

    #[inline(always)]
    fn trace_ray<F>(&self, rng: &mut R, mut on_ray: F) -> usize
    where
        F: FnMut(RayResult) -> (),
    {
        let l = self.choose_light(rng);
        let mut rays = 0;
        let mut ray = Some(Ray::new(l, rng));
        while ray.is_some() {
            let (result, new_ray) = ray.unwrap().collision_list(&self.objects, self.viewport);
            rays += 1;
            if result.is_some() {
                (on_ray)(result.unwrap());
            }
            ray = new_ray;
        }
        rays
    }

    #[inline(always)]
    fn render_thread<F>(&self, constraint: RenderConstraint, mut on_ray: F) -> usize
    where
        F: FnMut(RayResult) -> (),
    {
        let mut local_rng = R::from_entropy();
        match constraint {
            RenderConstraint::Count(n) => (0..n).fold(0usize, |r, _| {
                r + self.trace_ray(&mut local_rng, &mut on_ray)
            }),
            RenderConstraint::TimeMS(t) => {
                let start = time::Instant::now();
                let mut rays = 0;
                while start.elapsed().as_millis() < t as u128 {
                    rays += (0..BATCH_SIZE).fold(0usize, |r, _| {
                        r + self.trace_ray(&mut local_rng, &mut on_ray)
                    });
                }
                rays
            }
        }
    }

    /// Starts the ray tracing process
    ///
    /// Naturally this call is very expensive. It also consumes the Renderer
    /// and populates the provided `Image` with the rendered image data.
    /// The total number of rays cast into the scene is returned.
    ///
    /// # Parameters:
    /// __constraint__: A `RenderConstraint` object, which defines the condition to stop rendering.
    /// __threads__: the number of cpu threads the render will use. The renderer will automatically spawn this many threads.
    /// __image__: a `Image` to be populated with image data.
    ///
    /// The renderer will automatically use the GPU if a `VulkanImage` is passed into the image slot, otherwise the image will be rendered solely in software..
    ///
    /// # Returns:
    /// `usize` The first value is the number of rays added to the scene. This is needed for some image export functions.
    /// The second value is the resulting image as an `Image` object, which has serveral export functions.
    pub fn render<C>(
        self,
        constraint: C,
        threads: usize,
        image: &mut Arc<impl RenderImage + 'static>,
    ) -> usize
    where
        C: Into<RenderConstraint>,
    {
        let threads = match threads {
            0 => num_cpus::get(),
            i => i,
        };

        let con = match constraint.into() {
            RenderConstraint::Count(n) => RenderConstraint::Count(n / (threads)),
            RenderConstraint::TimeMS(t) => RenderConstraint::TimeMS(t),
        };

        Arc::get_mut(image)
            .expect("need exclusive access to image")
            .prepare_render(self.total_light_power);

        let s = Arc::new(self);
        let join_handles: Vec<_> = (0..threads)
            .map(|_| {
                let local_s = s.clone();
                let local_image = image.clone();
                thread::spawn(move || local_s.render_thread(con, |r| local_image.draw_line(r)))
            })
            .collect();

        let rays_cast = join_handles
            .into_iter()
            .fold(0, |r, h| r + h.join().unwrap());

        Arc::get_mut(image).unwrap().finish_render();

        rays_cast
    }
}

#[cfg(test)]
mod tests {
    use super::object::Segment;
    use super::Scene;
    use crate::image::Image;
    use crate::material::hqz_legacy;
    use crate::sampler::Sampler;
    use crate::scene::Light;
    use std::sync::Arc;

    #[test]
    fn nrt_works() {
        let obj = Segment::line_from_points((0.0, 0.0), (10.0, 10.0), hqz_legacy(0.3, 0.3, 0.3));

        let l = Light {
            power: Sampler::new_const(1.0),
            location: (10.0, 10.0).into(),
            polar_angle: Sampler::new_range(360.0, 0.0),
            polar_distance: Sampler::new_const(0.0),
            ray_angle: Sampler::new_range(360.0, 0.0),
            wavelength: Sampler::new_blackbody(5800.0),
        };

        let r = Scene::new(1920, 1080).with_light(l).with_object(obj);

        let mut img = Arc::new(Image::new(1920, 1080));

        r.render(super::RenderConstraint::TimeMS(1000), 1, &mut img);
    }
}