r2rs-stats 0.1.1

use std::fmt::Debug;

use nalgebra::DMatrix;
use num_traits::Float;
use r2rs_nmath::{distribution::FBuilder, traits::Distribution};
use strafe_trait::{
    Assumption, Concept, Conclusion, NoSignificantFeature, RejectionStatus, SomeSignficantFeature,
    Statistic, StatisticalTest,
};
use strafe_type::{Alpha64, FloatConstraint, ModelMatrix, Rational64};

use crate::funcs::{regression_sum_of_squares, residual_sum_of_squares};

#[derive(Copy, Clone, Debug)]
pub struct SignificanceOfRegressionTest {
    alpha: Alpha64,
    scale: Option<f64>,
    df: Option<Rational64>,
}

impl StatisticalTest for SignificanceOfRegressionTest {
    type Input = (ModelMatrix, DMatrix<f64>, DMatrix<f64>, DMatrix<f64>);
    type Output = Result<SignificanceOfRegressionStatistic, Box<dyn std::error::Error>>;

    fn assumptions() -> Vec<Box<dyn Assumption>> {
        Vec::new()
    }

    fn null_hypotheses() -> Vec<Box<dyn Conclusion>> {
        vec![Box::new(NoSignificantFeature {})]
    }

    fn alternate_hypotheses() -> Vec<Box<dyn Conclusion>> {
        vec![Box::new(SomeSignficantFeature {})]
    }

    fn test(&mut self, (x, y, b, w): &Self::Input) -> Self::Output {
        let x = x.matrix();
        let mut weighted_x = x.clone();
        weighted_x
            .column_iter_mut()
            .for_each(|mut row| row.component_mul_assign(w));
        let weighted_y = y.component_mul(w);

        let n = x.shape().0;
        let k = x.shape().1 - 1;

        let regression_mean_square =
            regression_sum_of_squares(&weighted_x, &weighted_y, b) / k as f64;

        let residual_mean_square =
            residual_sum_of_squares(&weighted_x, &weighted_y, b) / (n - k - 1) as f64;

        let f = if residual_mean_square == 0.0 {
            f64::infinity()
        } else {
            regression_mean_square / residual_mean_square
        };

        let mut f_distr_builder = FBuilder::new();
        f_distr_builder.with_df1(k);
        f_distr_builder.with_df2(n - k - 1);
        let f_distr = f_distr_builder.build();
        let p = f_distr.probability(f, false).unwrap();

        Ok(SignificanceOfRegressionStatistic {
            f,
            p,
            alpha: self.alpha,
        })
    }
}

#[derive(Copy, Clone, Debug)]
pub struct SignificanceOfRegressionStatistic {
    alpha: Alpha64,
    f: f64,
    p: f64,
}

impl Statistic for SignificanceOfRegressionStatistic {
    fn alpha(&self) -> Alpha64 {
        self.alpha
    }

    fn statistic(&self) -> f64 {
        self.f
    }

    fn probability_value(&self) -> f64 {
        self.p
    }

    fn conclusion(&self) -> RejectionStatus {
        if self.p < self.alpha.unwrap() {
            RejectionStatus::RejectInFavorOf(Box::new(SomeSignficantFeature::new()))
        } else {
            RejectionStatus::FailToReject(Box::new(NoSignificantFeature::new()))
        }
    }
}

impl Default for SignificanceOfRegressionTest {
    fn default() -> Self {
        Self {
            alpha: 0.05.into(),
            scale: None,
            df: None,
        }
    }
}

impl SignificanceOfRegressionTest {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_alpha<A: Into<Alpha64>>(self, alpha: A) -> Self {
        Self {
            alpha: alpha.into(),
            ..self
        }
    }

    pub fn with_scale(self, scale: f64) -> Self {
        Self {
            scale: Some(scale),
            ..self
        }
    }

    pub fn with_df<R: Into<Rational64>>(self, df: R) -> Self {
        Self {
            df: Some(df.into()),
            ..self
        }
    }
}