use ndarray::{Array1, Array2, Axis};
use robust_rs_core::error::RobustError;
use super::chi2::{chi2_cdf, chi2_quantile};
use super::distances_from;
use super::linalg::mean_covariance;
use crate::util::median;
pub(crate) fn consistency_factor(alpha: f64, p: f64) -> f64 {
if alpha >= 1.0 {
return 1.0;
}
let q = chi2_quantile(alpha, p);
let denom = chi2_cdf(q, p + 2.0);
if denom > 0.0 {
alpha / denom
} else {
1.0
}
}
pub(crate) struct Reweighted {
pub location: Array1<f64>,
pub scatter: Array2<f64>,
pub distances: Array1<f64>,
pub weights: Array1<f64>,
}
pub(crate) fn hard_reweight(
x: &Array2<f64>,
raw_location: &Array1<f64>,
raw_scatter: &Array2<f64>,
quantile: f64,
adjusted: bool,
) -> Result<Option<Reweighted>, RobustError> {
let (n, p) = x.dim();
let raw_dist = distances_from(x, raw_location, raw_scatter)?;
let cut2 = if adjusted {
let mut d2: Vec<f64> = raw_dist.iter().map(|&d| d * d).collect();
let med = median(&mut d2);
med * chi2_quantile(quantile, p as f64) / chi2_quantile(0.5, p as f64)
} else {
chi2_quantile(quantile, p as f64)
};
let mut weights = Array1::<f64>::zeros(n);
let mut keep: Vec<usize> = Vec::with_capacity(n);
for i in 0..n {
if raw_dist[i] * raw_dist[i] <= cut2 {
weights[i] = 1.0;
keep.push(i);
}
}
if keep.len() < p + 1 {
return Ok(None);
}
let xs = x.select(Axis(0), &keep);
let (location, cov_raw) = mean_covariance(&xs);
let alpha = keep.len() as f64 / n as f64;
let scatter = &cov_raw * consistency_factor(alpha, p as f64);
let distances = distances_from(x, &location, &scatter)?;
Ok(Some(Reweighted {
location,
scatter,
distances,
weights,
}))
}