#Tyler Moore
#Huffman Coding Example
#CS2123
#Last Modified 9/28/2016
#huffman method adapted from code by Magnus Lie Hetland

from heapq import heapify, heappush, heappop
def huffman(seq, frq):
    """
    Input: seq: sequence of symbols to encode
           frq: frequency or probability distribution of symbols in seq
    Output: huffman tree
    """
    trees = list(zip(frq, seq))
    heapify(trees)              # A min-heap based on freq
    while len(trees) > 1:       # Until all are combined
        fa, a = heappop(trees)  # Get the two smallest trees
        fb, b = heappop(trees)
        heappush(trees, (fa+fb, [a, b])) # Combine and re-add
    return trees[0][-1]

def codes(tree, prefix = ""):
    if len(tree) == 1:
        return [tree, prefix]
    return codes(tree[0],prefix+"0")+ \
           codes(tree[1],prefix+"1")

def codesd(tr):
    cmap = {}
    def codesh(tree, prefix = ""):
        if len(tree) == 1:
            cmap[tree]=prefix
        else:
            codesh(tree[0],prefix+"0")
            codesh(tree[1],prefix+"1")
    codesh(tr,"")
    return cmap

if __name__=="__main__":
    seq = "abcdefghi"
    frq = [4,5,6,9,11,12,15,16,20]
    htree = huffman(seq,frq)
    print(htree)
    print(codes(htree))
    ch = codesd(htree)
    print(ch)
    text = "abbafabgee"
    print(text, " encodes to:")
    print("".join([ch[x] for x in text]))
    """
[['i',[['a','b'],'e']],[['f','g'],[['c','d'],'h']]]
['i','00','a','0100','b','0101','e','011','f','100',
'g','101','c','1100','d','1101','h','111']
{'a':'0100','c':'1100','b':'0101','e':'011',
'd':'1101','g':'101','f':'100','i':'00','h':'111'}
abbafabgee  encodes to:
010001010101010010001000101101011011
    """