from typing import Dict, List, Tuple, Iterable

Leaf = str
Tree = Tuple[Iterable["Tree"] | Leaf, Iterable["Tree"] | Leaf] | Leaf
Queue = List[Tuple[Tree, float | int]]

CodeBook = Dict[str, str]


def build_huffman_tree(occurrence_frequencies: Dict[str, float | int]) -> Tree:
    """
        arguments:
          - `occurrence_frequencies` can either count the number of occurrences
            of each character in the alphabet or their frequencies.

        return:
          a structure representing the binary Huffman encoding tree.
    """
    queue: Queue = [(x, px) for x, px in occurrence_frequencies.items()]

    while len(queue) > 1:
        # combine two smallest elements
        a, pa = __extract_min(queue)
        b, pb = __extract_min(queue)

        # insert new node
        chars: Tree = (a, b)
        weight = pa + pb
        queue.append((chars, weight))

    x, _ = __extract_min(queue)  # only root node left
    return x


def __extract_min(queue: Queue):  # -> Tuple[Tree, float | int]:
    sort_key = [px for _, px in queue]
    return queue.pop(sort_key.index(min(sort_key)))


def generate_code(huff_tree: Tree) -> CodeBook:
    """
        arguments:
          - `huff_tree` is a Huffman tree with embedded encoding, e.g. built
            with `build_huffman_tree`.

        return:
          a codebook, i.e. a dictionary where characters are keys and
          associated binary strings are values.
    """
    def aux(t, prefix: str) -> CodeBook:
        if isinstance(t, Leaf):
            return {t: prefix}
        else:
            lchild, rchild = t
            codebook = {}

            codebook.update(aux(lchild, prefix + '0'))
            codebook.update(aux(rchild, prefix + '1'))
            return codebook

    return aux(huff_tree, "")


def compress(text: str, codebook: CodeBook) -> str:
    return "".join(codebook[ch] for ch in text if ord(ch) < 128)


def build_decoding_dict(codebook: CodeBook) -> CodeBook:
    return {y: x for (x, y) in codebook.items()}


def decompress(bits: str, decodebook: CodeBook) -> str:
    prefix = ""
    result = []
    for bit in bits:
        prefix += bit
        if prefix in decodebook:
            result.append(decodebook[prefix])
            prefix = ""
    assert prefix == ""  # must finish last codeword
    return "".join(result)
