from .consts import OFFSET

def bytes_to_unicode_original():
    # Can call this through
    # from transformers.models.gpt2.tokenization_gpt2 import bytes_to_unicode
    """
    Returns list of utf-8 byte and a mapping to unicode strings. We specifically avoids mapping to whitespace/control
    characters the bpe code barfs on.

    The reversible bpe codes work on unicode strings. This means you need a large # of unicode characters in your vocab
    if you want to avoid UNKs. When you're at something like a 10B token dataset you end up needing around 5K for
    decent coverage. This is a significant percentage of your normal, say, 32K bpe vocab. To avoid that, we want lookup
    tables between utf-8 bytes and unicode strings.
    """
    bs = (
        list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
    )
    cs = bs[:]
    n = 0
    for b in range(2**8):
        if b not in bs:
            bs.append(b)
            cs.append(2**8 + n)
            n += 1
    cs = [chr(n) for n in cs]
    return dict(zip(bs, cs))

def bytes_to_unicode_special(special_bytes, offset=OFFSET):
    bs = (
        list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
    )
    cs = bs[:]
    n = 0
    for b in range(2**8):
        if b not in bs:
            bs.append(b)
            cs.append(2**8 + n)
            n += 1
    
    for b in special_bytes:
        bs.append(offset+b)
        cs.append(2**8 + n)
        n += 1

    cs = [chr(n) for n in cs]
    return dict(zip(bs, cs))

byte2unic = bytes_to_unicode_original()
unic2byte = {value: key for key, value in byte2unic.items()}

def decode_uniced(scrambled):
    unic_nums = [unic2byte[a] for a in scrambled]
    char_byte = [chr(a) for a in unic_nums]
    full = "".join(char_byte)
    return full.encode("latin-1").decode()


def find_aether_bytes(encoding, range_start, range_end):
    aether_bytes = set()
    for i in range(range_start, range_end):
        pieces = chr(i).encode(encoding)
        for piece in pieces:
            aether_bytes.add(piece)
    return aether_bytes