package algs55; // section 5.5
import stdlib.*;
import algs24.MinPQ;
/* ***********************************************************************
 *  Compilation:  javac Huffman.java
 *  Execution:    java Huffman - < input.txt   (compress)
 *  Execution:    java Huffman + < input.txt   (expand)
 *  Dependencies: BinaryIn.java BinaryOut.java
 *  Data files:   http://algs4.cs.princeton.edu/55compression/abra.txt
 *                http://algs4.cs.princeton.edu/55compression/tinytinyTale.txt
 *
 *  Compress or expand a binary input stream using the Huffman algorithm.
 *
 *  % java Huffman - < abra.txt | java BinaryDump 60
 *  010100000100101000100010010000110100001101010100101010000100
 *  000000000000000000000000000110001111100101101000111110010100
 *  120 bits
 *
 *  % java Huffman - < abra.txt | java Huffman +
 *  ABRACADABRA!
 *
 *************************************************************************/

public class Huffman {
  private static BinaryIn binaryIn;
  private static BinaryOut binaryOut;

  // alphabet size of extended ASCII
  private static final int R = 256;

  // Huffman trie node
  private static class Node implements Comparable<Node> {
    public final char ch;
    public final int freq;
    public final Node left, right;

    public Node(char ch, int freq, Node left, Node right) {
      this.ch    = ch;
      this.freq  = freq;
      this.left  = left;
      this.right = right;
    }

    // is the node a leaf node?
    private boolean isLeaf() {
      assert (left == null && right == null) || (left != null && right != null);
      return (left == null && right == null);
    }

    // compare, based on frequency
    public int compareTo(Node that) {
      return this.freq - that.freq;
    }
  }


  // compress bytes from standard input and write to standard output
  public static void compress() {
    // read the input
    String s = binaryIn.readString();
    char[] input = s.toCharArray();

    // tabulate frequency counts
    int[] freq = new int[R];
    for (int i = 0; i < input.length; i++)
      freq[input[i]]++;

    // build Huffman trie
    Node root = buildTrie(freq);
    writeTrie(root);
    printTrie (root);

    // build code table
    String[] st = new String[R];
    buildCode(st, root, "");

    // print number of bytes in original uncompressed message
    binaryOut.write(input.length);
    StdOut.print ("Length: " + input.length);

    // use Huffman code to encode input
    for (int i = 0; i < input.length; i++) {
      String code = st[input[i]];
      for (int j = 0; j < code.length(); j++) {
        if (code.charAt(j) == '0') {
          binaryOut.write(false);
        }
        else if (code.charAt(j) == '1') {
          binaryOut.write(true);
        }
        else throw new IllegalStateException("Illegal state");
      }
    }

    // close output stream
    binaryOut.close();
  }

  // expand Huffman-encoded input from standard input and write to standard output
  public static void expand() {

    // read in Huffman trie from input stream
    Node root = readTrie();
    printTrie (root);

    // number of bytes to write
    int length = binaryIn.readInt();

    StdOut.print ("Length: " + length);
    // decode using the Huffman trie
    for (int i = 0; i < length; i++) {
      Node x = root;
      while (!x.isLeaf()) {
        boolean bit = binaryIn.readBoolean();
        if (bit) x = x.right;
        else     x = x.left;
      }
      binaryOut.write(x.ch);
    }
    binaryOut.close();
  }

  // build the Huffman trie given frequencies
  private static Node buildTrie(int[] freq) {

    // initialze priority queue with singleton trees
    MinPQ<Node> pq = new MinPQ<>();
    for (char i = 0; i < R; i++)
      if (freq[i] > 0)
        pq.insert(new Node(i, freq[i], null, null));

    // merge two smallest trees
    while (pq.size() > 1) {
      Node left  = pq.delMin();
      Node right = pq.delMin();
      Node parent = new Node('\0', left.freq + right.freq, left, right);
      pq.insert(parent);
    }
    return pq.delMin();
  }


  // make a lookup table from symbols and their encodings
  private static void buildCode(String[] st, Node x, String s) {
    if (!x.isLeaf()) {
      buildCode(st, x.left,  s + '0');
      buildCode(st, x.right, s + '1');
    }
    else {
      st[x.ch] = s;
    }
  }

  // write bitstring-encoded trie to standard output
  private static void writeTrie(Node x) {
    //StdOut.println ("wrote " + x.isLeaf ());
    if (x.isLeaf()) {
      //StdOut.println ("wrote " + x.ch);
      binaryOut.write(true);
      binaryOut.write(x.ch, 8);
    } else {
      binaryOut.write(false);
      writeTrie(x.left);
      writeTrie(x.right);
    }
  }

  private static Node readTrie() {
    boolean isLeaf = binaryIn.readBoolean();
    //StdOut.println ("read " + isLeaf);
    if (isLeaf) {
      char ch = binaryIn.readChar(8);
      //StdOut.println ("read " + ch);
      return new Node(ch, -1, null, null);
    } else {
      Node left = readTrie();
      Node right = readTrie();
      return new Node('\0', -1, left, right);
    }
  }
  private static void printTrie(Node x) { printTrie(x, ""); }
  private static void printTrie(Node x, String pre) {
    if (x.isLeaf()) {
      StdOut.format ("%c %s\n", x.ch, pre);
    }
    if (x.left!=null)  printTrie(x.left, pre + "0");
    if (x.right!=null) printTrie(x.right, pre + "1");
  }

  // seems to be broken...
  public static void main(String[] args) {
    String txtFile = "data/abra.txt";
    String outFile = "/tmp/abra.txt";
    String binFile = "/tmp/abra.bin";

    //args = new String[] { "-" }; BinaryStdIn.fromFile (txtFile); BinaryStdOut.toFile (binFile);
    //args = new String[] { "+" }; BinaryStdIn.fromFile (binFile); BinaryStdOut.toFile (outFile);

    //args = new String[] { "-" }; binaryIn = new BinaryIn(txtFile); binaryOut = new BinaryOut(binFile);
    //args = new String[] { "+" }; binaryIn = new BinaryIn(binFile); binaryOut = new BinaryOut(outFile);

    if      (args[0].equals("-")) compress();
    else if (args[0].equals("+")) expand();
    else throw new Error("Illegal command line argument");
  }

}