import Foundation

// Пример:
// Входная строка 'aabbbc'
// Коды:
// a (ASCII 97) -> 01
// b (ASCII 98) -> 1
// c (ASCII 99) -> 00
// Результат: 01 01 1 1 1 00

/// Класс узла дерева Хаффмана
// Реализуем протокол Comparable для сортировки узлов
class Node: Comparable {
    let char : Character?
    let weight : UInt
    
    let left : Node?
    let right : Node?
    
    init(char : Character, weight: UInt) {
        self.char = char
        self.weight = weight
        left = nil
        right = nil
    }
    
    init(left: Node, right: Node) {
        self.weight = left.weight + right.weight
        self.left = left
        self.right = right
        self.char = nil
    }

    
    var isLeaf : Bool {
        return left == nil && right == nil
    }
    
    
    static func < (lhs: Node, rhs: Node) -> Bool {
        return lhs.weight < rhs.weight
    }
    
    static func == (lhs: Node, rhs: Node) -> Bool {
        return lhs.weight == rhs.weight
    }
}


/// Строит таблицу частот для каждого символа в строке
func generateFrequencyTable(_ value: String) -> [Character: UInt] {
    // Создание хэш-таблицы
    var frequencyTable: [Character: UInt] = [:]
    
    // Подсчёт вероятности для каждого символа
    // Вероятность равна количеству символов в строке
    for char in value {
        frequencyTable[char, default: 0] += 1
    }

    return frequencyTable
}


/// Строит дерево Хаффмана из списка узлов
func generateTree(_ map : [Character: UInt]) -> Node {
    // Преобразует таблицу частот в список узлов и сортирует его
    // создаём заполненный массив
    var list : [Node] = Array(map.map{ key, value in Node(char: key, weight: value) })
    // сортируем
    list.sort()
    
    // Создаём новый мутабельный массив
    // Изменяемость зависит от типа переменной
    // let - не изменяемая
    // var - изменяемая
    while list.count > 1 {
        let left = list.removeLast()
        let right = list.removeLast()
        
        list.append(Node(left: left, right: right))
        list.sort()
    }
    return list.first!
}

/// Генерирует таблицу кодов для каждого символа на основе дерева Хаффмана
func generateCodesIterative(root: Node) -> [Character: String] {
    var codes = [Character: String]()
    var stack : [(Node, String)] = [(root, "")]
    
    while !stack.isEmpty {
        let (node, path) = stack.removeLast()
        
        // если листок дерева, то присваивание символу код
        // иначе продолжаем собирать код для символа
        if node.isLeaf, let character = node.char {
            codes[character] = path.isEmpty ? "0" : path
        } else {
            if let right = node.right {
                stack.insert((right, path + "1"), at: 0)
            }
            
            if let left = node.left {
                stack.insert((left, path + "0"), at: 0)
            }
        }
    }
    
    return codes
}

/// Кодирует строку с использованием таблицы кодов
func encode(source: String, codes : [Character : String]) -> String {
    return source.map { codes[$0]! }.joined()
}

/// Декодирует строку с использованием дерева Хаффмана
func decode(source: String, tree: Node) -> String {
    var buffer : [Character] = []
    var node: Node? = tree
    
    // Если дерево состоит только из одного узла (один уникальный символ)
    if node?.isLeaf == true, let char = node?.char {
        return String(repeating: char, count: source.count)
    }
    
    for char in Array(source) {
        node = if char == "0" {
            node?.left
        }else {
            node?.right
        }
        
        if let c = node?.char{
            buffer.append(c)
            node = tree
        }
    }
    
    return String(buffer)
}

func main(){
    let source = "aabbbc";
    let frequencyMap = generateFrequencyTable(source)
    let tree = generateTree(frequencyMap)
    
    let codesMap = generateCodesIterative(root: tree)
    let encodedSource = encode(source: source, codes: codesMap)
    
    print("The '\(source)' to encode '\(encodedSource)'")
    print("Decode '\(encodedSource)' to '\(decode(source: encodedSource, tree: tree))'")
    
    // Вычисляем коэффициент сжатия
    let originalBits = source.count * 8 // Предполагаем 8 бит на символ в ASCII
    let compressedBits = encodedSource.count
    let compressionRatio = Double(originalBits) / Double(compressedBits)
    print("Коэффициент сжатия: \(String(format: "%.2f", compressionRatio))")
}

main()