use std::collections::btree_map::{BTreeMap, Entry};
use std::ops::{Add, Sub};

type Variable = usize;
type Coefficient = isize;

#[derive(Clone)]
struct LinComb(Vec<(Variable, Coefficient)>);
struct CanonicalLinComb(BTreeMap<Variable, Coefficient>);

impl Add for LinComb {
    type Output = LinComb;

    fn add(self, other: LinComb) -> Self::Output {
        LinComb(self.0.into_iter().chain(other.0.into_iter()).collect())
    }
}

impl Sub for LinComb {
    type Output = LinComb;

    fn sub(self, other: LinComb) -> Self::Output {
        LinComb(
            self.0
                .into_iter()
                .chain(other.0.into_iter().map(|(var, coeff)| (var, -coeff)))
                .collect(),
        )
    }
}

impl LinComb {
    fn as_canonical<'a>(&'a self) -> CanonicalLinComb {
        CanonicalLinComb(self.0.clone().into_iter().fold(
            BTreeMap::new(),
            |mut acc, (val, coeff)| {
                // if we're adding 0 times some variable, we can ignore this term
                if coeff != 0 {
                    match acc.entry(val) {
                        Entry::Occupied(o) => {
                            // if the new value is non zero, update, else remove the term entirely
                            if o.get() + coeff != 0 {
                                *o.into_mut() += coeff;
                            } else {
                                o.remove();
                            }
                        }
                        Entry::Vacant(v) => {
                            // We checked earlier but let's make sure we're not creating zero-coeff terms
                            assert!(coeff != 0);
                            v.insert(coeff);
                        }
                    }
                }

                acc
            },
        ))
    }
}

fn main() {
    let a = LinComb(vec![(0, 1), (1, 1)]);
    let b = LinComb(vec![(0, 1), (1, -1)]);
    for i in (a + b).as_canonical().0.iter() {
        println!("{:?}", i);
    }

    let a = LinComb(vec![(0, 1), (1, 1)]);
    let b = LinComb(vec![(0, 1), (1, -1)]);
    for i in (a - b).as_canonical().0.iter() {
        println!("{:?}", i);
    }
}