package huffman;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map.Entry;
import java.util.PriorityQueue;
import java.util.Queue;

public class Main {
    public static void main(String[] args) throws Exception {
        File f = new File("illiad.txt");
        HuffmanEncoder he = new HuffmanEncoder();
        String orig_text = new String(HuffmanEncoder.readFile(f));

        System.out.println("Getting Frequencies...");
        String freq = he.getFrequencies(f);

        System.out.println("Building Tree...");
        HuffTree ht = he.buildTree(f);

        System.out.println("Encoding File...");
        String encoded = he.encodeFile(f, ht);

        System.out.println("Decoding File...");
        String decoded = he.decodeFile(encoded, ht);

        System.out.println("Getting Dict File...");
        String dict = he.traverseHuffmanTree(ht);

        System.out.println(freq);
        System.out.println(dict);
        System.out.println(orig_text.equals(decoded));
    }
}

class HuffmanEncoder implements HuffmanCoding {
    public String getFrequencies(File inputFile) throws FileNotFoundException {
        HashMap<Character, Integer> char_table = new HashMap<Character, Integer>();
        byte[] fileContent = readFile(inputFile);
        for (byte b : fileContent)
        {
            if (!char_table.containsKey((char)b))
                char_table.put((char)b, 0);
            char_table.replace((char)b, char_table.get((char)b) + 1);
        }

        String ret = "";
        for (Entry<Character, Integer> e : char_table.entrySet())
            ret += e.getKey() + " " + e.getValue() + "\n";
        return ret;
    }

    public HuffTree buildTree(File inputFile) throws FileNotFoundException, Exception {
        HashMap<Character, Integer> char_table = new HashMap<Character, Integer>();
        byte[] fileContent = readFile(inputFile);
        for (byte b : fileContent) {
            if (!char_table.containsKey((char)b))
                char_table.put((char)b, 0);
            char_table.replace((char)b, char_table.get((char)b) + 1);
        }
        HuffTree tree = new HuffTree(char_table);
        return tree;
    }

    public String encodeFile(File inputFile, HuffTree huffTree) throws FileNotFoundException {
        byte[] fileContent = readFile(inputFile);
        byte[] output = new byte[fileContent.length * Byte.MAX_VALUE];
        String en;
        int j = 0;
        for (int i = 0; i < fileContent.length; i++)
        {
            en = huffTree.translate((char)fileContent[i]);
            for (byte c : en.getBytes())
                output[j++] = c;
        }
        byte[] out_string = new byte[j];
        System.arraycopy(output, 0, out_string, 0, j);
        return new String(out_string);
    }

    public String decodeFile(String code, HuffTree huffTree) throws Exception {
        byte[] output = new byte[code.length()];
        String curr = "";
        int j = 0;
        for (byte b : code.getBytes()) {
            curr += (char)b;
            if (huffTree.translate(curr) != null) {
                output[j++] = (byte)huffTree.translate(curr).charValue();
                curr = "";
            }
        }
        byte[] out_string = new byte[j];
        System.arraycopy(output, 0, out_string, 0, j);
        return new String(out_string);
    }

    public String traverseHuffmanTree(HuffTree huffTree) throws Exception {
        String ret = "";
        HashMap<Character, String> dict = huffTree.getDict();
        for (Entry<Character, String> e : dict.entrySet())
            ret += e.getKey() + " " + e.getValue() + "\n";
        return ret;
    }

    public static byte[] readFile(File inputFile) throws FileNotFoundException {
        String total_text = "";
        BufferedReader br = new BufferedReader(new FileReader(inputFile));
        try {
            String line = null;
            while ((line = br.readLine()) != null)
                total_text += line;
        }
        catch (IOException e){e.printStackTrace();}
        return total_text.getBytes();
    }
}

class HuffTree {
    Node home = null;
    HashMap<Character, String> dict;
    HashMap<String, Character> un_dict;

    public String translate(char c) {
        if (!dict.containsKey(c))
            return null;
        return dict.get(c);
    }

    public Character translate(String s) {
        if (!un_dict.containsKey(s))
            return null;
        return un_dict.get(s);
    }

    public HuffTree(HashMap<Character, Integer> char_table)
    {
        List<Entry<Character, Integer>> freq_list = sortChars(char_table);
        Queue<Node> queue = new PriorityQueue<Node>();
        for(Entry<Character, Integer> e : freq_list)
            queue.add(new Node(e.getKey(), e.getValue()));
        while (queue.size() > 1) {
            Node b = queue.poll();
            Node a = queue.poll();
            Node new_node = new Node(a,b);
            queue.add(new_node);
        }
        home = queue.poll();
        dict = home.getDict();
        un_dict = new HashMap<String, Character>();
        for (Entry<Character, String> e : dict.entrySet())
            un_dict.put(e.getValue(), e.getKey());
    }

    public List<Entry<Character, Integer>> sortChars(HashMap<Character, Integer> char_table)
    {
        Comparator<Entry<Character, Integer>> valueComparator = new Comparator<Entry<Character,Integer>>() {
            @Override
            public int compare(Entry<Character, Integer> e1, Entry<Character, Integer> e2) {
                return -Integer.compare(e1.getValue(), e2.getValue());
            }
        };
        List<Entry<Character, Integer>> listOfEntries = new ArrayList<Entry<Character, Integer>>(char_table.entrySet());
        Collections.sort(listOfEntries, valueComparator);
        return listOfEntries;
    }

    public HashMap<Character, String> getDict() {
        return this.dict;
    }

    public HashMap<String, Character> getUnDict() {
        return this.un_dict;
    }

    public void print() {
        home.printNode(0);
    }
}

class Node implements Comparable<Node> {
    char value = 0;
    int weight;
    Node parent = null;
    Node a = null;
    Node b = null;
    String encoding = "";

    public Node(char value, int weight) {
        this.value = value;
        this.weight = weight;
    }

    public Node(Node a, Node b) {
        this.a = a;
        a.parent = this;
        this.b = b;
        b.parent = this;
        this.weight = a.weight + b.weight;
    }

    public String getEncoding(String e)
    {
        String s = "";
        if (parent != null) {
            if (parent.a == this)
                s += "1";
            else
                s += "0";
            s = this.parent.getEncoding(e) + s;
        }
        return s;
    }

    public HashMap<Character, String> getDict() {
        HashMap<Character, String> dict = new HashMap<Character, String>();
        if (value != 0)
            dict.put(value, this.getEncoding(""));
        else {
            dict.putAll(a.getDict());
            dict.putAll(b.getDict());
        }
        return dict;
    }

    public void printNode(int t) {
        for (int i = 0; i < t; i++)
            System.out.print("\t");
        System.out.print("(E=" + this.getEncoding(""));
        if (this.value == 0)
            System.out.print(", W="+ this.weight +")\n");
        else
            System.out.print(", V=" + this.value + ", W="+this.weight+")\n");
        if (this.a != null)
            a.printNode(t + 1);
        if (this.b != null)
            b.printNode(t + 1);
    }

    public int compareTo(Node arg0) {
        if (this.weight != arg0.weight)
            return Integer.compare(this.weight, arg0.weight);
        else
            return -Character.compare(this.value, arg0.value);
    }
}