use std::collections::BTreeMap;
use std::ops::{Index, Mul};
use num::Zero;

pub struct SparseVector<T> {
    len: usize,
    data: BTreeMap<usize, T>,
    zero: T,
}

impl<T: Zero> SparseVector<T> {
    pub fn insert(&mut self, index: usize, value: T) {
        self.data.insert(index, value);
    }
}

impl<T: Zero> SparseVector<T> {
    pub fn zero(len: usize) -> SparseVector<T> {
        SparseVector {
            len: len,
            data: BTreeMap::new(),
            zero: T::zero(),
        }
    }
}

impl<T: Zero> Index<usize> for SparseVector<T> {
    type Output = T;

    fn index(&self, idx: usize) -> &T {
        match self.data.get(&idx) {
            Some(x) => x,
            None => &self.zero,
        }
    }
}

impl<T: Mul<Output=T> + Clone + Zero> Mul<SparseVector<T>> for SparseVector<T> {
    type Output = SparseVector<T>;

    fn mul(self, rhs: SparseVector<T>) -> SparseVector<T> {
        let data = self.data.into_iter().filter_map(|(key, value)| {
            match rhs.data.get(&key) {
                Some(r) => Some((key, r.clone() * value)),
                None => None,
            }
        }).collect();
        SparseVector {
            len: self.len,
            data: data,
            zero: self.zero,
        }
    }
}

    #[test]
    fn mul_test() {
        let mut lhs = SparseVector::zero(3);
        lhs.insert(1, 3);
        lhs.insert(2, 3);

        let mut rhs = SparseVector::zero(3);
        rhs.insert(1,2);
        rhs.insert(0,1);

        let product = lhs * rhs;
        assert_eq!(product[0], 0);
        assert_eq!(product[1], 6);
        assert_eq!(product[2], 0);
    }