extern crate itertools;
use itertools::Itertools;
use itertools::Unfold;
use std::cmp::Eq;
use std::cmp::Ordering;
use std::fmt::Debug;
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
  æøå
"#;

#[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 }

impl ToString for Direction {
    fn to_string(&self) -> String {
        match *self {
            Direction::Left => "0".to_string(),
            Direction::Right => "1".to_string()
        }
    }
}

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_char2<'a, T: Debug+Eq+Copy, I>(tree: &Tree<T>, path: &mut I) -> Option<T>
    where I: Iterator<Item=&'a Direction> {
    match (tree) {
        &Tree::Leaf(_, ref c) => Some(*c),
        &Tree::Branch(_, ref left, ref right) =>
            match (path.next()) {
                Some(&Direction::Left) =>
                    find_char2(left.as_ref(), path),
                Some(&Direction::Right) =>
                    find_char2(right.as_ref(), path),
                _ => None
            },
    }
}

fn decode<'a, T: Debug+Eq+Copy, I>(tree: &Tree<T>, path: &'a mut I)
    -> Vec<T>
    where I: Iterator<Item=&'a Direction> {
    Unfold::new((tree, path),
                |&mut (ref tree, ref mut path)| {
                    find_char2(tree, path)
                }).collect::<Vec<_>>()
}

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)| Box::new(Tree::Leaf(xs.len() as i32, 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 code = MESSAGE
        .chars()
        .map(|c| make_path(left.as_ref(), &c))
        .collect::<Vec<_>>();

    println!("{}", code.iter()
             .flatten().map(|d| d.to_string())
             .collect::<String>());

    let restored_message0 =
        decode(left.as_ref(), &mut code.iter().flatten())
         .into_iter().collect::<String>();

    let restored_message = code.iter()
        .flat_map(|p| find_char2(left.as_ref(), &mut p.iter()))
        .collect::<String>();

    println!("{}", MESSAGE);
    println!("Restored message");
    println!("{}", restored_message0);
}