use ndarray::Array1;

/// Functions to apply to residuals for robust regression.
///
/// These functions are applied to the resiudals as part of the `cost` function:
/// $\text{cost}(r) = \sqrt{f^2 \rho(z)}$,
/// where $r$ is the residual, $\rho$ is the loss function,
/// $f$ is the scaling factor, and $z = \frac{r^2}{f^2}$.
#[derive(Clone, Debug, Copy)]
pub enum Loss {
    /// Linear: $\rho(z) = z$
    Linear,
    /// Soft L1: $\rho(z) = 2 (\sqrt{1 + z} - 1)$
    SoftL1(f64),
    /// Huber: $\rho(z) = z$ if $z <= 1$, else $2 \sqrt{z} - 1$
    Huber(f64),
    /// Cauchy: $\rho(z) = \ln(1 + z)$
    Cauchy(f64),
    /// Arctan: $\rho(z) = \arctan(z)$
    Arctan(f64),
}

impl Loss {
    pub fn softl1(scaling_factor: f64) -> Self {
        Self::SoftL1(scaling_factor)
    }
    pub fn huber(scaling_factor: f64) -> Self {
        Self::Huber(scaling_factor)
    }
    pub fn cauchy(scaling_factor: f64) -> Self {
        Self::Cauchy(scaling_factor)
    }
    pub fn arctan(scaling_factor: f64) -> Self {
        Self::Arctan(scaling_factor)
    }

    /// Apply function to array of residuals.
    pub fn apply(&self, res: &mut Array1<f64>) {
        match self {
            Self::Linear => (),
            Self::SoftL1(s) => {
                let s2 = s * s;
                let s2_inv = 1.0 / s2;
                res.mapv_inplace(|ri| {
                    (s2 * (2.0 * (((ri * ri * s2_inv) + 1.0).sqrt() - 1.0))).sqrt()
                })
            }
            Self::Huber(s) => {
                let s2 = s * s;
                let s2_inv = 1.0 / s2;
                res.mapv_inplace(|ri| {
                    if ri * ri * s2_inv <= 1.0 {
                        ri
                    } else {
                        (s2 * (2.0 * (ri / s).abs() - 1.0)).sqrt()
                    }
                })
            }
            Self::Cauchy(s) => {
                let s2 = s * s;
                let s2_inv = 1.0 / s2;
                res.mapv_inplace(|ri| (s2 * (1.0 + (ri * ri * s2_inv)).ln()).sqrt())
            }
            Self::Arctan(s) => {
                let s2 = s * s;
                let s2_inv = 1.0 / s2;
                res.mapv_inplace(|ri| (s2 * (ri * ri * s2_inv).atan()).sqrt())
            }
        }
    }
}
