import sys
import struct

########################################################################################
# python lz_78_xor.py <d/c> <caminho_para_o_ficheiro_original>
# 
# c: comprimir
# d: descomprimir
# 
# Os ficheiros gerados ficarão na mesma pasta que o script
########################################################################################

def xor(bitstream, largura, altura):
    matriz = []
    for i in range(altura):
        linha = [int(b) for b in bitstream[i * largura : (i + 1) * largura]]
        matriz.append(linha)
    
    nova_imagem = ""
    # A primeira linha mantém-se igual (não tem linha anterior)
    nova_imagem += "".join(map(str, matriz[0]))
    
    for i in range(1, altura):
        nova_linha = []
        for j in range(largura):
            res = matriz[i][j] ^ matriz[i-1][j]
            nova_linha.append(str(res))
        nova_imagem += "".join(nova_linha)
        
    return nova_imagem

def descodificar_xor(bitstream_transformado, largura, altura):
    matriz_temp = []
    for i in range(altura):
        linha = [int(b) for b in bitstream_transformado[i * largura : (i + 1) * largura]]
        matriz_temp.append(linha)
    
    matriz_original = []
    
    matriz_original.append(matriz_temp[0])
    
    for i in range(1, altura):
        linha_recuperada = []
        for j in range(largura):
            pixel_original = matriz_temp[i][j] ^ matriz_original[i-1][j]
            linha_recuperada.append(pixel_original)
        matriz_original.append(linha_recuperada)
    
    bitstream_final = ""
    for linha in matriz_original:
        bitstream_final += "".join(map(str, linha))
        
    return bitstream_final

def lz78_compression(image):
    dictionary = {
        0: ""
    }

    output = []

    symbol = ""

    for i in image:
        if (symbol + i) in dictionary.values():
            symbol += i
        else:
            if symbol == "":
                output.append([0, i])
                dictionary[len(dictionary)] = i
            else:
                output.append([list(dictionary.keys())[list(dictionary.values()).index(symbol)], i])
                dictionary[len(dictionary)] = symbol + i
            symbol = ""

    if symbol != "":
        idx = list(dictionary.keys())[list(dictionary.values()).index(symbol)]
        output.append([idx, ""])

    return output, dictionary

def lz78_decompression(compressed_image):
    dictionary = {
        0: ""
    }

    output = ""

    for i in compressed_image:
        output += dictionary.get(i[0]) + i[1]
        dictionary[len(dictionary)] = dictionary.get(i[0]) + i[1]

    return output

def clean_image_data(image_data):
    
    clean_text = ""

    for line in image_data:
        line = line.strip()
        line = line.replace(" ", "")
        if not line or line.startswith('#'):
            continue
        clean_text += line

    return clean_text

def chunkstring(string, length):
    return [string[i : length+i] for i in range(0, len(string), length)]

#Comprimir
if sys.argv[1] == "c":

    with open(sys.argv[2], "r") as image:
        lines = image.readlines()
        image_size = lines[1].split()

        if lines[0].strip() != 'P1':
            raise ValueError("This is not a pbm file")
        
        clean_text = clean_image_data(lines[2:]) #Limpa os espaços em branco para comprimir melhor

        clean_text = xor(clean_text, int(image_size[0]), int(image_size[1]))

        output, dic = lz78_compression(clean_text)

        print(f"Output: {output}\n")
        print(f"Dictionary: {dic}")

        #Compressão em texto (human-readable)
        with open(sys.argv[2] + "_compressed", "w") as save_file:
            save_file.write(f"{output} \n{image_size[0]} {image_size[1]}")
        
        #Compressão em binário (compressão "a sério")
        with open(sys.argv[2] + ".bin", "wb") as save_file:
            largura = int(image_size[0])
            altura = int(image_size[1])
            save_file.write(struct.pack('II', largura, altura))

            for indice, simbolo in output:
                simbolo_byte = ord(simbolo) if simbolo != "" else 0
                save_file.write(struct.pack('HB', indice, simbolo_byte))
#Descomprimir
elif sys.argv[1] == "d":

    #Descomprimir do ficheiro de texto
    with open(sys.argv[2] + "_compressed", "r") as image:
        lines = image.readlines()

        compressed_image = eval(lines[0])
        image_size = lines[1]
        largura = int(lines[1].split(" ")[0])
        altura = int(lines[1].split(" ")[1])
        
        output = lz78_decompression(compressed_image)

        output = descodificar_xor(output, largura, altura)
        
        chuncked_output = chunkstring(output, largura)

        with open(sys.argv[2] + "_decompressed", "w") as save_file:
            save_file.write(f"P1\n{image_size}\n")
            for line in chuncked_output:
                save_file.write(f"{line}\n")
    
    #Descomprimir do binário
    with open(sys.argv[2] + ".bin", "rb") as image:
        header = image.read(8)
        largura, altura = struct.unpack('II', header)
        
        compressed_image_bin = []
        while True:
            chunk = image.read(3)
            if not chunk:
                break
            indice, simbolo_byte = struct.unpack('HB', chunk)
            simbolo = chr(simbolo_byte) if simbolo_byte != 0 else ""
            compressed_image_bin.append([indice, simbolo])

        output = lz78_decompression(compressed_image_bin)

        output = descodificar_xor(output, largura, altura)

        chuncked_output = chunkstring(output, largura)

        with open(sys.argv[2]+ "_bin_decompressed", "w") as save_file:
            save_file.write(f"P1\n{largura} {altura}\n")
            for line in chuncked_output:
                save_file.write(f"{line}\n")