import heapq

import os

class HuffmanNode:

    def __init__(self, char, freq):

        self.char = char

        self.freq = freq

        self.left = None

        self.right = None

    def __lt__(self, other):

        return self.freq < other.freq

def build_huffman_tree(frequencies):

    heap = [HuffmanNode(char, freq) for char, freq in frequencies.items()]

    heapq.heapify(heap)

    while len(heap) > 1:

        left = heapq.heappop(heap)

        right = heapq.heappop(heap)

        merged = HuffmanNode(None, left.freq + right.freq)

        merged.left = left

        merged.right = right

        heapq.heappush(heap, merged)

    return heap[0]

def build_huffman_codes(node, current_code, huffman_codes):

    if node is None:

        return

    if node.char is not None:

        huffman_codes[node.char] = current_code

        return

    build_huffman_codes(node.left, current_code + '0', huffman_codes)

    build_huffman_codes(node.right, current_code + '1', huffman_codes)

def compress(input_file, output_file):

    with open(input_file, 'rb') as f:

        data = f.read()

    frequencies = {}

    for byte in data:

        if byte not in frequencies:

            frequencies[byte] = 0

        frequencies[byte] += 1

    root = build_huffman_tree(frequencies)

    huffman_codes = {}

    build_huffman_codes(root, '', huffman_codes)

    compressed_data = ''

    for byte in data:

        compressed_data += huffman_codes[byte]

    padding = 8 - len(compressed_data) % 8

    compressed_data += '0' * padding

    with open(output_file, 'wb') as f:

        # Write frequency information

        f.write(bytes([len(frequencies)]))

        for byte, freq in frequencies.items():

            f.write(bytes([byte, (freq >> 24) & 0xFF, (freq >> 16) & 0xFF, (freq >> 8) & 0xFF, freq & 0xFF]))

        # Write compressed data

        for i in range(0, len(compressed_data), 8):

            byte = compressed_data[i:i+8]

            f.write(bytes([int(byte, 2)]))

def build_huffman_tree_from_bytes(frequencies):

    heap = [HuffmanNode(byte, freq) for byte, freq in frequencies.items()]

    heapq.heapify(heap)

    while len(heap) > 1:

        left = heapq.heappop(heap)

        right = heapq.heappop(heap)

        merged = HuffmanNode(None, left.freq + right.freq)

        merged.left = left

        merged.right = right

        heapq.heappush(heap, merged)

    return heap[0]

def decompress(input_file, output_file):

    with open(input_file, 'rb') as f:

        frequencies_count = ord(f.read(1))

        frequencies = {}

        for _ in range(frequencies_count):

            byte = ord(f.read(1))

            freq = (ord(f.read(1)) << 24) + (ord(f.read(1)) << 16) + (ord(f.read(1)) << 8) + ord(f.read(1))

            frequencies[byte] = freq

        root = build_huffman_tree_from_bytes(frequencies)

        huffman_codes = {}

        build_huffman_codes(root, '', huffman_codes)

        reverse_huffman_codes = {v: k for k, v in huffman_codes.items()}

        decompressed_data = bytearray()

        code = ''

        while True:

            byte = f.read(1)

            if len(byte) == 0:

                break

            byte = ord(byte)

            for i in range(7, -1, -1):

                if byte & (1 << i):

                    code += '1'

                else:

                    code += '0'

                if code in reverse_huffman_codes:

                    decompressed_data.append(reverse_huffman_codes[code])

                    code = ''

    with open(output_file, 'wb') as f:

        f.write(decompressed_data)

if __name__ == "__main__":

    # input_file = 'input.txt'

    compressed_file = 'compressed.bin'

    decompressed_file = 'decompressed.txt'

    # 压缩文件

    # compress(input_file, compressed_file)

    # 解压缩文件

    decompress(compressed_file, decompressed_file)