import java.util.TreeSet;
import java.util.HashMap;
import java.util.ArrayList;
import java.util.PriorityQueue;
import java.util.Iterator;
import java.util.Scanner;

public class HuffmanCode
{
    private static ArrayList<Tree> lists;
    public static void main (String args [])
    {
        Scanner scan = new Scanner (System.in);
        System.out.println("Input line to compress");
        String a = scan.nextLine();
        
        HashMap <String, Integer> map = new HashMap<String, Integer>();
        for (int i = 0; i <a.length(); i++)
        {
            String value = String.valueOf(a.charAt(i));
            if (map.containsKey(value))
            {
                int count = map.get(value);
                map.put (value, count+1);
            }
            else
            {
                map.put(value, 1);
            }
        }
        TreeSet<String> treeSet =  new TreeSet<String> (map.keySet());
        ArrayList<String> list = new ArrayList<String>();
        System.out.println ("Frequency Count");
        for (String s : treeSet)
        {
            System.out.println (s + "   " + map.get(s));
            list.add(s);
        }
        System.out.println();
        Tree hTree = buildHuffmanTree(map);
        hTree.outputTree(hTree.getRoot(), 0);
        HashMap<String, String> codeMap = generateCodes(hTree, list);
        TreeSet<String> wCodes = new TreeSet<String> (codeMap.keySet());
        System.out.println ("Word Codes :");
        for (String s : wCodes)
        {
            System.out.println (s + "   " + codeMap.get(s));
        }
        System.out.println();
        System.out.println ("Huffmann Enconding of : " + a + ":" + encode(a, codeMap));
        System.out.println ("Decoded String = " +  decode(encode(a, codeMap), hTree));
    }

    public static String decode (String input, Tree t)
    {
        return t.getCode(t.getRoot(), input);
    }

    public static String encode (String input, HashMap<String, String> map)
    {
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i <input.length(); i++)
        {
            char c = input.charAt(i);
            String key = String.valueOf(c);
            sb.append (map.get(key));
        }
        return sb.toString();
    }

    public static Tree buildHuffmanTree (HashMap<String, Integer> map)
    {
        TreeSet<String> treeSet =  new TreeSet<String> (map.keySet());
        PriorityQueue<Tree> pQueue = new PriorityQueue<Tree>(); // a binary heap
        for (String s : treeSet)
        {
            Tree t = new Tree();
            t.insertData (map.get(s));
            t.getRoot().setValue(s);
            pQueue.add(t);
        }

        System.out.println("Huffman Forest");
        Iterator <Tree> iterator = pQueue.iterator();
        while (iterator.hasNext())
        {
            Tree t = iterator.next();
            System.out.print (t.getRoot().getValue() + ":"  + t.getRoot().getData() + "   ");
        }
        System.out.println();

        while (pQueue.size() > 1)
        {
            Tree a = pQueue.poll();
            Tree b = pQueue.poll();
            Tree c = new Tree();
            c.insertData (a.getRoot().getData() + b.getRoot().getData());
            c.insertBranch (a.getRoot(), c.getRoot());
            c.insertBranch (b.getRoot(), c.getRoot());
            String value = a.getRoot().getValue() + b.getRoot().getValue();
            c.getRoot().setValue(value);
            pQueue.add(c);
        }
        return pQueue.peek();
    }

    public static HashMap<String, String> generateCodes (Tree t, ArrayList<String> lists)
    {
        HashMap <String, String> Map = new HashMap<String, String>();
        for (int i = 0; i <lists.size(); i++)
        {
            String code = t.codeGen (lists.get(i), t.getRoot());
            Map.put (lists.get(i), code);
        }
        return Map;
    }
}