from numpy.random import Generator
from tasks import register
import string
import numpy as np

chars = string.ascii_letters + string.digits

@register()
def copy(rng: Generator, l: list[int]):
    l = rng.integers(*l)
    inp = ''.join(rng.choice(list(chars), size=l))
    return inp + '=', inp, None

@register()
def reverse(rng: Generator, l: list[int]):
    l = rng.integers(*l)
    inp = ''.join(rng.choice(list(chars), size=l))
    return inp + '=', inp[::-1], None

@register()
def MQAR(rng: Generator, l: list[int], m: list[int]=[3,6]):
    l = rng.integers(*l)
    inp = ''.join(rng.choice(list(chars), size=l))
    n = rng.integers(2, l // 2)
    out = ''
    mask = []
    for i in range(n):
        mm = rng.integers(*m)
        qi = rng.integers(mm, l)
        out += inp[qi-mm:qi] + inp[qi] + ';'
        mask += [0] * (mm) + [1, 0]

    return inp + '=', out, mask

@register()
def capitalize(rng: Generator, l: list[int]):
    l = rng.integers(*l)
    inp = ''.join(rng.choice(list(chars), size=l))

    tbl = str.maketrans(string.ascii_lowercase + string.ascii_uppercase, string.ascii_uppercase + string.ascii_lowercase)
    return inp + '=', inp.translate(tbl), None

@register()
def capitalize_reverse(rng: Generator, l: list[int]):
    l = rng.integers(*l)
    inp = ''.join(rng.choice(list(chars), size=l))

    tbl = str.maketrans(string.ascii_lowercase + string.ascii_uppercase, string.ascii_uppercase + string.ascii_lowercase)
    return inp + '=', inp[::-1].translate(tbl), None

if __name__ == '__main__':
    rng = np.random.default_rng(0)
    print(copy(rng, [10, 11]))
    print(MQAR(rng, [8, 9], [3, 4]))
    print(capitalize_reverse(rng, [10, 11]))
