from pathlib import Path
import re
import random


def load_wikitext(path, samples=100000):
    """
    Loads the wikitext corpus by random choosing with replacement from the full corpus
    :param path: str, path to the data folder
    :param samples: int, number of documents to return
    :return: list of str, wikitext corpus
    """
    heading_pattern = '( \n [=\s].*[=\s] \n)'
    train_data = Path(f'{path}datasets/data_wikitext-103-raw/wiki.train.raw').read_text(encoding='utf-8')
    train_split = re.split(heading_pattern, train_data)
    train_headings = [x[7:-7] for x in train_split[1::2]]
    train_articles = [x for x in train_split[2::2]]
    return random.choices(train_articles, k=samples)