bwtzip

# Author: Kieran Moran
# Student ID: 31482244
import sys

from st2sa import st2sa
from st2sa import read_file
from st2sa import clear_file


def bwt(sa, text):
    bwt_array = []
    for i in sa:
        bwt_array.append(text[i - 2])
    return bwt_array


def eliasCodeword(number):
    bit_array = [bit for bit in bin(number)[2:]]
    cur_len = len(bit_array)
    len_components = []
    while cur_len > 1:
        l_component = [bit for bit in bin(cur_len - 1)[2:]]
        l_component[0] = '0'
        cur_len = len(l_component)
        len_components += l_component
    bit_array = len_components + bit_array
    return bit_array


def buildCharArray(bwt):
    unique_char_arr = [0]*91
    for i in bwt:
        unique_char_arr[ord(i)-36] += 1
    return unique_char_arr


def numberOfUniqueChar(unique_char_arr):
    return sum(1 if element > 0 else 0 for element in unique_char_arr)


class Huffman:
    def __init__(self):
        # end of list is the smallest count
        self.root_nodes = []

    class Node:
        def __init__(self, count, letters):
            self.count = count
            self.upper_edge = None
            self.letters = letters
            self.children = []

    def build(self, char_count_arr):
        for i in range(len(char_count_arr)):
            if char_count_arr[i] > 0:
                self.root_nodes.append(self.Node(char_count_arr[i], [chr(i+36)]))
        self.root_nodes.sort(key=lambda node: node.count, reverse=True)

        while len(self.root_nodes) > 1:
            left = self.root_nodes.pop()
            left.upper_edge = '0'
            right = self.root_nodes.pop()
            right.upper_edge = '1'


            total_count = left.count + right.count
            new_letters = left.letters + right.letters
            new_node = self.Node(total_count, new_letters)
            new_node.children.append(left)
            new_node.children.append(right)

            new_index = None
            for i in range(len(self.root_nodes)):
                if total_count > self.root_nodes[i].count:
                    new_index = i
                    break

            if new_index is None:
                self.root_nodes.append(new_node)
            else:
                self.root_nodes.insert(new_index, new_node)

    def findLetter(self, letter):
        output_code = []
        cur_node = self.root_nodes[0]
        while len(cur_node.children) > 0:
            for node in cur_node.children:
                if letter in node.letters:
                    output_code.append(node.upper_edge)
                    cur_node = node
        return output_code


class WriteBin:
    def __init__(self, file):
        self.write_buf = []
        self.file= open(file, 'ab')

    def add_bit(self, bit):
        self.write_buf.append(bit)
        if len(self.write_buf) == 8:
            byte = int(''.join(self.write_buf), 2)
            self.file.write(bytes([byte]))
            self.write_buf = []

    def close(self):
        if len(self.write_buf) > 0:
            while len(self.write_buf) < 8:
                self.write_buf.append('0')
            byte = int(''.join(self.write_buf), 2)
            self.file.write(bytes([byte]))
            self.write_buf = []
        self.file.close()

def asciiBits(char):
    number = ord(char)
    output_arr = []
    while number > 0:
        bit = number % 2
        output_arr.append(str(bit))
        number = number // 2

    # edge case
    if char == '$':
        output_arr.append('0')
    return output_arr[::-1]




if __name__ == '__main__':
    clear_file('bwtencoded.bin')
    if len(sys.argv) != 2:
        print("Usage: python bwtzip.py <text filename>")
        sys.exit(1)

    text_file = sys.argv[1]
    text = read_file(text_file)
    sa = st2sa(text)
    text += ('$')
    bwt = bwt(sa, text)
    char_arr = buildCharArray(bwt)
    huffman = Huffman()
    huffman.build(char_arr)

    compressed_file = WriteBin('bwtencoded.bin')

    # add bwt length
    for bit in eliasCodeword(len(bwt)):
        compressed_file.add_bit(bit)

    # add n unique chars
    for bit in eliasCodeword(numberOfUniqueChar(char_arr)):
        compressed_file.add_bit(bit)

    # add unique chars & number of chars
    for i in range(1, len(char_arr)):
        if char_arr[i] > 0:
            letter = chr(i+36)
            huff_array = huffman.findLetter(letter)
            for bit in asciiBits(letter):
                compressed_file.add_bit(bit)
            for bit in eliasCodeword(len(huff_array)):
                compressed_file.add_bit(bit)
            for bit in huff_array:
                compressed_file.add_bit(bit)

    # do $ last
    print( asciiBits('$'))
    for bit in asciiBits('$'):
        compressed_file.add_bit(bit)
    for bit in eliasCodeword(3):
        compressed_file.add_bit(bit)
    for bit in huffman.findLetter('$'):
        compressed_file.add_bit(bit)

    run_len = 1
    i = 0
    while i < len(bwt):
        if i == len(bwt) - 1 or bwt[i] != bwt[i + 1]:
            for bit in huffman.findLetter(bwt[i]):
                compressed_file.add_bit(bit)
            for bit in eliasCodeword(run_len):
                compressed_file.add_bit(bit)
            run_len = 1
        else:
            run_len += 1
        i += 1

    compressed_file.close()