Documentation
use std::ops::{Sub, SubAssign};

use crate::arr::{
    core::{adaptative_result_shape, dim_expansion_for_shape, plain_index},
    shape_iterator::ShapeIterator,
    vec::core::ArrVec,
};

impl<T> ArrVec<T>
where
    T: Default + Clone + Copy + Sub<Output = T>,
{
    fn sub_to_left(lhs: &mut Self, rhs: &Self) {
        let out_shape = adaptative_result_shape(&lhs.shape, &rhs.shape);
        assert_eq!(lhs.shape, out_shape);
        let rhs_expanded_shape = dim_expansion_for_shape(&rhs.shape, out_shape.len());
        for index in ShapeIterator::new(&out_shape) {
            let out_i = plain_index(&out_shape, &out_shape, &index);
            let lhs_i = plain_index(&out_shape, &out_shape, &index);
            let rhs_i = plain_index(&rhs_expanded_shape, &out_shape, &index);
            lhs.buff[out_i] = lhs.buff[lhs_i] - rhs.buff[rhs_i];
        }
    }

    fn sub_to_new(lhs: &Self, rhs: &Self) -> Self {
        let out_shape = adaptative_result_shape(&lhs.shape, &rhs.shape);
        let mut out = ArrVec::<T>::new(T::default(), out_shape);
        Self::sub_to_out(lhs, rhs, &mut out);
        out
    }

    fn sub_to_out(lhs: &Self, rhs: &Self, out: &mut Self) {
        let lhs_expanded_shape = dim_expansion_for_shape(&lhs.shape, out.shape.len());
        let rhs_expanded_shape = dim_expansion_for_shape(&rhs.shape, out.shape.len());
        for index in ShapeIterator::new(&out.shape) {
            let out_i = plain_index(&out.shape, &out.shape, &index);
            let lhs_i = plain_index(&lhs_expanded_shape, &out.shape, &index);
            let rhs_i = plain_index(&rhs_expanded_shape, &out.shape, &index);
            out.buff[out_i] = lhs.buff[lhs_i] - rhs.buff[rhs_i];
        }
    }
}

impl<T> Sub for &ArrVec<T>
where
    T: Default + Clone + Copy + Sub<Output = T>,
{
    type Output = ArrVec<T>;

    fn sub(self, rhs: Self) -> Self::Output {
        ArrVec::<T>::sub_to_new(self, rhs)
    }
}

impl<T> Sub for ArrVec<T>
where
    T: Default + Clone + Copy + Sub<Output = T>,
{
    type Output = ArrVec<T>;

    fn sub(self, rhs: Self) -> Self::Output {
        ArrVec::<T>::sub_to_new(&self, &rhs)
    }
}

impl<T> SubAssign for ArrVec<T>
where
    T: Default + Clone + Copy + Sub<Output = T>,
{
    fn sub_assign(&mut self, rhs: Self) {
        ArrVec::<T>::sub_to_left(self, &rhs);
    }
}

impl<T> SubAssign<&ArrVec<T>> for ArrVec<T>
where
    T: Default + Clone + Copy + Sub<Output = T>,
{
    fn sub_assign(&mut self, rhs: &Self) {
        ArrVec::<T>::sub_to_left(self, rhs);
    }
}