// 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/.

//! Wrappers for stochastically sampled variables.

use crate::geom::Point;
use crate::spectrum::blackbody_wavelength;

use rand::prelude::*;
use rand_distr::num_traits::{Float, FromPrimitive};
use rand_distr::uniform::SampleUniform;
use rand_distr::Distribution;
use rand_distr::{Normal, StandardNormal};

use std::sync::Arc;

/// Wrapper for a bounded stocastically sampled value
///
/// a `Sampler` encloses the distribution function and bounds for a random
/// variable of type `T`, such that is can be sampled with an external random
/// number generator of type `R`
///
/// Several constructors are provided for common distributions of values, however
/// if a custom one is needed it can be provided as a closure and set of bounds to
/// `Sampler::from_fn`
#[derive(Clone)]
pub struct Sampler<T: Copy, R: Rng> {
    sample: Arc<dyn Fn(&mut R) -> T + Send + Sync>,
    bounds: (T, T),
}

impl<T, R> Sampler<T, R>
where
    R: Rng,
    T: Copy + Send + Sync + 'static,
{
    /// Creates a new `Sampler` with a constant distibution.
    ///
    /// When sampled the `Sampler` will always return `c`
    pub fn new_const<A>(c: A) -> Self
    where
        A: Into<T>,
    {
        let c = c.into();
        Self {
            sample: Arc::new(move |_| c),
            bounds: (c, c),
        }
    }
}

impl<T, R> Sampler<T, R>
where
    R: Rng,
    T: Copy + Float + SampleUniform + Send + Sync + 'static,
{
    /// Creates a new `Sampler` with a uniform distibution.
    ///
    /// When sampled the `Sampler` will always return a value between `a` and `b` (inclusive)
    pub fn new_range<A, B>(a: A, b: B) -> Self
    where
        A: Into<T>,
        B: Into<T>,
    {
        let a = a.into();
        let b = b.into();
        if a == b {
            return Self::new_const(a);
        }
        let lower = T::min(a, b);
        let upper = T::max(a, b);
        Self {
            sample: Arc::new(move |r| r.gen_range(lower..=upper)),
            bounds: (lower, upper),
        }
    }
}

impl<T, R> Sampler<T, R>
where
    R: Rng,
    T: Copy + Float + From<f32> + Send + Sync + 'static,
    StandardNormal: Distribution<T>,
{
    /// Creates a new `Sampler` with a normal (gaussian) distibution.
    ///
    /// Real gaussian's have infinate bounds, this can cause problems in the raytracing engine, so
    /// this implementation wraps at 3 standard deviations.
    ///
    /// When sampled the `Sampler` will always return a value between `mean - 3 * std_dev` and `mean + 3 * std_dev` (inclusive)
    pub fn new_gaussian<A, B>(mean: A, std_dev: B) -> Self
    where
        A: Copy + Into<T>,
        B: Copy + Into<T>,
    {
        let width: T = std_dev.into() * 3.0.into();
        let mean = mean.into();
        let g = Normal::new(0.0.into(), std_dev.into()).unwrap();
        Self {
            sample: Arc::new(move |r| (g.sample(r) % width) + mean),
            bounds: (mean - width, mean + width),
        }
    }
}

impl<T, R> Sampler<T, R>
where
    R: Rng,
    T: Copy + Float + SampleUniform + FromPrimitive + From<f32> + Send + Sync + 'static,
{
    /// Creates a new `Sampler` with a blackbody radiation curve distibution.
    ///
    /// This is only really useful for creating realistic white lights of a given colour temprature, or stars.
    ///
    /// This sampler should be used carefully; it can lead to unexpected results when used on a value other than `Light.wavelength`.
    /// It is also __substancially__ slower than the other `Sampler` types. When Sampled it will return a value between
    /// 0 and `T::MAX` (probably larger than you want) This can cause all sorts of weird problems for collision detection, so _please_
    /// don't use this on a `SamplerPoint`
    pub fn new_blackbody<A>(temperature: A) -> Self
    where
        A: Into<T>,
    {
        let t = temperature.into();
        Self {
            sample: Arc::new(move |r| blackbody_wavelength(t, r.gen_range(0.0.into()..1.0.into()))),
            bounds: (0.0.into(), T::max_value()),
        }
    }
}

impl<T, R> Sampler<T, R>
where
    T: Copy,
    R: Rng,
{
    /// Creates a new `Sampler` with a distribution dictated by `f`
    ///
    /// `f` must always return a value within the bounds specified by `lower` & `upper`
    ///
    /// This is very low level access to how the renderer works; you're in control,
    /// if you break it it's your fault!
    pub fn from_fn(f: Arc<dyn Fn(&mut R) -> T + Send + Sync>, lower: T, upper: T) -> Self {
        Self {
            sample: f,
            bounds: (lower, upper),
        }
    }

    /// Get a sampled value from this `Sampler` using random number generator `rng`
    #[inline(always)]
    pub fn sample(&self, rng: &mut R) -> T {
        (self.sample)(rng)
    }

    /// Get the bounds of a possible sampled value (useful for spacial partitioning and limit checks etc)
    #[inline(always)]
    pub fn bounds(&self) -> (T, T) {
        self.bounds
    }
}

/// A single value of type `T` will be interpreted as a constant distribution if turned into a sampler.
///
/// # Example:
/// ```
/// extern crate rand;
///
/// use rustic_zen::sampler::Sampler;
/// use rand::prelude::*;
/// let mut r = rand::thread_rng();
///
/// let s: Sampler<f64, ThreadRng> = 5.0.into();
///
/// assert_eq!(s.sample(&mut r), 5.0); // will always be true
/// ```
impl<T, R> From<T> for Sampler<T, R>
where
    T: Copy + From<T> + Send + Sync + 'static,
    R: Rng,
{
    fn from(value: T) -> Self {
        Self::new_const(value)
    }
}

/// A tuple of type `(T,T)` will be interpreted as a uniform distribution if turned into a sampler.
///
/// # Example:
/// ```
/// extern crate rand;
///
/// use rustic_zen::sampler::Sampler;
/// use rand::prelude::*;
/// let mut r = rand::thread_rng();
///
/// let s: Sampler<f64, ThreadRng> = (5.0,10.0).into();
///
/// assert!(s.sample(&mut r) >= 5.0); // will always be true
/// assert!(s.sample(&mut r) <= 10.0); // will always be true
/// ```
impl<T, R> From<(T, T)> for Sampler<T, R>
where
    T: Copy + Float + SampleUniform + Send + Sync + From<f32> + 'static,
    R: Rng,
{
    fn from(value: (T, T)) -> Self {
        Self::new_range(value.0, value.1)
    }
}

/// Wrapper around two `Sampler`s to make a Stochasically sampled point.
#[derive(Clone)]
pub struct SamplerPoint<R: Rng> {
    x: Sampler<f64, R>,
    y: Sampler<f64, R>,
}

/// A tuple of type `(Sampler::<T>,Sampler<T>)` will be interpreted as a uniform distribution if turned into a sampler.
///
/// # Example:
/// ```
/// extern crate rand;
///
/// use rustic_zen::sampler::SamplerPoint;
/// use rustic_zen::sampler::Sampler;
/// use rand::prelude::*;
/// let mut r = rand::thread_rng();
///
/// // results in a rounds soft distribution, useful for area lights.
/// let _p: SamplerPoint<ThreadRng> = (Sampler::new_gaussian(0.0, 3.0), Sampler::new_gaussian(0.0, 3.0)).into();
///
/// ```
///
/// This is the only way to construct a `SamplerPoint`. This shorthand can be extended using the `Sampler`'s into traits:
/// ```
/// extern crate rand;
///
/// use rustic_zen::sampler::SamplerPoint;
/// use rustic_zen::geom::Point;
/// use rand::prelude::*;
/// let mut r = rand::thread_rng();
///
/// // results in a fixed point at 0.0, 0.0
/// let _p: SamplerPoint<ThreadRng> = (0.0, 0.0).into();
///
/// // results in a uniformly distributed square from 0.0, 0.0 to 10.0, 10.0
/// let _p: SamplerPoint<ThreadRng> = ((0.0, 10.0), (0.0, 10.0)).into();
/// ```
impl<R> From<Point> for SamplerPoint<R>
where
    R: Rng,
{
    fn from(value: Point) -> Self {
        Self {
            x: value.x.into(),
            y: value.y.into(),
        }
    }
}

impl<A, B, R> From<(A, B)> for SamplerPoint<R>
where
    A: Into<Sampler<f64, R>>,
    B: Into<Sampler<f64, R>>,
    R: Rng,
{
    fn from(value: (A, B)) -> Self {
        Self {
            x: value.0.into(),
            y: value.1.into(),
        }
    }
}

impl<R> SamplerPoint<R>
where
    R: Rng,
{
    #[inline(always)]
    pub(crate) fn get(&self, rng: &mut R) -> Point {
        Point {
            x: self.x.sample(rng),
            y: self.y.sample(rng),
        }
    }
}

#[cfg(test)]
mod tests {
    type RandGen = rand_pcg::Pcg64Mcg;
    use rand::prelude::*;

    use super::Sampler;

    #[test]
    fn const_bounds() {
        // test bounds with 1000 random numbers
        for _ in 0..1000 {
            let mut stdrng = RandGen::from_entropy();
            let f: f64 = stdrng.gen();
            let s = Sampler::<f64, RandGen>::new_const(f);
            let (a, b) = s.bounds();
            assert_eq!(a, b);
            assert_eq!(a, f);
            assert_eq!(b, f);
        }
    }

    #[test]
    fn range_bounds() {
        for _ in 0..1000 {
            let mut stdrng = RandGen::from_entropy();
            let a: f64 = stdrng.gen();
            let b: f64 = stdrng.gen();
            let s = Sampler::<f64, RandGen>::new_range(a, b);

            let (c, d) = s.bounds();
            assert_eq!(a.min(b), c);
            assert_eq!(a.max(b), d);
        }
    }

    #[test]
    fn val_const() {
        let mut rng = RandGen::from_entropy();

        let mut stdrng = RandGen::from_entropy();
        let f: f64 = stdrng.gen();
        let s = Sampler::new_const(f);

        for _ in 0..1000 {
            let y: f64 = s.sample(&mut rng);
            assert_eq!(y, f);
        }
    }

    #[test]
    fn val_range() {
        let mut rng = RandGen::from_entropy();

        let mut stdrng = RandGen::from_entropy();
        let mut f1: f64 = stdrng.gen();
        let mut f2: f64 = stdrng.gen();
        if f1 < f2 {
            let tmp = f1;
            f1 = f2;
            f2 = tmp;
        }
        let s = Sampler::new_range(f1, f2);

        // This does actually run 100,000 times despite finishing so quickly
        for _ in 0..100000 {
            let y: f64 = s.sample(&mut rng);
            assert!(y <= f1);
            assert!(y >= f2);
        }
    }

    #[test]
    fn blackbody_works() {
        let mut rng = RandGen::from_entropy();
        // Get a sampled value from the range of valid wavelenghts
        let w: Sampler<f64, rand_pcg::Mcg128Xsl64> = Sampler::new_range(780.0, 360.0);

        // create sample with random wavelenght
        let s = Sampler::new_blackbody(w.sample(&mut rng));

        // Check s can be sampled without panicing
        let _: f64 = s.sample(&mut rng);
    }

    #[test]
    fn blackbody_white_light_64() {
        let mut rng = RandGen::from_entropy();
        let s = Sampler::new_blackbody(0.0);
        // Check s can be sampled without panicing
        let _: f64 = s.sample(&mut rng);
    }

    #[test]
    fn blackbody_white_light_32() {
        let mut rng = RandGen::from_entropy();
        let s = Sampler::new_blackbody(0.0);
        // Check s can be sampled without panicing
        let _: f32 = s.sample(&mut rng);
    }

    #[test]
    fn gaussian_at_zero() {
        let mut rng = RandGen::from_entropy();
        let s = Sampler::new_gaussian(0.0, 10.0);
        // Check s can be sampled without panicing
        let _: f64 = s.sample(&mut rng);
    }

    #[test]
    fn range_big_first() {
        let mut rng = RandGen::from_entropy();
        let s = Sampler::new_range(10.0, 0.0);
        // Check s can be sampled without panicing
        let _: f64 = s.sample(&mut rng);
    }

    #[test]
    fn range_small_first() {
        let mut rng = RandGen::from_entropy();
        let s = Sampler::new_range(0.0, 10.0);
        // Check s can be sampled without panicing
        let _: f64 = s.sample(&mut rng);
    }
}