r2rs-stats 0.1.1

// "Whatever you do, work at it with all your heart, as working for the Lord,
// not for human masters, since you know that you will receive an inheritance
// from the Lord as a reward. It is the Lord Christ you are serving."
// (Col 3:23-24)

use nalgebra::DMatrix;
use num_traits::Float;
use r2rs_base::traits::StatisticalSlice;

#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
pub enum NAMethod {
    #[default]
    Everything,
    AllObs,
    CompleteObs,
    NAOrComplete,
    PairwiseCompleteObs,
}

#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
pub enum Method {
    #[default]
    Pearson,
    Kendall,
    Spearman,
}

fn clamp(x: f64) -> f64 {
    if x > 1.0 {
        1.0
    } else if x < -1.0 {
        -1.0
    } else {
        x
    }
}

fn sign(x: f64) -> f64 {
    if x == 0.0 {
        0.0
    } else {
        x.signum()
    }
}

pub fn sd(x: &[f64]) -> f64 {
    var(
        &DMatrix::from_column_slice(x.len(), 1, x),
        NAMethod::Everything,
    )
    .sqrt()
}

pub fn var(x: &DMatrix<f64>, na_method: NAMethod) -> f64 {
    let x = DMatrix::from_column_slice(x.as_slice().len(), 1, x.as_slice());
    c_cov(&x, Some(&x), na_method, false)[(0, 0)]
}

pub fn cov(
    x: &DMatrix<f64>,
    y: &DMatrix<f64>,
    method: Method,
    na_method: NAMethod,
) -> DMatrix<f64> {
    if let Method::Pearson = method {
        c_cov(x, Some(y), na_method, false)
    } else if na_method == NAMethod::CompleteObs || na_method == NAMethod::NAOrComplete {
        let nas = x
            .row_iter()
            .zip(y.row_iter())
            .enumerate()
            .filter_map(|(i, (x, y))| {
                if x.iter().any(|x_i| x_i.is_nan()) || y.iter().any(|y_i| y_i.is_nan()) {
                    Some(i)
                } else {
                    None
                }
            })
            .collect::<Vec<_>>();

        let mut r_x = x.clone().remove_rows_at(&nas);
        for mut column in r_x.column_iter_mut() {
            for (&rank, x_i) in column.as_slice().rank().iter().zip(column.iter_mut()) {
                *x_i = rank;
            }
        }

        let mut r_y = y.clone().remove_rows_at(&nas);
        for mut column in r_y.column_iter_mut() {
            for (&rank, y_i) in column.as_slice().rank().iter().zip(column.iter_mut()) {
                *y_i = rank;
            }
        }

        c_cov(&r_x, Some(&r_y), na_method, method == Method::Kendall)
    } else if na_method != NAMethod::PairwiseCompleteObs {
        let mut r_x = x.clone();
        for mut column in r_x.column_iter_mut() {
            for (&rank, x_i) in column.as_slice().rank().iter().zip(column.iter_mut()) {
                *x_i = rank;
            }
        }

        let mut r_y = y.clone();
        for mut column in r_y.column_iter_mut() {
            for (&rank, y_i) in column.as_slice().rank().iter().zip(column.iter_mut()) {
                *y_i = rank;
            }
        }

        c_cov(&r_x, Some(&r_y), na_method, method == Method::Kendall)
    } else {
        panic!("Cannot handle Pairwise Complete Obs")
    }
}

pub fn cor(
    x: &DMatrix<f64>,
    y: &DMatrix<f64>,
    na_method: NAMethod,
    method: Method,
) -> DMatrix<f64> {
    if method == Method::Pearson {
        c_cor(x, Some(y), na_method, false)
    } else if na_method == NAMethod::CompleteObs || na_method == NAMethod::NAOrComplete {
        let nas = x
            .row_iter()
            .zip(y.row_iter())
            .enumerate()
            .filter_map(|(i, (x, y))| {
                if x.iter().any(|x_i| x_i.is_nan()) || y.iter().any(|y_i| y_i.is_nan()) {
                    Some(i)
                } else {
                    None
                }
            })
            .collect::<Vec<_>>();

        let mut r_x = x.clone().remove_rows_at(&nas);
        for mut column in r_x.column_iter_mut() {
            for (&rank, x_i) in column.as_slice().rank().iter().zip(column.iter_mut()) {
                *x_i = rank;
            }
        }

        let mut r_y = y.clone().remove_rows_at(&nas);
        for mut column in r_y.column_iter_mut() {
            for (&rank, y_i) in column.as_slice().rank().iter().zip(column.iter_mut()) {
                *y_i = rank;
            }
        }

        c_cor(&r_x, Some(&r_y), na_method, method == Method::Kendall)
    } else if na_method != NAMethod::PairwiseCompleteObs {
        let mut r_x = x.clone();
        for mut column in r_x.column_iter_mut() {
            for (&rank, x_i) in column.as_slice().rank().iter().zip(column.iter_mut()) {
                *x_i = rank;
            }
        }

        let mut r_y = y.clone();
        for mut column in r_y.column_iter_mut() {
            for (&rank, y_i) in column.as_slice().rank().iter().zip(column.iter_mut()) {
                *y_i = rank;
            }
        }

        c_cor(&r_x, Some(&r_y), na_method, method == Method::Kendall)
    } else {
        let mut r = DMatrix::zeros(x.ncols(), y.ncols());
        for i in 0..x.ncols() {
            for j in 0..y.ncols() {
                let mut x2 = x.column(i).iter().cloned().collect::<Vec<_>>();
                let mut y2 = y.column(j).iter().cloned().collect::<Vec<_>>();
                let nas = x2
                    .iter()
                    .zip(x2.iter())
                    .enumerate()
                    .filter_map(|(i, (x_i, y_i))| {
                        if x_i.is_nan() || y_i.is_nan() {
                            Some(i)
                        } else {
                            None
                        }
                    })
                    .collect::<Vec<_>>();
                x2 = x2
                    .into_iter()
                    .enumerate()
                    .filter_map(|(i, x_i)| if nas.contains(&i) { None } else { Some(x_i) })
                    .collect::<Vec<_>>()
                    .rank();
                y2 = y2
                    .into_iter()
                    .enumerate()
                    .filter_map(|(i, y_i)| if nas.contains(&i) { None } else { Some(y_i) })
                    .collect::<Vec<_>>()
                    .rank();
                r[(i, j)] = if !x2.is_empty() && !y2.is_empty() {
                    c_cor(
                        &DMatrix::from_vec(x2.len(), 1, x2),
                        Some(&DMatrix::from_vec(y2.len(), 1, y2)),
                        NAMethod::AllObs,
                        method == Method::Kendall,
                    )[(0, 0)]
                } else {
                    f64::nan()
                };
            }
        }
        r
    }
}

fn c_cov(
    x: &DMatrix<f64>,
    y: Option<&DMatrix<f64>>,
    na_method: NAMethod,
    kendall: bool,
) -> DMatrix<f64> {
    corcov(x, y, na_method, kendall, false)
}

fn c_cor(
    x: &DMatrix<f64>,
    y: Option<&DMatrix<f64>>,
    na_method: NAMethod,
    kendall: bool,
) -> DMatrix<f64> {
    corcov(x, y, na_method, kendall, true)
}

fn corcov(
    x: &DMatrix<f64>,
    y: Option<&DMatrix<f64>>,
    na_method: NAMethod,
    kendall: bool,
    cor: bool,
) -> DMatrix<f64> {
    let mut na_fail = false;
    let mut everything = false;
    let mut empty_err = true;
    let mut pair = false;
    match na_method {
        NAMethod::AllObs => na_fail = true,
        NAMethod::CompleteObs => {}
        NAMethod::PairwiseCompleteObs => pair = true,
        NAMethod::Everything => {
            everything = true;
            empty_err = false;
        }
        NAMethod::NAOrComplete => empty_err = false,
    };

    if let Some(y) = y {
        if everything {
            cov_na_2(x, y, cor, kendall)
        } else if !pair {
            let ind = complete2(x, y, na_fail);
            if empty_err && !ind.iter().any(|&i| i == 1.0) {
                panic!("No complete element pairs");
            }
            cov_complete2(x, y, &ind, cor, kendall)
        } else {
            cov_pairwise2(x, y, cor, kendall)
        }
    } else if everything {
        cov_na_2(x, x, cor, kendall)
    } else if !pair {
        let ind = complete2(x, x, na_fail);
        if empty_err && !ind.iter().any(|&i| i == 1.0) {
            panic!("No complete element pairs");
        }
        cov_complete2(x, x, &ind, cor, kendall)
    } else {
        cov_pairwise2(x, x, cor, kendall)
    }
}

fn mean_2(x: &DMatrix<f64>, has_na: &[bool]) -> Vec<f64> {
    let mut xm = vec![0.0; x.nrows()];
    for i in 0..x.ncols() {
        let mut tmp;
        if has_na[i] {
            tmp = f64::nan();
        } else {
            let xx = &x.column(i);
            let mut sum = 0.0;
            for k in 0..x.nrows() {
                sum += xx[k];
            }
            tmp = sum / x.nrows() as f64;
            if tmp.is_finite() {
                sum = 0.;
                for k in 0..x.nrows() {
                    sum += xx[k] - tmp;
                }
                tmp += sum / x.nrows() as f64;
            }
        }
        xm[i] = tmp;
    }
    xm
}

fn cov_na_2(x: &DMatrix<f64>, y: &DMatrix<f64>, cor: bool, kendall: bool) -> DMatrix<f64> {
    let mut ans = DMatrix::repeat(x.ncols(), y.ncols(), f64::nan());

    let has_na_x = x
        .column_iter()
        .map(|c| c.iter().any(|c_i| c_i.is_nan()))
        .collect::<Vec<_>>();
    let has_na_y = y
        .column_iter()
        .map(|c| c.iter().any(|c_i| c_i.is_nan()))
        .collect::<Vec<_>>();

    let n = x.nrows();
    let n1 = n - 1;

    let mut xm = if kendall {
        vec![0.0; x.ncols()]
    } else {
        mean_2(x, &has_na_x)
    };
    let mut ym = if kendall {
        vec![0.0; x.ncols()]
    } else {
        mean_2(y, &has_na_y)
    };

    for i in 0..x.ncols() {
        if has_na_x[i] {
            for j in 0..y.ncols() {
                ans[(i, j)] = f64::nan();
            }
        } else {
            let xx = x.column(i);
            if !kendall {
                let xxm = xm[i];
                for j in 0..y.ncols() {
                    if has_na_y[j] {
                        ans[(i, j)] = f64::nan();
                    } else {
                        let yy = y.column(j);
                        let yym = ym[j];
                        let mut sum = 0.0;
                        for k in 0..n {
                            sum += (xx[k] - xxm) * (yy[k] - yym);
                        }
                        ans[(i, j)] = sum / n1 as f64;
                    }
                }
            } else {
                for j in 0..y.ncols() {
                    if has_na_y[j] {
                        ans[(i, j)] = f64::nan();
                    } else {
                        let yy = y.column(j);
                        let mut sum = 0.0;
                        for k in 0..n {
                            for n1 in 0..n {
                                sum += sign(xx[k] - xx[n1]) * sign(yy[k] - yy[n1]);
                            }
                        }
                        ans[(i, j)] = sum;
                    }
                }
            }
        }
    }

    if cor {
        let cov_sdev = |x: &DMatrix<f64>, has_na_x: &[bool], xm: &mut [f64]| {
            for i in 0..x.ncols() {
                if !has_na_x[i] {
                    let xx = x.column(i);
                    let mut sum = 0.0;
                    if !kendall {
                        let xxm = xm[i];
                        for k in 0..n {
                            sum += (xx[k] - xxm) * (xx[k] - xxm);
                        }
                        sum /= n1 as f64;
                    } else {
                        for k in 0..n {
                            for n1 in 0..n {
                                if xx[k] != xx[n1] {
                                    sum += 1.0;
                                }
                            }
                        }
                    }
                    xm[i] = sum.sqrt();
                }
            }
        };

        cov_sdev(x, &has_na_x, &mut xm);
        cov_sdev(y, &has_na_y, &mut ym);

        let mut _sd_0 = false;
        for i in 0..x.ncols() {
            if !has_na_x[i] {
                for j in 0..y.ncols() {
                    if !has_na_y[j] {
                        if xm[i] == 0.0 || ym[j] == 0.0 {
                            _sd_0 = true;
                            ans[(i, j)] = f64::nan();
                        } else {
                            ans[(i, j)] /= xm[i] * ym[j];
                            ans[(i, j)] = clamp(ans[(i, j)]);
                        }
                    }
                }
            }
        }
    }

    ans
}

fn complete2(x: &DMatrix<f64>, y: &DMatrix<f64>, na_fail: bool) -> Vec<f64> {
    let mut ind = vec![1.0; x.nrows()];
    for j in 0..x.ncols() {
        let z = x.column(j);
        for i in 0..x.nrows() {
            if z[i].is_nan() {
                if na_fail {
                    panic!("Missing observations in cov/cor")
                } else {
                    ind[i] = 0.0;
                }
            }
        }
    }
    for j in 0..y.ncols() {
        let z = y.column(j);
        for i in 0..x.nrows() {
            if z[i].is_nan() {
                if na_fail {
                    panic!("Missing observations in cov/cor")
                } else {
                    ind[i] = 0.0;
                }
            }
        }
    }
    ind
}

fn cov_complete2(
    x: &DMatrix<f64>,
    y: &DMatrix<f64>,
    ind: &[f64],
    cor: bool,
    kendall: bool,
) -> DMatrix<f64> {
    let mut ans = DMatrix::repeat(x.ncols(), y.ncols(), f64::nan());

    let n = x.nrows();
    let n1 = n - 1;
    let mut xm = x
        .column_iter()
        .map(|c| c.as_slice().mean())
        .collect::<Vec<_>>();
    let mut ym = y
        .column_iter()
        .map(|c| c.as_slice().mean())
        .collect::<Vec<_>>();

    for i in 0..x.ncols() {
        let xx = x.column(i);
        let xxm = xm[i];
        for j in 0..y.ncols() {
            let yy = y.column(j);
            let mut sum = 0.0;
            if !kendall {
                let yym = ym[j];
                for k in 0..n {
                    if ind[k] != 0.0 {
                        sum += (xx[k] - xxm) * (yy[k] - yym);
                    }
                }
                ans[(i, j)] = sum / n1 as f64;
            } else {
                for k in 0..n {
                    if ind[k] != 0.0 {
                        for n1 in 0..n {
                            if ind[n1] != 0.0 {
                                sum += sign(xx[k] - xx[n1]) * sign(yy[k] - yy[n1]);
                            }
                        }
                    }
                }
                ans[(i, j)] = sum;
            }
        }
    }

    if cor {
        let cov_sdev = |x: &DMatrix<f64>, xm: &mut [f64], ind: &[f64]| {
            for (i, xx) in x.column_iter().enumerate() {
                let mut sum = 0.0;
                if !kendall {
                    let xxm = xm[i];
                    for k in 0..n {
                        if ind[k] != 0.0 {
                            sum += (xx[k] - xxm) * (xx[k] - xxm);
                        }
                    }
                    sum /= n1 as f64;
                } else {
                    for k in 0..n {
                        if ind[k] != 0.0 {
                            for n1 in 0..n {
                                if ind[n1] != 0.0 && xx[k] != xx[n1] {
                                    sum += 1.0;
                                }
                            }
                        }
                    }
                }
                xm[i] = sum.sqrt();
            }
        };

        cov_sdev(x, &mut xm, ind);
        cov_sdev(y, &mut ym, ind);

        let mut _sd_0 = false;
        for i in 0..x.ncols() {
            for j in 0..y.ncols() {
                if xm[i] == 0.0 || ym[j] == 0.0 {
                    _sd_0 = true;
                    ans[(i, j)] = f64::nan();
                } else {
                    ans[(i, j)] /= xm[i] * ym[j];
                    ans[(i, j)] = clamp(ans[(i, j)]);
                }
            }
        }
    }

    ans
}

fn cov_pairwise2(x: &DMatrix<f64>, y: &DMatrix<f64>, cor: bool, kendall: bool) -> DMatrix<f64> {
    let mut ans = DMatrix::repeat(x.ncols(), y.ncols(), f64::nan());
    let mut _sd_0 = false;

    let n = x.nrows();
    let mut n1 = n - 1;

    for i in 0..x.ncols() {
        let xx = x.column(i);
        for j in 0..y.ncols() {
            let yy = y.column(j);
            let mut xmean = 0.0;
            let mut ymean = 0.0;

            let mut nobs = 0;
            for k in 0..n {
                if !xx[k].is_nan() && !yy[k].is_nan() {
                    nobs += 1;
                    if !kendall {
                        xmean += xx[k];
                        ymean += yy[k];
                    }
                }
            }

            if nobs >= 2 {
                let mut xsd = 0.0;
                let mut ysd = 0.0;
                let mut sum = 0.0;

                if !kendall {
                    xmean /= nobs as f64;
                    ymean /= nobs as f64;
                    n1 = nobs - 1;
                }

                for k in 0..n {
                    if !xx[k].is_nan() && !yy[k].is_nan() {
                        if !kendall {
                            let xm = xx[k] - xmean;
                            let ym = yy[k] - ymean;

                            sum += xm * ym;
                            if cor {
                                xsd += xm * xm;
                                ysd += ym * ym;
                            }
                        } else {
                            for n1 in 0..k {
                                if !xx[n1].is_nan() && !yy[n1].is_nan() {
                                    let xm = sign(xx[k] - xx[n1]);
                                    let ym = sign(yy[k] - yy[n1]);

                                    sum += xm * ym;
                                    if cor {
                                        xsd += xm * xm;
                                        ysd += ym * ym;
                                    }
                                }
                            }
                        }
                    }
                }

                if cor {
                    if xsd == 0.0 || ysd == 0.0 {
                        _sd_0 = true;
                        sum = f64::nan();
                    } else {
                        if !kendall {
                            xsd /= n1 as f64;
                            ysd /= n1 as f64;
                            sum /= n1 as f64;
                        }
                        sum /= xsd.sqrt() * ysd.sqrt();
                        sum = clamp(sum);
                    }
                } else if !kendall {
                    sum /= n1 as f64;
                }

                ans[(i, j)] = sum;
            } else {
                ans[(i, j)] = f64::nan();
            }
        }
    }

    ans
}

pub fn cov2cor(v: &DMatrix<f64>) -> DMatrix<f64> {
    let is = v
        .diagonal()
        .iter()
        .map(|i| (1.0 / i).sqrt())
        .collect::<Vec<_>>();

    let mut ret = v.clone();
    ret.column_iter_mut()
        .for_each(|mut c| c.iter_mut().zip(is.iter()).for_each(|(c_i, i)| *c_i *= i));
    ret.row_iter_mut()
        .for_each(|mut r| r.iter_mut().zip(is.iter()).for_each(|(r_i, i)| *r_i *= i));

    ret
}