'''
- Author: Zhengxiang (Jack) Wang 
- GitHub: https://github.com/jaaack-wang
- Website: https://jaaack-wang.eu.org
- About: Code for generating train, dev, test, and gen sets 
containing string pairs applying the following string transducation
functions: identity (w --> w), reveral (w --> w^R), total reduplication  
(w --> ww), and quadratic copying (w --> w^|w|).  

- Note that for all the datasets generated by the three string functions,
the corresponding input strings are identical.
'''
import random
import itertools
from string import ascii_lowercase

from os import makedirs
from os.path import join

import sys
import pathlib
# import from local script
sys.path.insert(0, str(pathlib.Path(__file__).parent))
from utils import save_ds_in_txt

# define the alphbet. in this case, = 26 English letters
alphabet = ascii_lowercase

# where to store the generated data 
data_folder = "./data"

# define the in- and out-of- distribution ranges
in_distribution_ranges = [(6, 15)]
out_distribution_ranges = [(1, 5), (16, 30)]

# define the sizes for the train, dev, test, and gen sets
train_size, dev_size, test_size, gen_size = 1000, 1000, 5000, 5000

# random seed (may not help reproducibility though)
random.seed(8741983)


def get_identity_pair(w):
    '''Returns a pair of identical strings.'''
    return [w, w]


def get_rev_pair(w):
    '''Returns a pair of an input string and its reverse.'''
    return [w, w[::-1]]


def get_total_red_pair(w):
    '''Returns a pair of an input string and its total reduplication.'''
    return [w, w+w]


def get_quad_copy_pair(w):
    '''Returns a pair of an input string and its total reduplication.'''
    return [w, w * len(w)]


def all_words_of_length(length, alphabet):
    '''Returns all possible strings given the string length and an alphabet.'''
    return [''.join(list(b)) for b 
            in itertools.product(alphabet, repeat=length)]


def n_words_of_length(max_num_per_seq, 
                      length, alphabet, 
                      always_unique=False):
    '''Returns ::max_num_per_seq:: strings given the string length and an alphabet.
    When the possible number of strings is less than ::always_unique::, the number
    of returned strings equals ::max_num_per_seq:: if ::always_unique:: is set False 
    (by default) and thus the strings may contain duplicates; otherwise, simple return
    all the possible unique strings given the string length and the alphabet.'''
    
    if max_num_per_seq >= pow(len(alphabet), length):

        out = all_words_of_length(length, alphabet)
        if max_num_per_seq == len(out) or always_unique:
            return out
        
        multiplier = max_num_per_seq // len(out)
        if max_num_per_seq % len(out) != 0:
            multiplier = multiplier + 1

        out = out * multiplier
        random.shuffle(out)
        return out[:max_num_per_seq]
        

    out = set()
    while len(out) < max_num_per_seq:
        word = "".join(random.choices(alphabet, k=length))
        out.add(word)
    
    return list(out)


def generate_datasets(train_X, dev_X, test_X, gen_X, 
                      str_func, str_func_name, data_folder):
    '''Generates train/dev/test/gen sets for a given str_func
    and saves them in ::data_folder/str_func_name::. '''
    folder = join(data_folder, str_func_name)
    makedirs(folder, exist_ok=True)
    
    for X, fname in zip([train_X, dev_X, test_X, gen_X], 
                        ["train.txt", "dev.txt", "test.txt", "gen.txt"]):
        data = list(map(str_func, X))
        fp = join(folder, fname)
        save_ds_in_txt(data, fp)


def main():
    '''The main function generates all the needed datasets for the three given
    string transduction functions according to the configurations specified at
    the top of the current script.'''
    train_X, dev_X, test_X, gen_X = [], [], [], []
    total = sum((train_size, dev_size, test_size))
    split_1, split_2 = train_size, train_size + dev_size
    
    for (l, h) in in_distribution_ranges:
        for n in range(l, h+1):
            data = n_words_of_length(total, n, alphabet)
            train_X.extend(data[:split_1])
            dev_X.extend(data[split_1:split_2])
            test_X.extend(data[split_2:])

    for (l, h) in out_distribution_ranges:
        for n in range(l, h+1):
            gen_X.extend(n_words_of_length(gen_size, n, alphabet))
    
    str_funcs = [get_identity_pair, get_rev_pair,
                 get_total_red_pair, get_quad_copy_pair]
    str_func_names = ["identity", "rev", "total_red", "quad_copy"]
    

    for str_func, str_func_name in zip(str_funcs, str_func_names):
        generate_datasets(train_X, dev_X, test_X, gen_X, 
                          str_func, str_func_name, data_folder=data_folder)


if __name__ == "__main__":
    main()
