###########################
# DO NOT MODIFY THIS PART #
# BUT YOU SHOULD READ IT  #
###########################
from heapq import *

class Node:
    """
    Huffman tree node definition.
    """
    def __init__(self, symbol=None, count=0, left=None, right=None):
        """
        initialization
          symbol   : symbol to be coded
          count    : count of symbol
          left     : left child node
          right    : right child node
        """
        self.__left = left
        self.__right = right
        self.__symbol = symbol
        self.__count = count
        self.__code_word = ''

    def setLeft(self, l):
        """
        sets the left child of current node
        """
        self.__left = l
    
    def setRight(self, r):
        """
        sets the right child of current node
        """
        self.__right = r
    
    def getLeft(self):
        """
        returns the left child of current node
        """
        return self.__left
    
    def getRight(self):
        """
        returns the right child of current node
        """
        return self.__right

    def setSymbol(self, symbol):
        """
        sets coding symbol of current node
        """
        self.__symbol = symbol

    def getSymbol(self):
        """
        returns coding symbol of current node
        """
        return self.__symbol

    def setCount(self, count):
        """
        sets count of the symbol
        """
        self.__count = count

    def getCount(self):
        """
        returns count of the symbol
        """
        return self.__count
    
    def setCodeWord(self, code_word):
        """
        sets code-word of the symbol
        """
        self.__code_word = code_word

    def getCodeWord(self):
        """
        returns code-word of the symbol
        """
        return self.__code_word

    def __lt__(self, node):
        return self.__count < node.getCount()

    def __repr__(self):
        return "symbol: {}, count: {}, code-word: {}".format(self.__symbol, self.__count, self.__code_word)

###########################
# DO NOT MODIFY THIS PART #
# BUT YOU SHOULD READ IT  #
###########################



##############################
# FINISH THE BELOW FUNCTIONS #
##############################
def buildDictionary(message):
    """
    counts the occurrence of every symbol in the message and store it in a python dictionary
      parameter:
        message: input message string
      return:
        python dictionary, key = symbol, value = occurrence
    """
    _dict = dict()
    for c in message:
        if c not in _dict.keys():
            _dict[c] = 1
        else:
            _dict[c] += 1
    return _dict


def buildHuffmanTree(word_dict):
    """
    uses the word dictionary to generate a huffman tree using a min heap
      parameter:
        word_dict  : word dictionary generated by buildDictionary()
      return:
        root node of the huffman tree
    """
    tree = []
    for kv in word_dict.items():
        tree.append(Node(kv[0],kv[1]))
    heapify(tree)

    while True:
        if len(tree) == 1:
            break
        node1 = heappop(tree)
        node2 = heappop(tree)
        temp = Node(None, node1.getCount() + node2.getCount())
        temp.setLeft(node1)
        temp.setRight(node2)
        heappush(tree, temp)

    return tree[0]

def assignCodeWord(root, code_word=''):
    """
    recursively assigns code-word to the nodes in the huffman tree
      parameter:
        root       : root node of the huffman tree
        code_word  : code-word for the root node
      return:
        no return
    """
    if root.getSymbol() != None:
        if code_word == '':
            code_word = '0'
        root.setCodeWord(code_word)
        return
    else:
        assignCodeWord(root.getLeft(), code_word+'0')
        assignCodeWord(root.getRight(), code_word+'1')


def huffmanEncode(message):
    """
    converts the input message into huffman code
      parameter:
        message    : input message string
      return:
        a tuple, the first element is the huffman code string for the input message,
        the second element is the huffman tree root node
    """
    keyVal = buildDictionary(message)
    root = buildHuffmanTree(keyVal)
    assignCodeWord(root)

    stack = []
    _dict = dict()
    stack.append(root)
    while len(stack) != 0:
        v = stack.pop()
        if v.getSymbol() != None:
            _dict[v.getSymbol()] = v.getCodeWord()
        l = v.getLeft()
        r = v.getRight()
        if l != None:
            stack.append(l)
        if r != None:
            stack.append(r)

    encoded_message = ''
    for c in message:
        encoded_message += _dict[c]

    # print(encoded_message)
    return (encoded_message, root)


def huffmanDecode(message, huffman_tree):
    """
    decode the message
      parameter:
        message      : input huffman code string
        huffman_tree : huffman tree root node
      return:
        decoded message
    """
    decoded_message = ''
    current = huffman_tree

    if current.getSymbol() != None:
        for c in message:
            decoded_message += current.getSymbol()
    else:
        for c in message:
            if c == '0':
                current = current.getLeft()
            else:
                current = current.getRight()

            if current.getSymbol() != None:
                decoded_message += current.getSymbol()
                current = huffman_tree

    return decoded_message

def main():
    """
    main process goes here
    """
    message = input("Enter a message: ")
    if message == '':
        exit("No message entered. Exiting...")
    encoded_message, root = huffmanEncode(message)
    print("The encoded message is: " + encoded_message)
    decoded_message = huffmanDecode(encoded_message, root)
    print("The decoded message is: " + decoded_message)

##############################
# FINISH THE ABOVE FUNCTIONS #
##############################


###########################
# DO NOT MODIFY THIS PART #
###########################
if __name__ == "__main__":
    main()
###########################
# DO NOT MODIFY THIS PART #
###########################