extern crate itertools;
use itertools::Itertools;
use std::cmp::Eq;
use std::cmp::Ordering;
use std::collections::binary_heap::BinaryHeap;

static MESSAGE: &'static str = r#"
  Hello! What's going on? This is a program.
  The brown fox jumped over the lazy dog.
  xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
  \o/
"#;

#[derive(Debug, Eq, PartialEq)]
enum Tree<T> {
    Leaf(i32, T),
    Branch(i32, Box<Tree<T>>, Box<Tree<T>>),
}

#[derive(Debug, Eq, PartialEq)]
enum Direction { Left, Right }

fn cost<T>(t: &Tree<T>) -> i32 {
    match *t {
        Tree::Leaf(cost, _) => cost,
        Tree::Branch(cost, _, _) => cost,
    }
}

impl<T: Eq> Ord for Tree<T> {
    fn cmp(&self, other: &Self) -> Ordering {
        cost(other).cmp(&cost(self))
    }
}

impl<T: Eq> PartialOrd for Tree<T> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

fn make_path<T: Eq>(tree: &Tree<T>, elem: &T) -> Vec<Direction> {
    fn find<T: Eq>(n: &Tree<T>, elem: &T) -> Option<Vec<Direction>> {
        match *n {
            Tree::Leaf(_, ref e) =>
                if e == elem { Some(vec!()) } else { None },
            Tree::Branch(_, ref left, ref right) =>
                find(left.as_ref(), elem)
                .map(|mut path| { path.push(Direction::Left); path })

                .or(find(right.as_ref(), elem)
                    .map(|mut path| { path.push(Direction::Right); path})),
        }
    };
    let mut result = find(tree, elem).unwrap();
    result.reverse();
    result
}

fn find_char<T: Eq+Copy>(tree: &Tree<T>, path: &[Direction]) -> Option<T> {
    match (tree, path.split_first()) {
        (&Tree::Branch(_, ref left, _), Some((&Direction::Left, ds))) =>
            find_char(left.as_ref(), ds),
        (&Tree::Branch(_, _, ref right), Some((&Direction::Right, ds))) =>
            find_char(right.as_ref(), ds),
        (&Tree::Leaf(_, ref c), None) => Some(*c),
        _ => None
    }
}

fn main() {
    if MESSAGE.is_empty() {
        panic!("Empty string");
    }

    let mut counts = MESSAGE
        .chars()
        .sorted().into_iter()
        .group_by(|x| *x)
        .map(|(elem, xs)| (elem, xs.len() as i32))
        .map(|(elem, cost)| Box::new(Tree::Leaf(cost, elem)))
        .collect::<BinaryHeap<_>>();

    let mut left;
    loop {
        left = counts.pop().unwrap();
        let right = match counts.pop() {
            Some(right) => right,
            None => break
        };

        counts.push(Box::new(
            Tree::Branch(cost(left.as_ref()) + cost(right.as_ref()),
                         left, right)));
    }

    let restored_message = MESSAGE
        .chars()
        .map(|c| make_path(left.as_ref(), &c))
        .flat_map(|p| find_char(left.as_ref(), &p))
        .collect::<String>();

    println!("{}", MESSAGE);
    println!("{}", restored_message);
}