// Courtesy of Jean-Christophe Filliatre (INF411).
// With slight changes by Francois Pottier.

import java.io.IOException;
import java.io.StringReader;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;
import java.util.PriorityQueue;

// ----------------------------------------------------------------------------

// The parent class. A tree is either a leaf or a binary node.

abstract class HuffmanTree implements Comparable<HuffmanTree> {

  // The cumulated frequency of this (sub-)tree.
  int freq;

  // A total ordering, based on frequency.
  @Override
  public int compareTo(HuffmanTree that) {
    return this.freq - that.freq;
  }

  // Assume s is a string of 0's and 1's.
  // Consider the suffix of s that begins at offset i.
  // This suffix defines a path in this tree.
  // find(s, i) follows this path until it reaches a leaf
  // and returns the character found at that leaf.
  abstract char find(String s, int i);

  // traverse is used to turn a tree into a dictionary, which maps
  // every character to its encoding.
  abstract void traverse(String prefix, Map<Character, String> m);

  // encode is used to turn a tree into a string.
  abstract void encode(StringBuilder sb);
}

// ----------------------------------------------------------------------------

// The Leaf class.

class Leaf extends HuffmanTree {
  final char c;
  Leaf(char c) {
    this.c = c;
    this.freq = 0;
  }
  @Override
  void traverse(String prefix, Map<Character, String> m) {
    m.put(this.c, prefix);
  }
  @Override
  char find(String s, int i) {
    return this.c;
  }
  @Override
  void encode(StringBuilder sb) {
    sb.append('0');
    sb.append(this.c);
  }
}

// ----------------------------------------------------------------------------

// The Node class.

class Node extends HuffmanTree {
  HuffmanTree left, right;
  Node(HuffmanTree left, HuffmanTree right) {
    this.left = left;
    this.right = right;
    this.freq = left.freq + right.freq;
  }
  @Override
  void traverse(String prefix, Map<Character, String> m) {
    this.left.traverse(prefix + '0', m);
    this.right.traverse(prefix + '1', m);
  }
  @Override
  char find(String s, int i) {
    assert (0 <= i && i < s.length());
    assert (s.charAt(i) == '0' || s.charAt(i) == '1');
    return (s.charAt(i) == '0' ? this.left : this.right).find(s, i+1);
  }
  @Override
  void encode(StringBuilder sb) {
    sb.append('1');
    this.left.encode(sb);
    this.right.encode(sb);
  }
}

// ----------------------------------------------------------------------------

// Construction of the tree (for a given alphabet). Encoding, decoding.

public class Huffman {
  
  private HuffmanTree tree;
  private Map<Character, String> codes;
  
  Huffman(Collection<Leaf> alphabet) {
    // The alphabet must have size at least 2, so that the Huffman tree
    // has at least two leaves, and every path down to a leaf has length
    // at least 1. Thus, every symbol is encoded as a nonempty bit string.
    if (alphabet.size() <= 1) throw new IllegalArgumentException();
    this.tree = buildTree(alphabet);
    this.codes = new HashMap<Character, String>();
    this.tree.traverse("", this.codes);
  }
  
  // Out of a collection of leaves, each of which represents a
  // symbol and its frequency, construct an optimal Huffman tree.
  HuffmanTree buildTree(Collection<Leaf> alphabet) {
    PriorityQueue<HuffmanTree> q = new PriorityQueue<HuffmanTree>();
    q.addAll(alphabet);
    while (q.size() > 1) {
      HuffmanTree left = q.poll();
      HuffmanTree right = q.poll();
      q.add(new Node(left, right));
    }
    return q.poll();
  }
  
  String encode(String msg) {
    StringBuilder sb = new StringBuilder();
    for (int i = 0; i < msg.length(); i++)
      sb.append(this.codes.get(msg.charAt(i)));
    return sb.toString();
  }
  
  String decode(String msg) {
    StringBuilder sb = new StringBuilder();
    int i = 0;
    while (i < msg.length()) {
      char c = this.tree.find(msg, i);
      sb.append(c);
      i += this.codes.get(c).length();
    }
    return sb.toString();
  }
  
  public String toString() {
    StringBuilder sb = new StringBuilder();
    for (Entry<Character, String> e: this.codes.entrySet())
      sb.append("[" + e.getKey() + " " + e.getValue() + "]");
    return sb.toString();
  }
   
  public String encodeTree() {
    StringBuilder sb = new StringBuilder();
    this.tree.encode(sb);
    return sb.toString();
  }
  
  private HuffmanTree decodeTree(StringReader r) throws IOException {
    int c = r.read();
    if (c == -1) throw new IllegalArgumentException("bad tree");
    if (c == '0') return new Leaf((char)r.read());
    HuffmanTree left = decodeTree(r);
    HuffmanTree right = decodeTree(r);
    return new Node(left, right);
  }
  
  public void decodeTree(String s) throws IOException {
    this.tree = decodeTree(new StringReader(s));
    this.codes = new HashMap<Character, String>();
    this.tree.traverse("", this.codes);
  }
  
}

// ----------------------------------------------------------------------------

// Testing.

class HuffmanTest {
  
  static Collection<Leaf> buildAlphabet(String text) {
    HashMap<Character, Leaf> index = new HashMap<Character, Leaf>();
    for (int i = 0; i < text.length(); i++) {
      Character c = text.charAt(i);
      Leaf l = index.get(c);
      if (l == null) { l = new Leaf(c); index.put(c, l); }
      l.freq++;
    }
    return index.values();
  }

  static void printAlphabet(Collection<Leaf> alphabet) {
    for (Leaf l : alphabet)
      System.out.print("[" + l.c + " " + l.freq + "]");
    System.out.println();
  }

  static Huffman buildHuffman(String text) {
    System.out.println("Text:");
    System.out.println(text);
    Collection<Leaf> alphabet = buildAlphabet(text);
    System.out.println("Alphabet:");
    printAlphabet(alphabet);
    System.out.println("Huffman encoding:");
    Huffman h = new Huffman(alphabet);
    System.out.println(h);
    return h;
  }

  static void test (Huffman h, String message) {
    System.out.println("Test:");
    System.out.println("  message: " + message);
    System.out.println("    (" + 8 * message.length() + " bits, in ASCII)");
    String encoded = h.encode(message);
    System.out.println("  encoded: " + encoded);
    System.out.println("    (" + encoded.length() + " bits)");
    String decoded = h.decode(encoded);
    System.out.println("  decoded: " + decoded);
    assert (decoded.equals(message));
    System.out.println("  OK");
  }
  
  public static void main(String[] args) {
    String text;
    Huffman h;

    text = "les bases de la programmation et de l'algorithmique";
    h = buildHuffman(text);
    test(h, "bases");
    test(h, "goal");
    test(h, "");
    test(h, text);

    text = "Mississippi";
    h = buildHuffman(text);
    test(h, text);
    test(h, "sissi");

    text = "ab";
    h = buildHuffman(text);
    test(h, "aaabbb");
    test(h, "baba");
  }
}