5 курс / lab4

import heapq
import math
import os
from collections import Counter


class HuffmanNode:
    def __init__(self, char=None, freq=0):
        self.char = char
        self.freq = freq
        self.left = None
        self.right = None
        self.code = ""

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


class HuffmanCoder:
    def __init__(self):
        self.text = ""
        self.frequencies = {}
        self.fixed_length_codes = {}
        self.huffman_codes = {}
        self.reverse_huffman_codes = {}
        self.huffman_tree = None

    def load_from_file(self, filename):
        try:
            with open(filename, 'r', encoding='utf-8') as file:
                self.text = file.read()
            self._calculate_frequencies()
            print(f"Файл успешно загружен.")
            print(f"Количество символов: {len(self.text)}")
            print(f"Количество уникальных символов: {len(self.frequencies)}")
            return True
        except FileNotFoundError:
            print(f"Файл '{filename}' не найден.")
            return False
        except Exception as e:
            print(f"Ошибка при загрузке файла: {e}")
            return False

    def _calculate_frequencies(self):
        self.frequencies = {}
        for char in self.text:
            self.frequencies[char] = self.frequencies.get(char, 0) + 1

    def show_text(self):
        if not self.text:
            print("Текст не загружен.")
            return

        print("=" * 50)
        print("СОДЕРЖИМОЕ ТЕКСТОВОГО ФАЙЛА")
        print("=" * 50)
        print(self.text[:500] + ("..." if len(self.text) > 500 else ""))
        print("=" * 50)

    def show_frequencies(self):
        if not self.frequencies:
            print("Частоты символов не рассчитаны.")
            return

        print("\n" + "=" * 30)
        print("СИМВОЛЫ АЛФАВИТА (сортировка по частоте)")
        print("=" * 30)

        sorted_freq = sorted(self.frequencies.items(), key=lambda x: (-x[1], x[0]))

        for char, freq in sorted_freq:
            if char == '\n':
                display_char = "\\n"
            elif char == '\t':
                display_char = "\\t"
            elif char == ' ':
                display_char = "' '"
            elif ord(char) < 32:
                display_char = f"chr({ord(char)})"
            else:
                display_char = char

            print(f"Символ: {display_char:<10} Частота: {freq}")

    def generate_fixed_length_codes(self):
        if not self.frequencies:
            print("Сначала загрузите файл.")
            return

        self.fixed_length_codes = {}
        chars = sorted(self.frequencies.keys())
        num_bits = math.ceil(math.log2(len(chars))) if len(chars) > 1 else 1

        for i, char in enumerate(chars):
            binary_code = format(i, f'0{num_bits}b')
            self.fixed_length_codes[char] = binary_code

        print("\n" + "=" * 30)
        print("КОДЫ ФИКСИРОВАННОЙ ДЛИНЫ")
        print("=" * 30)
        print(f"Количество символов: {len(chars)}")
        print(f"Длина кода каждого символа: {num_bits} бит")

        for char, code in self.fixed_length_codes.items():
            if char == '\n':
                display_char = "\\n"
            elif char == '\t':
                display_char = "\\t"
            elif char == ' ':
                display_char = "' '"
            elif ord(char) < 32:
                display_char = f"chr({ord(char)})"
            else:
                display_char = char

            print(f"Символ: {display_char:<10} Код: {code}")

    def _build_huffman_tree(self):
        if not self.frequencies:
            return None

        priority_queue = []
        for char, freq in self.frequencies.items():
            node = HuffmanNode(char, freq)
            heapq.heappush(priority_queue, (freq, id(node), node))

        while len(priority_queue) > 1:
            freq1, id1, node1 = heapq.heappop(priority_queue)
            freq2, id2, node2 = heapq.heappop(priority_queue)

            merged_node = HuffmanNode(freq=freq1 + freq2)
            merged_node.left = node1
            merged_node.right = node2

            heapq.heappush(priority_queue, (freq1 + freq2, id(merged_node), merged_node))

        _, _, root = heapq.heappop(priority_queue)
        return root

    def _generate_codes(self, node, code=""):
        if node is None:
            return

        if node.char is not None:
            self.huffman_codes[node.char] = code
            self.reverse_huffman_codes[code] = node.char
            return

        self._generate_codes(node.left, code + "0")
        self._generate_codes(node.right, code + "1")

    def generate_huffman_codes(self):
        if not self.frequencies:
            print("Сначала загрузите файл.")
            return

        self.huffman_codes = {}
        self.reverse_huffman_codes = {}

        self.huffman_tree = self._build_huffman_tree()
        self._generate_codes(self.huffman_tree)

        print("\n" + "=" * 30)
        print("КОДЫ ХАФФМАНА")
        print("=" * 30)

        sorted_codes = sorted(self.huffman_codes.items(), key=lambda x: (len(x[1]), x[1]))

        total_bits = 0
        total_freq = sum(self.frequencies.values())

        for char, code in sorted_codes:
            if char == '\n':
                display_char = "\\n"
            elif char == '\t':
                display_char = "\\t"
            elif char == ' ':
                display_char = "' '"
            elif ord(char) < 32:
                display_char = f"chr({ord(char)})"
            else:
                display_char = char

            freq = self.frequencies[char]
            total_bits += len(code) * freq
            print(f"Символ: {display_char:<10} Код: {code:<15} Частота: {freq}")

        avg_length = total_bits / total_freq if total_freq > 0 else 0
        print(f"\nСредняя длина кода: {avg_length:.2f} бит")
        print(f"Общее количество бит для кодирования текста: {total_bits} бит")

    def compress_with_fixed_length(self, filename):
        if not self.text:
            print("Сначала загрузите файл.")
            return

        if not self.fixed_length_codes:
            self.generate_fixed_length_codes()

        compressed_bits = ""
        for char in self.text:
            compressed_bits += self.fixed_length_codes.get(char, "")

        bytes_array = []
        for i in range(0, len(compressed_bits), 8):
            byte = compressed_bits[i:i + 8]
            if len(byte) < 8:
                byte = byte.ljust(8, '0')
            bytes_array.append(int(byte, 2))

        with open(filename, 'wb') as file:
            file.write(bytes(bytes_array))

        original_size = len(self.text.encode('utf-8'))
        compressed_size = len(bytes_array)

        print(f"\nФайл сжат с использованием кодов фиксированной длины.")
        print(f"Сохранено в: {filename}")
        print(f"Исходный размер: {original_size} байт")
        print(f"Сжатый размер: {compressed_size} байт")
        print(
            f"Коэффициент сжатия: {original_size / compressed_size:.2f}" if compressed_size > 0 else "Коэффициент сжатия: ∞")

    def compress_with_huffman(self, filename):
        if not self.text:
            print("Сначала загрузите файл.")
            return

        if not self.huffman_codes:
            self.generate_huffman_codes()

        compressed_bits = ""
        for char in self.text:
            compressed_bits += self.huffman_codes.get(char, "")

        bytes_array = []
        for i in range(0, len(compressed_bits), 8):
            byte = compressed_bits[i:i + 8]
            if len(byte) < 8:
                byte = byte.ljust(8, '0')
            bytes_array.append(int(byte, 2))

        with open(filename, 'wb') as file:
            file.write(bytes(bytes_array))

        original_size = len(self.text.encode('utf-8'))
        compressed_size = len(bytes_array)

        print(f"\nФайл сжат с использованием кодов Хаффмана.")
        print(f"Сохранено в: {filename}")
        print(f"Исходный размер: {original_size} байт")
        print(f"Сжатый размер: {compressed_size} байт")
        print(
            f"Коэффициент сжатия: {original_size / compressed_size:.2f}" if compressed_size > 0 else "Коэффициент сжатия: ∞")

    def compare_file_sizes(self, original_file, fixed_file, huffman_file):
        try:
            original_size = os.path.getsize(original_file)
            fixed_size = os.path.getsize(fixed_file)
            huffman_size = os.path.getsize(huffman_file)

            print("\n" + "=" * 30)
            print("СРАВНЕНИЕ РАЗМЕРОВ ФАЙЛОВ")
            print("=" * 30)
            print(f"Исходный файл ({original_file}): {original_size} байт")
            print(f"Сжатый файл ({fixed_file}): {fixed_size} байт")
            print(f"Сжатый файл ({huffman_file}): {huffman_size} байт")

            print("\nЭффективность сжатия:")
            if original_size > 0:
                fixed_ratio = (fixed_size - original_size) / original_size * 100
                huffman_ratio = (original_size - huffman_size) / original_size * 100

                print(
                    f"- Фиксированная длина: {'увеличение' if fixed_ratio > 0 else 'сжатие'} на {abs(fixed_ratio):.1f}%")
                print(f"- Код Хаффмана: {'увеличение' if huffman_ratio < 0 else 'сжатие'} на {abs(huffman_ratio):.1f}%")

                if huffman_size > 0:
                    print(f"\nКоэффициент сжатия Хаффмана: {original_size / huffman_size:.2f}")
        except FileNotFoundError as e:
            print(f"Ошибка: {e}")


def main():
    coder = HuffmanCoder()

    while True:
        print("\n" + "=" * 50)
        print("ПРОГРАММА КОДИРОВАНИЯ ХАФФМАНА")
        print("=" * 50)
        print("1. Открыть текстовый файл")
        print("2. Вывести содержимое текстового файла")
        print("3. Вывести символы алфавита с частотами")
        print("4. Сгенерировать коды для символов алфавита")
        print("   4.1 Коды фиксированной длины")
        print("   4.2 Коды Хаффмана")
        print("5. Сжать файл (фиксированная длина)")
        print("6. Сжать файл (Хаффман)")
        print("7. Сравнить размеры файлов")
        print("8. Выход")
        print("=" * 50)

        choice = input("Выберите пункт меню: ")

        if choice == '1':
            filename = input("Введите имя файла: ")
            coder.load_from_file(filename)

        elif choice == '2':
            coder.show_text()

        elif choice == '3':
            coder.show_frequencies()

        elif choice == '4.1':
            coder.generate_fixed_length_codes()

        elif choice == '4.2':
            coder.generate_huffman_codes()

        elif choice == '5':
            if not coder.text:
                print("Сначала загрузите файл.")
                continue

            output_file = input(
                "Введите имя выходного файла (по умолчанию: fixed_compressed.bin): ") or "fixed_compressed.bin"
            coder.compress_with_fixed_length(output_file)

        elif choice == '6':
            if not coder.text:
                print("Сначала загрузите файл.")
                continue

            output_file = input(
                "Введите имя выходного файла (по умолчанию: huffman_compressed.bin): ") or "huffman_compressed.bin"
            coder.compress_with_huffman(output_file)

        elif choice == '7':
            original_file = input("Введите имя исходного файла: ")
            fixed_file = input("Введите имя файла с фиксированным кодированием: ")
            huffman_file = input("Введите имя файла с кодированием Хаффмана: ")
            coder.compare_file_sizes(original_file, fixed_file, huffman_file)

        elif choice == '8':
            print("Выход из программы.")
            break

        else:
            print("Неверный выбор. Пожалуйста, выберите пункт от 1 до 8.")


if __name__ == "__main__":
    main()