trs-dataframe 0.10.2

use super::ColumnFrame;
use crate::Key;
use data_value::DataValue;
use std::collections::HashMap;
use std::ops::{AddAssign, DivAssign, MulAssign, SubAssign};

macro_rules! op_assign {
    ($item: ident, $op: ident, $oop: tt) => {
        impl $item for ColumnFrame {
            fn $op(&mut self, rhs: Self) {
                if self.data_frame.ncols() != rhs.data_frame.ncols() || self.data_frame.ncols() == 1 {
                    let indexes = self.index.select(rhs.keys());
                    for (right_idx, left_index) in rhs.index.indexes().into_iter().zip(indexes.indexes()) {
                        let right_col = rhs.data_frame.column(right_idx);
                        if right_col.len() == 1 {
                            // fixme this has to be fixed in data_value crate
                            let right_value = &right_col[0];
                            for i in 0..self.data_frame.nrows() {
                                self.data_frame[(i, left_index)].$op(right_value);
                            }
                        } else {
                            // fixme same here
                            for (alpha, rhs) in self
                                .data_frame
                                .column_mut(left_index)
                                .iter_mut()
                                .zip(right_col.iter())
                            {
                                alpha.$op(rhs);
                            }
                        }
                    }
                } else {
                    self.data_frame.zip_mut_with(&rhs.data_frame, |x, y| {
                        let z = x.clone();
                        *x = z $oop y.clone();
                    });
                }
            }
        }

        impl $item<HashMap<String, DataValue>> for ColumnFrame {
            fn $op(&mut self, rhs: HashMap<String, DataValue>) {
                let keys = rhs.keys().cloned().map(Key::from).collect::<Vec<_>>();
                let indexes = self.index.select(keys.as_slice());
                let int_idx =indexes.indexes();
                for left_index in int_idx.iter(){
                    let key_right = self.index.get_key(*left_index).expect("BUG: Defined above!");
                    let right_value = rhs.get(key_right.name()).cloned().unwrap_or_default();
                    for i in 0..self.data_frame.nrows() {
                        self.data_frame[(i, *left_index)].$op(&right_value);
                    }

                }
            }
        }
    };
}

op_assign!(AddAssign, add_assign, +);
op_assign!(SubAssign, sub_assign, -);
op_assign!(MulAssign, mul_assign, *);
op_assign!(DivAssign, div_assign, /);

#[cfg(test)]
mod test {
    use crate::column_frame;

    use super::*;
    use data_value::DataValue;
    use rstest::*;
    #[rstest]
    #[case(
        column_frame!("a" => vec![1, 2, 3]),
        column_frame!("a" => vec![1, 2, 3]),
        column_frame!("a" => vec![2, 4, 6])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        column_frame!("a" => vec![2]),
        column_frame!("a" => vec![3, 4, 5], "b" => vec![1, 2, 3])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        column_frame!("a" => vec![2, 3]),
        column_frame!("a" => vec![3, 5, 3], "b" => vec![1, 2, 3])
    )]
    fn test_add_assign(
        #[case] mut a: ColumnFrame,
        #[case] b: ColumnFrame,
        #[case] expected: ColumnFrame,
    ) {
        a += b;
        assert_eq!(a, expected);
    }
    #[rstest]
    #[case(
        column_frame!("a" => vec![1, 2, 3]),
        column_frame!("a" => vec![1, 2, 3]),
        column_frame!("a" => vec![0f32, 0f32, 0f32])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        column_frame!("a" => vec![2]),
        column_frame!("a" => vec![-1f32, 0f32, 1f32], "b" => vec![1, 2, 3])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        column_frame!("a" => vec![2, 3]),
        column_frame!("a" => vec![DataValue::from(-1f32), DataValue::from(-1f32), DataValue::from(3)], "b" => vec![1, 2, 3])
    )]
    fn test_sub_assign(
        #[case] mut a: ColumnFrame,
        #[case] b: ColumnFrame,
        #[case] expected: ColumnFrame,
    ) {
        a -= b;
        assert_eq!(a, expected);
    }

    #[rstest]
    #[case(
        column_frame!("a" => vec![1, 2, 3]),
        column_frame!("a" => vec![1, 2, 3]),
        column_frame!("a" => vec![1f32, 4f32, 9f32])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        column_frame!("a" => vec![2]),
        column_frame!("a" => vec![2f32, 4f32, 6f32], "b" => vec![1, 2, 3])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        column_frame!("a" => vec![2, 3]),
        column_frame!("a" => vec![DataValue::from(2f32), DataValue::from(6f32), DataValue::from(3)], "b" => vec![1, 2, 3])
    )]
    fn test_mul_assign(
        #[case] mut a: ColumnFrame,
        #[case] b: ColumnFrame,
        #[case] expected: ColumnFrame,
    ) {
        a *= b;
        assert_eq!(a, expected);
    }

    #[rstest]
    #[case(
        column_frame!("a" => vec![1, 2, 3]),
        column_frame!("a" => vec![1, 2, 3]),
        column_frame!("a" => vec![1f32, 1f32, 1f32])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        column_frame!("a" => vec![2]),
        column_frame!("a" => vec![0.5f32, 1f32, (3f32/2f32)], "b" => vec![1, 2, 3])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        column_frame!("a" => vec![2, 3]),
        column_frame!("a" => vec![DataValue::from(0.5f32), DataValue::from(2f32/3f32), DataValue::from(3)], "b" => vec![1, 2, 3])
    )]
    fn test_div_assign(
        #[case] mut a: ColumnFrame,
        #[case] b: ColumnFrame,
        #[case] expected: ColumnFrame,
    ) {
        a /= b;
        assert_eq!(a, expected);
    }

    #[rstest]
    #[case(
        column_frame!("a" => vec![1, 2, 3]),
        crate::data_value::stdhashmap!("a" => 2, "b" => 3),
        column_frame!("a" => vec![2f32, 4f32, 6f32])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        crate::data_value::stdhashmap!("a" => 2),
        column_frame!("a" => vec![2f32, 4f32, 6f32], "b" => vec![1, 2, 3])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        crate::data_value::stdhashmap!("a" => 2, "b" => 3),
        column_frame!("a" => vec![2f32, 4f32, 6f32], "b" => vec![3f32, 6f32, 9f32])
    )]
    fn test_mul_assign_map(
        #[case] mut a: ColumnFrame,
        #[case] b: HashMap<String, DataValue>,
        #[case] expected: ColumnFrame,
    ) {
        a *= b;
        assert_eq!(a, expected);
    }

    #[rstest]
    #[case(
        column_frame!("a" => vec![1, 2, 3]),
        crate::data_value::stdhashmap!("a" => 2, "b" => 3),
        column_frame!("a" => vec![0.5f32, 1f32, 3f32/2f32])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        crate::data_value::stdhashmap!("a" => 2),
        column_frame!("a" => vec![0.5f32, 1f32, 3f32/2f32], "b" => vec![1, 2, 3])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        crate::data_value::stdhashmap!("a" => 2, "b" => 3),
        column_frame!("a" => vec![0.5f32, 1f32, 3f32/2f32], "b" => vec![1f32/3f32, 2f32/3f32, 1f32])
    )]
    fn test_div_assign_map(
        #[case] mut a: ColumnFrame,
        #[case] b: HashMap<String, DataValue>,
        #[case] expected: ColumnFrame,
    ) {
        a /= b;
        assert_eq!(a, expected);
    }

    #[rstest]
    #[case(
        column_frame!("a" => vec![1, 2, 3]),
        crate::data_value::stdhashmap!("a" => 2, "b" => 3),
        column_frame!("a" => vec![3, 4, 5])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        crate::data_value::stdhashmap!("a" => 2),
        column_frame!("a" => vec![3, 4, 5], "b" => vec![1, 2, 3])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        crate::data_value::stdhashmap!("a" => 2, "b" => 3),
        column_frame!("a" => vec![3, 4, 5], "b" => vec![4, 5, 6])
    )]
    fn test_add_assign_map(
        #[case] mut a: ColumnFrame,
        #[case] b: HashMap<String, DataValue>,
        #[case] expected: ColumnFrame,
    ) {
        a += b;
        assert_eq!(a, expected);
    }

    #[rstest]
    #[case(
        column_frame!("a" => vec![1, 2, 3]),
        crate::data_value::stdhashmap!("a" => 1, "b" => 3),
        column_frame!("a" => vec![0f32, 1f32, 2f32])
    )]
    #[case(
        column_frame!("a" => vec![1, 2, 3], "b" => vec![1, 2, 3]),
        crate::data_value::stdhashmap!("a" => 1),
        column_frame!("a" => vec![0f32, 1f32, 2f32], "b" => vec![1, 2, 3])
    )]
    #[case(
        column_frame!("a" => vec![10, 20, 30], "b" => vec![10, 20, 30]),
        crate::data_value::stdhashmap!("a" => 2, "b" => 3),
        column_frame!("a" => vec![8f32, 18f32, 28f32], "b" => vec![7f32, 17f32, 27f32])
    )]
    fn test_sub_assign_map(
        #[case] mut a: ColumnFrame,
        #[case] b: HashMap<String, DataValue>,
        #[case] expected: ColumnFrame,
    ) {
        a -= b;
        assert_eq!(a, expected);
    }
}