use faer::Mat;
use rayon::prelude::*;
use crate::core::math::stats::z_scores_to_pval;
use crate::prelude::*;
#[derive(Clone, Debug)]
pub struct DiffCorRes<T> {
pub r_a: Vec<T>,
pub r_b: Vec<T>,
pub z_score: Vec<T>,
pub p_vals: Vec<T>,
}
pub fn calculate_diff_correlation<T>(
mat_a: &Mat<T>,
mat_b: &Mat<T>,
no_sample_a: usize,
no_sample_b: usize,
spearman: bool,
) -> DiffCorRes<T>
where
T: BixverseFloat,
{
assert_symmetric_mat!(mat_a);
assert_symmetric_mat!(mat_b);
let mut cors_a: Vec<T> = Vec::new();
let mut cors_b: Vec<T> = Vec::new();
let upper_triangle_indices = upper_triangle_indices(mat_a.ncols(), 1);
for (&r, &c) in upper_triangle_indices
.0
.iter()
.zip(upper_triangle_indices.1.iter())
{
cors_a.push(*mat_a.get(r, c));
cors_b.push(*mat_b.get(r, c));
}
let original_cor_a = cors_a.to_vec();
let original_cor_b = cors_b.to_vec();
cors_a.par_iter_mut().for_each(|x| *x = x.atanh());
cors_b.par_iter_mut().for_each(|x| *x = x.atanh());
let constant = if spearman {
T::from_f32(1.06).unwrap()
} else {
T::from_f32(1.0).unwrap()
};
let denominator: T = ((constant
/ (T::from_usize(no_sample_a).unwrap() - T::from_f32(3.0).unwrap()))
+ (constant / (T::from_usize(no_sample_b).unwrap() - T::from_f32(3.0).unwrap())))
.sqrt();
let z_scores: Vec<T> = cors_a
.par_iter()
.zip(cors_b.par_iter())
.map(|(a, b)| (*a - *b) / denominator)
.collect();
let p_values = z_scores_to_pval(&z_scores, "twosided");
DiffCorRes {
r_a: original_cor_a,
r_b: original_cor_b,
z_score: z_scores,
p_vals: p_values,
}
}