# Instruções de uso:
# Para codificar uma frase, salve ela como txt e rode "python3 huffman.py <nomedoarquivo.txt>".
# Para decodificar uma frase, salve a frase codificada como txt (ou utilize o arquivo codificado.txt
# gerado pelo comando anterior) e utilize o json gerado pelo comando anterior, rodando
# "python3 huffman.py <arquivocodificado.txt> data.json"

import json
import sys
import re
import os
from itertools import groupby
from collections import defaultdict
from heapq import heappop, heappush, heapify

codigo = {}
filaP = []
codificado = []

class Node(object):
        left = None
        right = None
        item = None
        freq = 0

        def __init__(self, i, f):
                self.item = i
                self.freq = f

        def setChildren(self, ln, rn):
                self.left = ln
                self.right = rn

        def __cmp__(self, a):
                return cmp(self.freq, a.freq)
        def __lt__(self, a):
                return self.freq < a.freq

def construirArvore(texto):
    global filaP
    filaP = [Node(a,len(list(b))) for a,b in groupby(sorted(texto))]

    heapify(filaP)

    while len(filaP) > 1:
        l = heappop(filaP)
        r = heappop(filaP)
        n = Node(None, r.freq+l.freq)
        n.setChildren(l,r)
        heappush(filaP, n)

def construirCodigo(s, node):
    if node.item:
        if not s:
            codigo[node.item] = "0"
        else:
            codigo[node.item] = s
    else:
        construirCodigo(s+"0", node.left)
        construirCodigo(s+"1", node.right)

def encode(texto,codigo):
        codificado = ""
        for i in texto:
                codificado = codificado + codigo.get(i)
        return codificado

def decode(codificado, codigo):
        texto = ""
        buffer = ""
        for i in codificado:
                buffer = buffer + i
                if buffer in codigo.values():
                        texto = texto + str([k for k, v in codigo.items() if v == buffer])
                        buffer = ""
        texto = texto.replace('[','').replace(']','').replace("'",'')
        return texto

def preencherJson(node):
	return {"ocorrencia":node.freq, "caracter":node.item, "NodeDaDireita":checarNoDir(node), "NodeDaEsquerda":checarNoEsq(node)}

def checarNoEsq(node):
	if node.left != None:
		return preencherJson(node.left)
	else:
		return None

def checarNoDir(node):
	if node.right != None:
		return preencherJson(node.right)
	else:
		return None

def huffman(texto):
	global codificado
	print("Texto original: " + "'" + texto + "'")
	construirArvore(texto)
	print("Árvore de Huffman:")
	construirCodigo("", filaP[0])
	codificado = encode(texto, codigo)
	print(codigo)
	print("Texto comprimido:")
	print(codificado)
	with open("codificado.txt", "w") as text_file:
		print("{}".format(codificado), file=text_file)

def huffman_inverso(texto):
	print("Texto comprimido: " + "'" + texto + "'")
	print("Mensagem original:")
	print(decode(texto,codigo))

# Inicio:

if len(sys.argv) == 1: # ve se tem um argumento
	print("Não foi passado um arquivo de texto.\n Para codificar, use: python3 huffman.py FraseAserCodificada.txt")
	print("Para descodificar, use: python3 fraseCodificada.txt arvore.json")
else:
	if len(sys.argv) == 2:
		with open (sys.argv[1], "r") as myfile:
			texto=myfile.read().replace('\n', '')
		huffman(texto)
		data = {"codigoFrase":codificado, "arvore":preencherJson(filaP[0])}
		with open('data.json', 'w') as outfile:
			json.dump(data, outfile, indent=4)
	else:
		with open (sys.argv[1], "r") as myfile:
			texto=myfile.read().replace('\n', '')

		with open(sys.argv[2], 'r') as json_file:
			#data = json.load(json_file)
			caracteres = []
			for line in json_file:
				if "ocorrencia" in line:
					result = re.search(": (.*),", line)
					freq_temp = result.group(1)
				if "caracter" in line:
					result = re.search(": (.*),", line)
					if "null" not in result.group(1):
						#print(freq_temp+" "+result.group(1))
						#filaP.append(Node(result.group(1),freq_temp))
						for x in range (0, int(freq_temp)):
							caracteres.append(result.group(1))
		caracteres.sort()
		construirArvore(caracteres)

		construirCodigo("", filaP[0])
		huffman_inverso(texto)