import struct
from functools import partial
from math import log2
word_alphabet = "@ABCDEFHJKLMNPQRTUVWXY0123456789"

def expand(string):
    """
    Expands each character in the string to a 5-byte 1-bit-per-byte
    representation, where the 5 bits represent the index into the alphabet.

    e.g. '@' => b'\0\0\0\0\0', 'C' => b'\0\0\0\1\1', etc.
    """

    result = []
    assert all(c in word_alphabet for c in string)
    for c in string:
        idx = word_alphabet.index(c)
        bits = bin(idx)[2:].rjust(5, '0')
        for bit in bits:
            result.append(int(bit))
    return bytes(result)

def contract(bytes, alphabet=None):
    """
    Contracts the 1-byte-per-bit representation into regular text again.

    If no alphabet is given, uses the word alphabet (i.e. is the inverse of expand).
    If an alphabet is given (e.g. 0123456789abcdef), uses that alphabet.
    The hex alphabet is used to go from IR phobos => hex phobos (as printed)
    """
    if not alphabet:
        alphabet = word_alphabet
        bits = 5
    else:
        bits = int(log2(len(alphabet)))

    chars = ""
    idx = 0
    for n, byte in enumerate(bytes):
        n %= bits
        idx += (byte << ((bits - 1) - n))
        if (n == bits - 1):
            chars += alphabet[idx]
            idx = 0
    return chars

contract_phobos = partial(contract, alphabet="0123456789abcdef")
