GABELIA

import sys
import heapq as hp
import os
import PrefCode

''' class for heap tree '''


class Node:
    def __init__(self, index, freq):
        self.index = index
        self.freq = freq
        self.left = None
        self.right = None

    def __lt__(self, other):
        return self.freq < other.freq


def rec(elem, len, res):
    if elem.index != None:
        res.append((elem.index, len))
        return
    rec(elem.left, len+1, res)
    rec(elem.right, len+1, res)


def huffman(arr):
    arr = [(arr[x], x) for x in range(len(arr))]

    # create queue
    pq = []
    for key in arr:
        node = Node(freq=key[0], index=key[1])
        hp.heappush(pq, node)

    # create tree
    while len(pq) > 1:
        f = hp.heappop(pq)
        s = hp.heappop(pq)

        merged = Node(freq=f.freq + s.freq, index=None)
        merged.left = f
        merged.right = s

        hp.heappush(pq, merged)

    # evaluate code lengths
    res = []
    rec(hp.heappop(pq), 0, res)
    return res


def craft_value(arr):
    # arr = [(freq, len)]
    sum = 0
    for _, len in arr:
        sum += 2**(0-len)
    return sum


def build_code(arr):
    # Create tuple of (length, index, code)
    for i in range(n):
        arr[i] = (arr[i][1], arr[i][0], '')

    # sort by length of codes
    arr = sorted(arr, key=lambda x: x[0])

    dlm = ''
    for i in range(len(arr)):
        tmp = list(arr[i])

        tmp[2] = dlm + '0' * (arr[i][0] - len(dlm))
        dlm = "{0:b}".format((int(tmp[2], 2) + 1))
        dlm = '0' * (len(tmp[2]) - len(dlm)) + dlm

        arr[i] = tmp

    # sort by length of codes
    arr = sorted(arr, key=lambda x: x[1])

    return arr


if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("Usage : python3 Huffman.py src-fname dest-fname")
        exit()

    src_fname = sys.argv[1]
    dst_fname = sys.argv[2]

    with open(src_fname, 'r') as f:
        n = int(f.readline())
        arr = list(map(float, f.readline().split()))

    # find length of codes with probabilities
    arr = huffman(arr)

    # sort with indices
    arr.sort(key=lambda x: x[0])

    # write length of codes in file
    tmp_fname = 'studentFiles/tmp'
    with open(tmp_fname, 'w+') as tf:
        # total codes
        tf.write('{}\n'.format(len(arr)))

        for _, len in arr:
            tf.write('{} '.format(len))

    PrefCode.main(tmp_fname, dst_fname)
    os.remove(tmp_fname)