huffman.py

"""

Code for compressing and decompressing using Huffman compression.

"""

from nodes import HuffmanNode, ReadNode


def get_bit(byte, bit_num):

    """ Return bit number bit_num from right in byte.

    @param int byte: a given byte
    @param int bit_num: a specific bit number within the byte
    @rtype: int

    >>> get_bit(0b00000101, 2)
    1
    >>> get_bit(0b00000101, 1)
    0
    """

    return (byte & (1 << bit_num)) >> bit_num





def byte_to_bits(byte):

    """ Return the representation of a byte as a string of bits.

    @param int byte: a given byte
    @rtype: str
    >>> byte_to_bits(14)
    '00001110'

    """

    return "".join([str(get_bit(byte, bit_num))
                    for bit_num in range(7, -1, -1)])




def bits_to_byte(bits):

    """ Return int represented by bits, padded on right.

    @param str bits: a string representation of some bits
    @rtype: int
    
    >>> bits_to_byte("00000101")
    5
    >>> bits_to_byte("101") == 0b10100000
    True
    """

    return sum([int(bits[pos]) << (7 - pos)
                for pos in range(len(bits))])




# ====================

# Functions for compression

def make_freq_dict(text):

    """ Return a dictionary that maps each byte in text to its frequency.

    @param bytes text: a bytes object
    @rtype: dict(int,int)

    >>> d = make_freq_dict(bytes([65, 66, 67, 66]))
    >>> d == {65: 1, 66: 2, 67: 1}
    True
    """
    dic = {}

    for byte in text:
        if byte in dic:
            dic[byte] = dic[byte] + 1
        else:
            dic[byte] = 1
    return dic

    
#    dic = {}
#    for byte in text:
#        dic[byte] = text.count(byte)
#    return dic
#text.count means it has to check again. colours of cars in a car park analogy.
#Uncomment when it creates as its checking text. One iteration. Increasing time to /
#20 mins for book.txt


def huffman_tree(freq_dict):

    """ Return the root HuffmanNode of a Huffman tree corresponding
    to frequency dictionary freq_dict.


    @param dict(int,int) freq_dict: a frequency dictionary
    @rtype: HuffmanNode

    >>> freq = {2: 6, 3: 4}
    >>> t = huffman_tree(freq)
    >>> result1 = HuffmanNode(None, HuffmanNode(3), HuffmanNode(2))
    >>> result2 = HuffmanNode(None, HuffmanNode(2), HuffmanNode(3))
    >>> t == result1 or t == result2
    True

    """

    l = []
    a = 0

    for key in freq_dict:
        a = HuffmanNode(key)
        a.number = freq_dict[key] 
        l.append(a)
        a = 0

    l.sort(key=lambda x: x.number, reverse=False)

    while len(l) != 1:
        r = HuffmanNode()
        r.left = l.pop(0)
        r.right = l.pop(0)
        r.number = r.right.number + r.left.number
        l.append(r)
        l.sort(key=lambda x: x.number, reverse=False)
        
    tree = l.pop()



    def delete_num(tree):

        """
        """

        if tree == None:
            return 

        if tree != None:
            tree.number = None
            delete_num(tree.left)
            delete_num(tree.right)


    delete_num(tree)
    return tree
            
        


def get_codes(tree):

    """ Return a dict mapping symbols from Huffman tree to codes.


    @param HuffmanNode tree: a Huffman tree rooted at node 'tree'
    @rtype: dict(int,str)


    >>> tree = HuffmanNode(None, HuffmanNode(3), HuffmanNode(2))
    >>> d = get_codes(tree)
    >>> d == {3: "0", 2: "1"}
    True
    """

    code = ''
    d = {}

    def helper(tree, code):
        """
        Append to dictionary and add codes

        @param HuffmanNode tree: a Huffman tree rooted at node 'tree'
        @param Code: a string of '0' and '1'

        @rtype: dict(int,str)

        >>> tree = HuffmanNode(None, HuffmanNode(3), HuffmanNode(2))
        >>> d = get_codes(tree)
        >>> d == {3: "0", 2: "1"}
        True
        """
        

        if tree is not None:
            if tree.left:
                helper(tree.left, code + '0')
            if tree.right:
                helper(tree.right, code + '1')
            if tree.is_leaf():
                d[tree.symbol] = code

    
    helper(tree, code)
    return d





def number_nodes(tree):

    """ Number internal nodes in tree according to postorder traversal;
    start numbering at 0.

    @param HuffmanNode tree:  a Huffman tree rooted at node 'tree'
    @rtype: NoneType

    >>> left = HuffmanNode(None, HuffmanNode(3), HuffmanNode(2))
    >>> right = HuffmanNode(None, HuffmanNode(9), HuffmanNode(10))
    >>> tree = HuffmanNode(None, left, right)
    >>> number_nodes(tree)
    >>> tree.left.number
    0
    >>> tree.right.number
    1
    >>> tree.number
    2
    """
        

    def node_helper(tree, n = 0):
        """
        """
        
        if tree.left != None and not tree.left.is_leaf():
            n = node_helper(tree.left, n)
        if tree.right != None and not tree.right.is_leaf():
            n = node_helper(tree.right, n)
        tree.number = n
        return n + 1

        


    
    node_helper(tree, n = 0)



def avg_length(tree, freq_dict):

    """ Return the number of bits per symbol required to compress text

    made of the symbols and frequencies in freq_dict, using the Huffman tree.

    @param HuffmanNode tree: a Huffman tree rooted at node 'tree'
    @param dict(int,int) freq_dict: frequency dictionary
    @rtype: float

    >>> freq = {3: 2, 2: 7, 9: 1}
    >>> left = HuffmanNode(None, HuffmanNode(3), HuffmanNode(2))
    >>> right = HuffmanNode(9)
    >>> tree = HuffmanNode(None, left, right)
    >>> avg_length(tree, freq)
    1.9
    >>> tree = HuffmanNode(None, HuffmanNode(None, HuffmanNode(9, None, None), HuffmanNode(3, None, None)), HuffmanNode(2, None, None))
    >>> avg_length(tree, freq)
    1.3
    """

    bin_dict = get_codes(tree)
    length = 0
    occurance = 0

    for key in bin_dict:
        length = length + ((len(bin_dict[key])) * (freq_dict[key]))
        occurance += freq_dict[key]
    return length / occurance

        
        

def cut_str(str, cut_size):
     """
     """

     return [str[i:i+8] for i in range(0, len(str), cut_size)]



def generate_compressed(text, codes):

    """ Return compressed form of text, using mapping in codes for each symbol.


    @param bytes text: a bytes object
    @param dict(int,str) codes: mapping from symbols to codes
    @rtype: bytes
    
    >>> d = {0: "0", 1: "10", 2: "11"}
    >>> text = bytes([1, 2, 1, 0])
    >>> result = generate_compressed(text, d)
    >>> [byte_to_bits(byte) for byte in result]
    ['10111000']
    >>> text = bytes([1, 2, 1, 0, 2])
    >>> result = generate_compressed(text, d)
    >>> [byte_to_bits(byte) for byte in result]
    ['10111001', '10000000']
    """
    
    l = []
    f = []
    s = ''

    
    for byte in text:
        a = codes[byte]
        l.append(a)

    s = s.join(l)
    b = cut_str(s, 8)

    for ch in b:
        f.append(bits_to_byte(ch))

    return bytes(f)

    

    

def tree_to_bytes(tree):

    """ Return a bytes representation of the Huffman tree rooted at tree.


    @param HuffmanNode tree: a Huffman tree rooted at node 'tree'
    @rtype: bytes

    The representation should be based on the postorder traversal of tree
    internal nodes, starting from 0.
    Precondition: tree has its nodes numbered.

    >>> tree = HuffmanNode(None, HuffmanNode(3), HuffmanNode(2))
    >>> number_nodes(tree)
    >>> list(tree_to_bytes(tree))
    [0, 3, 0, 2]
    >>> left = HuffmanNode(None, HuffmanNode(3), HuffmanNode(2))
    >>> right = HuffmanNode(5)
    >>> tree = HuffmanNode(None, left, right)
    >>> number_nodes(tree)
    >>> list(tree_to_bytes(tree))
    [0, 3, 0, 2, 1, 0, 0, 5]
    """
    
    l = []

    def node_byte_helper(tree, n=0):
        """
        """

        if tree.left.is_leaf() and tree:
            l.append(0)
            l.append(tree.left.symbol)
        else:
            n = node_byte_helper(tree.left, n)
            l.append(1)
            l.append(tree.left.number)

        if tree.right.is_leaf() and tree:
            l.append(0)
            l.append(tree.right.symbol)
        else:
            n = node_byte_helper(tree.right, n)
            l.append(1)
            l.append(tree.right.number)
            
        return n + 1


    node_byte_helper(tree, n=0)

    return bytes(l)


def num_nodes_to_bytes(tree):

    """ Return number of nodes required to represent tree (the root of a
    numbered Huffman tree).
    
    @param HuffmanNode tree: a Huffman tree rooted at node 'tree'
    @rtype: bytes
    
    """

    return bytes([tree.number + 1])





def size_to_bytes(size):

    """ Return the size as a bytes object.

    @param int size: a 32-bit integer to convert to bytes
    @rtype: bytes
    >>> list(size_to_bytes(300))
    [44, 1, 0, 0]
    
    """

    return size.to_bytes(4, "little")

def compress(in_file, out_file):

    """ Compress contents of in_file and store results in out_file.

    @param str in_file: input file to compress
    @param str out_file: output file to store compressed result
    @rtype: NoneType

    """

    with open(in_file, "rb") as f1:
        text = f1.read()
    freq = make_freq_dict(text)
    tree = huffman_tree(freq)
    codes = get_codes(tree)
    number_nodes(tree)
    print("Bits per symbol:", avg_length(tree, freq))
    result = (num_nodes_to_bytes(tree) + tree_to_bytes(tree) +
              size_to_bytes(len(text)))
    result += generate_compressed(text, codes)
    with open(out_file, "wb") as f2:
        f2.write(result)


# ====================
# Functions for decompression


def generate_tree_general(node_lst, root_index):

    """ Return the root of the Huffman tree corresponding

    to node_lst[root_index].



    The function assumes nothing about the order of the nodes in node_lst.


    @param list[ReadNode] node_lst: a list of ReadNode objects
    @param int root_index: index in 'node_lst'
    @rtype: HuffmanNode

    >>> lst = [ReadNode(0, 5, 0, 7), ReadNode(0, 10, 0, 12), \
    ReadNode(1, 1, 1, 0)]
    >>> generate_tree_general(lst, 2)
    HuffmanNode(None, HuffmanNode(None, HuffmanNode(10, None, None), \
    HuffmanNode(12, None, None)), \
    HuffmanNode(None, HuffmanNode(5, None, None), HuffmanNode(7, None, None)))

    """

    tree = HuffmanNode(None)
    r = node_lst[root_index]    

    if r.l_type != 1:
        tree.left = HuffmanNode(r.l_data)
    else:
        tree.left = generate_tree_general(node_lst, r.l_data)

    if r.r_type != 1:
        tree.right = HuffmanNode(r.r_data)
    else:
        tree.right = generate_tree_general(node_lst, r.r_data)
        
    return tree



def generate_uncompressed(tree, text, size):

    """ Use Huffman tree to decompress size bytes from text.



    @param HuffmanNode tree: a HuffmanNode tree rooted at 'tree'

    @param bytes text: text to decompress

    @param int size: number of bytes to decompress from text.

    @rtype: bytes

    """

    new_dic = {}
    bit_list = []
    new_list = []
    final_list = []
    s = ''
    bit_checker = ''
    
    
    for key, value in get_codes(tree).items():
        new_dic[value] = (key)
        

    for byte in (text):
        s = s + (byte_to_bits(byte))

    while len(final_list) < size:
        for bit in s:
            bit_checker = bit_checker + bit
            if bit_checker in new_dic:
                final_list.append(new_dic[bit_checker])
                bit_checker = ''
            else:
                break
                
            
            

    return bytes(final_list)


def bytes_to_nodes(buf):

    """ Return a list of ReadNodes corresponding to the bytes in buf.


    @param bytes buf: a bytes object
    @rtype: list[ReadNode]
    >>> bytes_to_nodes(bytes([0, 1, 0, 2]))
    [ReadNode(0, 1, 0, 2)]

    """

    lst = []

    for i in range(0, len(buf), 4):
        l_type = buf[i]
        l_data = buf[i+1]
        r_type = buf[i+2]
        r_data = buf[i+3]
        lst.append(ReadNode(l_type, l_data, r_type, r_data))

    return lst



def bytes_to_size(buf):

    """ Return the size corresponding to the
    given 4-byte little-endian representation.

    @param bytes buf: a bytes object
    @rtype: int
    
    >>> bytes_to_size(bytes([44, 1, 0, 0]))
    300

    """

    return int.from_bytes(buf, "little")





def uncompress(in_file, out_file):

    """ Uncompress contents of in_file and store results in out_file.

    @param str in_file: input file to uncompress
    @param str out_file: output file that will hold the uncompressed results
    @rtype: NoneType

    """

    with open(in_file, "rb") as f:
        num_nodes = f.read(1)[0]
        buf = f.read(num_nodes * 4)
        node_lst = bytes_to_nodes(buf)
        tree = generate_tree_general(node_lst, num_nodes - 1)
        size = bytes_to_size(f.read(4))
        with open(out_file, "wb") as g:
            text = f.read()
            g.write(generate_uncompressed(tree, text, size))





if __name__ == "__main__":

    import time


    mode = input("Press c to compress or u to uncompress: ")
    if mode == "c":
        fname = input("File to compress: ")
        start = time.time()
        compress(fname, fname + ".huf")
        print("compressed {} in {} seconds."
              .format(fname, time.time() - start))
    elif mode == "u":
        fname = input("File to uncompress: ")
        start = time.time()
        uncompress(fname, fname + ".orig")
        print("uncompressed {} in {} seconds."
              .format(fname, time.time() - start))