package huffman;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;

public class Main {
    
    // for the compressed output file
    static FileOutputStream fos;
    
    // the Huffman code for each byte
    // position 256 is for "end of file" symbol
    static String[] codes = new String[257];
    
    // the bytes from the input file
    static byte[] bytesArray;
    
    // the following used to pack bits into bytes
    static int bitsindex = 0; 
    static int partialByte = 0;
    static int[] powers = {128, 64, 32, 16, 8, 4, 2, 1};    

    public static void main(String[] args) throws IOException {
        File file = new File("sample.txt");
        bytesArray = new byte[(int) file.length()];
        try (FileInputStream fis = new FileInputStream(file);) {
            fis.read(bytesArray);
            fis.close();
            fos = new FileOutputStream("compressed.bin");
        } 
        int [] freq = new int[257]; //byte array to count the frequency of each byte
        for (int i = 0; i < bytesArray.length; i++) {
            int b = bytesArray[i];
            //bytes starting with a 1 are interpreted as negative
            //The following trick fixes that
            freq[b & 0xFF]++;
        }
        freq[256]=1; //for eof
        int count = 0;
        for (int i = 0; i < 257; i++) {
            if (freq[i]>0) {
                count++;
            }
        }
        // if byte b is the ith byte with positive frequency,
        // values[i] is the value of byte b
        // posFreq[i] is the frequency of b
        int[] posFreq = new int[count]; 
        int[] values = new int[count];
        int index = 0;
        // initialize values and posFreq arrays
        for (int i = 0; i < 257; i++) {
            if (freq[i]>0) {
                posFreq[index]=freq[i];
                values[index++]= i;
            }                
        }
        int n = posFreq.length;
        System.out.println(n);
        Node[] q = new Node[n+1];
        // make n trees that are actually n leaves to be stored in a heap
        for (int i=1; i<=n; i++)
            q[i]=new Node(null, null, posFreq[i-1], values[i-1]);
        Heap h = new Heap(q, n);
        h.makeHeap();
        for (int i=1; i<n; i++) {
            // take the two smallest values from the heap
            // merge the trees and insert the new tree back in the heap
            Node x = h.extractMin();
            Node y = h.extractMin();
            Node z = new Node(x, y, x.freq+y.freq, 0);
            h.insert(z);
        }        
        h.a[1].printTree(0);
        //populate the codes array
        getCodes(h.a[1],"");
        printCodes();
        generateOutputFile(h.a[1], values.length, fos);
    }     
    public static void generateOutputFile(Node root, int size, FileOutputStream fos) throws IOException{
        // first byte is the number of different bytes occurring in the file to compress
        fos.write(size);
        // next output the bytes in order of the tree traversal
        outputBytes(root);
        outputTreeStructure(root);
        //output incomplete byte if necessary
        if (bitsindex != 0)
            fos.write(partialByte);
        bitsindex = 0;
        partialByte = 0;
        outputContents(); 
        if (bitsindex != 0) {
            fos.write(partialByte);
        }        
        fos.close();
    }
    public static void outputBytes(Node n) throws IOException{
        if (n.isLeaf()){ // if a leaf, output the byte
            if (n.value == 256) { //if eof, write two zero bytes
                fos.write(0);
                fos.write(0);
            } else 
            fos.write(n.value);
        }
        else {
            outputBytes(n.left);
            outputBytes(n.right);
        }
    }
    public static void outputTreeStructure(Node n) throws IOException{
        if (n.left == null) // leaf
            outputBit(1);
        else {
            outputBit(0);
            outputTreeStructure(n.left);
            outputTreeStructure(n.right);
        }
    }
    
    public static void outputBit(int n) throws IOException {
        // add a bit in the nth position of partialByte
        partialByte = partialByte+powers[bitsindex++]*n;
//        System.out.print(" "+String.format("0x%02X", partialByte));
        if (bitsindex == 8) {
            fos.write(partialByte);
            partialByte = 0;
            bitsindex = 0;
        }
    }

    public static void outputContents() throws IOException {
        String code;
        for (int i=0; i<bytesArray.length; i++) {
            code = codes[bytesArray[i] & 0xFF];
            for (int j=0; j<code.length(); j++)
                outputBit(code.charAt(j)=='0'?0:1);        
        }
        // add end of file code
        code = codes[256];
        for (int j = 0; j < code.length(); j++) {
            outputBit(code.charAt(j) == '0' ? 0 : 1);
        }      
    }
    public static void getCodes(Node n, String prefix) {
        if (n.isLeaf()) {
            codes[n.value] = prefix;
        } else {
            getCodes(n.left, prefix + "0");
            getCodes(n.right, prefix + "1");
        }
    }
    public static void printCodes() {
        for (int i=0; i<257; i++)
            if (codes[i] != null)
                System.out.println(codes[i]+" "+String.format("0x%03X", i)+" "+(char) (i & 0xFF));
    }
    public static int getLongestZeroesCode(Node n) { 
        int i=0;
        while (n.left != null) {
            i++;
            n=n.left;
        }
        return i;
    }   
    public static double weightedSum(Node n, int depth, String str){
      if (n.isLeaf()){         
          return n.freq*depth;
      }
      return weightedSum(n.left,depth+1,str+"0")+weightedSum(n.right,depth+1,str+"1");
    }
}