//! Multiscale Mean Estimator described in the [Hybrid Rendering for Real-Time Ray Tracing](https://media.contentapi.ea.com/content/dam/ea/seed/presentations/2019-ray-tracing-gems-chapter-25-barre-brisebois-et-al.pdf) and [Surfel-based global illumination on the web](https://juretriglav.si/surfel-based-global-illumination-on-the-web/)
//!
//! ## On `short_window_blend`
//!
//! This value controls how _fast_ the filter reacts to changes in the _short-memory_. Good values are `0.01..=0.1`.
//! Values closer to 0 create a _smoother_ signal at the cost of higher latency. A higher value make the signal react faster, but keeps more of the input-noise.

use crate::{
    Estimator,
    glam::{FloatExt, Vec3},
    smoothstep,
};

#[derive(Clone)]
#[repr(C, align(16))]
///Mean estimator for 3-component (color) vector.
pub struct MultiscaleEstimator3 {
    pub mean: Vec3,
    ///Variance based blend reduction
    pub vbbr: f32,
    pub short_mean: Vec3,
    pub inconsistency: f32,
    pub variance: Vec3,
    /// Sets how much the _short_ variance influences the overall
    /// variance. Basically, if you signal contains _fireflies_, this should be greater, and if it contains less, it should be smaller. `default=0.08`.
    pub short_window_blend: f32,
}

impl MultiscaleEstimator3 {
    ///Initializes the mean estimator, assuming that its _currently_ stable at the given
    /// mean value
    pub fn new(mean: Vec3) -> Self {
        MultiscaleEstimator3 {
            mean,
            vbbr: 0.0,
            short_mean: mean,
            inconsistency: 0.0,
            variance: Vec3::ZERO,
            short_window_blend: 0.08,
        }
    }
}

impl Default for MultiscaleEstimator3 {
    fn default() -> Self {
        MultiscaleEstimator3 {
            mean: Vec3::default(),
            short_mean: Vec3::default(),
            vbbr: 0.0,
            variance: Vec3::default(),
            inconsistency: 0.0,
            short_window_blend: 0.08,
        }
    }
}

impl Estimator for MultiscaleEstimator3 {
    type Val = Vec3;
    fn push(&mut self, new: Self::Val) -> Self::Val {
        let mut mean = self.mean;
        let mut short_mean = self.short_mean;
        let mut vbbr = self.vbbr;
        let mut variance = self.variance;
        let mut inconsistency = self.inconsistency;
        let short_window_blend = self.short_window_blend;

        //supress fireflies
        let y = {
            let dev = variance.max(Vec3::splat(1e-5f32)).sqrt();
            let high_threshold = 0.1 + short_mean + dev * 8.0;
            let overflow = Vec3::ZERO.max(new - high_threshold);
            new - overflow
        };

        let delta = y - short_mean;
        short_mean = short_mean.lerp(y, short_window_blend);

        let delta2 = y - short_mean;

        // This should be a longer window than shortWindowBlend to avoid bias
        // from the variance getting smaller when the short - term mean does .
        let variance_blend = short_window_blend * 0.5;
        variance = variance.lerp(delta * delta2, variance_blend);
        let dev = Vec3::splat(1e-5f32).max(variance).sqrt();

        let short_diff = mean - short_mean;

        let relative_diff =
            Vec3::new(0.299, 0.587, 0.114).dot(short_diff.abs() / Vec3::splat(1e-5f32).max(dev));
        inconsistency = inconsistency.lerp(relative_diff, 0.08);

        let variance_based_blend_reduction = Vec3::new(0.299, 0.587, 0.114)
            .dot(0.5 * short_mean / Vec3::splat(1e-5f32).max(dev))
            .clamp(1.0 / 32.0, 1.0);

        let mut catch_up_blend =
            smoothstep(0.0, 1.0, relative_diff * 0.02f32.max(inconsistency - 0.2))
                .clamp(1.0 / 256.0, 1.0);
        catch_up_blend = catch_up_blend * vbbr;
        vbbr = vbbr.lerp(variance_based_blend_reduction, 0.1);
        mean = mean.lerp(y, catch_up_blend.saturate());

        self.mean = mean;
        self.short_mean = short_mean;
        self.vbbr = vbbr;
        self.variance = variance;
        self.inconsistency = inconsistency;

        mean
    }
}

#[derive(Clone)]
#[repr(C, align(16))]
///Mean estimator for a single floating-point value
pub struct MultiscaleEstimator {
    pub mean: f32,
    ///Variance based blend reduction
    pub vbbr: f32,
    pub short_mean: f32,
    pub inconsistency: f32,
    pub variance: f32,
    /// Sets how much the _short_ variance influences the overall
    /// variance. Basically, if you signal contains _fireflies_, this should be greater, and if it contains less, it should be smaller. `default=0.08`.
    pub short_window_blend: f32,
}

impl MultiscaleEstimator {
    ///Initializes the mean estimator, assuming that its _currently_ stable at the given
    /// mean value.
    pub fn new(mean: f32) -> Self {
        MultiscaleEstimator {
            mean,
            vbbr: 0.0,
            short_mean: mean,
            inconsistency: 0.0,
            variance: 0.0,
            short_window_blend: 0.08,
        }
    }
}

impl Default for MultiscaleEstimator {
    fn default() -> Self {
        MultiscaleEstimator {
            mean: 0.0,
            short_mean: 0.0,
            vbbr: 0.0,
            variance: 0.0,
            inconsistency: 0.0,
            short_window_blend: 0.08,
        }
    }
}

impl Estimator for MultiscaleEstimator {
    type Val = f32;
    fn push(&mut self, new: Self::Val) -> Self::Val {
        let mut mean = self.mean;
        let mut short_mean = self.short_mean;
        let mut vbbr = self.vbbr;
        let mut variance = self.variance;
        let mut inconsistency = self.inconsistency;
        let short_window_blend = self.short_window_blend;

        //supress fireflies
        let y = {
            let dev = variance.max(1e-5f32).sqrt();
            let high_threshold = 0.1 + short_mean + dev * 8.0;
            let overflow = 0.0f32.max(new - high_threshold);
            new - overflow
        };

        let delta = y - short_mean;
        short_mean = short_mean.lerp(y, short_window_blend);

        let delta2 = y - short_mean;

        // This should be a longer window than shortWindowBlend to avoid bias
        // from the variance getting smaller when the short - term mean does .
        let variance_blend = short_window_blend * 0.5;
        variance = variance.lerp(delta * delta2, variance_blend);
        let dev = 1e-5f32.max(variance).sqrt();

        let short_diff = mean - short_mean;

        let relative_diff = short_diff.abs() / 1e-5f32.max(dev);
        inconsistency = inconsistency.lerp(relative_diff, 0.08);

        let variance_based_blend_reduction =
            (0.5 * short_mean / 1e-5f32.max(dev)).clamp(1.0 / 32.0, 1.0);

        let mut catch_up_blend =
            smoothstep(0.0, 1.0, relative_diff * 0.02f32.max(inconsistency - 0.2))
                .clamp(1.0 / 256.0, 1.0);
        catch_up_blend = catch_up_blend * vbbr;
        vbbr = vbbr.lerp(variance_based_blend_reduction, 0.1);
        mean = mean.lerp(y, catch_up_blend.saturate());

        self.mean = mean;
        self.short_mean = short_mean;
        self.vbbr = vbbr;
        self.variance = variance;
        self.inconsistency = inconsistency;

        mean
    }
}
