from torchtext import data, datasets
import spacy


spacy_de = spacy.load('de')
spacy_en = spacy.load('en')

def tokenize_de(text):
    return [tok.text for tok in spacy_de.tokenizer(text)]

def tokenize_en(text):
    return [tok.text for tok in spacy_en.tokenizer(text)]

BOS_WORD = '<s>'
EOS_WORD = '</s>'
BLANK_WORD = "<blank>"


def load_data():
    SRC = data.Field(tokenize=tokenize_de, 
                    pad_token=BLANK_WORD)
    TGT = data.Field(tokenize=tokenize_en, 
                    init_token=BOS_WORD, 
                    eos_token=EOS_WORD, 
                    pad_token=BLANK_WORD)

    MAX_LEN = 100
    filter_pred = lambda x: len(vars(x)['src']) <= MAX_LEN and len(vars(x)['trg']) <= MAX_LEN
    train, val, test = datasets.IWSLT.splits(exts=('.de', '.en'),
                                            fields=(SRC, TGT),
                                            filter_pred=filter_pred)

    MIN_FREQ = 2
    SRC.build_vocab(train.src, min_freq=MIN_FREQ)
    TGT.build_vocab(train.trg, min_freq=MIN_FREQ)

    return train, val, test, SRC, TGT
