import json
import random


def load_arxiv(path, samples=100000):
    """
    Loads the arxiv corpus by random choosing with replacement from the full corpus

    Can lead to freezing in jupyter notebooks, run only in a normal python file

    :param path: str, path to the data folder
    :param samples: int, number of documents to return
    :return: list of str, arxiv corpus
    """

    # preparing a generator as the arxiv file is very large
    def get_metadata():
        with open(f'{path}datasets/data_arxiv-metadata-oai-snapshot.json', 'r') as f:
            for line in f:
                yield line

    metadata = get_metadata()
    corpus = []

    # extracting all abstracts
    for paper in metadata:
        corpus.append(json.loads(paper)['abstract'])
    return random.choices(corpus, k=samples)