/// Extension methods that perform various matrix operations not included
/// in pyo3, numpy, or ndarray. Additionally this should be seen as a separate
/// mod than the utils as these are are more extentions than utilies.
use crate::utils::Array2R;
use itertools::Itertools;
use ndarray::Array2;
use num::rational::Ratio;
use std::collections::HashSet;
use std::iter::FromIterator;

/// Returns a reflection matrix generated by the vector x;
///
/// # Arguments
/// * `x` - Vector to generate reflection matrix from
pub fn reflection_matrix<'a>(x: &'a Array2R) -> Array2R {
    let n = x.shape()[1];

    Array2::<Ratio<i64>>::eye(n)
        - x.t()
            .dot(x)
            .mapv(|y| y * Ratio::new(2, 1) / (x.dot(&x.t())[[0, 0]]))
}

/// Returns a unique vec of roots rotated by reflection matrices
///
/// # Arguments
/// * `x` - Vector of vector like arrays that are to be reflected
/// * `reflmats` - Vector of reflection matrices.
pub fn reflect_weights<'a>(x: &'a Vec<Array2R>, reflmats: &'a Vec<Array2R>) -> Vec<Array2R> {
    let mut vecs: Vec<Array2R> = Vec::new();

    for i in x.iter() {
        // Tried `par_iter` here, overhead was slower.
        let reflections: Vec<Array2R> = reflmats.iter().map(|x| i.dot(x)).collect();
        vecs.extend(reflections);
    }
    vecs.extend(x.clone());
    vecs.iter().unique().cloned().collect()
}

/// Returns whether or not the given weight is dominant or not. Dominant weights
/// are defined by all coefficents being >= 0.
pub fn all_pos<'a>(x: &'a Array2R) -> bool {
    for &i in x.iter() {
        if i < Ratio::new(0, 1) {
            return false;
        }
    }
    true
}

/// Returns whether or not the given weight is dominant or not. Dominant weights
/// are defined by all coefficents being >= 0.
pub fn all_pos_filter<'a>(x: &'a Array2R, filter: Vec<usize>) -> bool {
    for &i in filter.iter() {
        if x[[0, i]] < Ratio::new(0, 1) {
            return false;
        }
    }
    true
}

///  Returns a set of unique arrays that are all positive after subtraction by `x`
fn select_pos_diff<'a>(x: &'a Array2R, arrays: &'a Vec<Array2R>) -> HashSet<Array2R> {
    HashSet::from_iter(
        arrays
            .iter()
            .map(|y| x - y)
            .filter(|y| all_pos(y))
            .clone()
            .into_iter(),
    )
}

/// Returns the unique set of positive arrays being subtracted by `arrays`
pub fn union_new_weights<'a>(
    x: &'a HashSet<Array2R>,
    arrays: &'a Vec<Array2R>,
) -> HashSet<Array2R> {
    let mut res = HashSet::new();
    for w in x.iter() {
        res = res.union(&select_pos_diff(w, arrays)).cloned().collect();
    }
    res.iter().cloned().collect()
}

#[cfg(test)]
mod test {
    use super::*;
    use ndarray::{array, Array, Dimension};
    use std::collections::HashSet;
    use std::iter::FromIterator;

    fn to_ratio<D>(x: Array<i64, D>) -> Array<Ratio<i64>, D>
    where
        D: Dimension,
    {
        x.mapv(|x| Ratio::new(x, 1))
    }
    #[test]
    fn test_reflection_matrix() {
        let root = to_ratio(array![[1, -1, 0, 0]]);
        let expected = to_ratio(array![
            [0, 1, 0, 0],
            [1, 0, 0, 0],
            [0, 0, 1, 0],
            [0, 0, 0, 1]
        ]);
        assert_eq!(reflection_matrix(&root), expected);
    }

    #[test]
    fn test_select_pos_diff() {
        let weights = vec![
            to_ratio(array![[1, 0, 1]]),
            to_ratio(array![[-1, 1, 1]]),
            to_ratio(array![[1, 1, -1]]),
            to_ratio(array![[-1, 2, -1]]),
        ];

        let result = select_pos_diff(&to_ratio(array![[1, 1, 0]]), &weights);
        let mut expected: HashSet<Array2R> = HashSet::new();
        expected.insert(to_ratio(array![[0, 0, 1]]));
        assert_eq!(result, expected);
    }

    #[test]
    fn test_reflect_weights() {
        let half: num::rational::Ratio<i64> = Ratio::new(1, 2);
        let weights = vec![
            to_ratio(array![[1, -1, 0, 0]]),
            to_ratio(array![[0, 1, -1, 0]]),
        ];

        let reflmats = vec![
            to_ratio(array![
                [0, 1, 0, 0],
                [1, 0, 0, 0],
                [0, 0, 1, 0],
                [0, 0, 0, 1]
            ]),
            to_ratio(array![
                [1, 0, 0, 0],
                [0, 0, 1, 0],
                [0, 1, 0, 0],
                [0, 0, 0, 1]
            ]),
            to_ratio(array![
                [1, 0, 0, 0],
                [0, 1, 0, 0],
                [0, 0, -1, 0],
                [0, 0, 0, 1]
            ]),
            array![
                [half, -half, -half, -half],
                [-half, half, -half, -half],
                [-half, -half, half, -half],
                [-half, -half, -half, half],
            ],
        ];

        let result: HashSet<Array2R> =
            HashSet::from_iter(reflect_weights(&weights, &reflmats).into_iter());
        let expected: HashSet<Array2R> = HashSet::from_iter(
            vec![
                to_ratio(array![[-1, 1, 0, 0]]),
                to_ratio(array![[1, 0, -1, 0]]),
                to_ratio(array![[1, -1, 0, 0]]),
                to_ratio(array![[0, -1, 1, 0]]),
                to_ratio(array![[0, 1, 1, 0]]),
                to_ratio(array![[0, 1, -1, 0]]),
            ]
            .into_iter(),
        );

        assert_eq!(result.difference(&expected).count(), 0);
    }
}
