class Corpus:

    def __init__(self, authors: list = None, books: list = None, tokens: list = None):
        """
        Create a new corpus object, which is empty by default, but which can be initialised optionally with a set of books.
        :param authors: A list of strings representing author names. If a set of books are being used to initialise the corpus, len(authors) must equal len(books) and len(tokens)
        :param books:  A list of strings representing book titles.
        :param tokens: A list of lists representing tokens in the books.
        """
        if authors and books and tokens:
            self.authors = authors
            self.books = books
            self.tokens = tokens
            self.texts = None
        else:
            self.authors = []
            self.books = []
            self.tokens = []
            self.texts = []

    def add_book(self, author: str, book: str, text: str):
        """
        Add a single book to the corpus. This can only be done if the corpus has not been initialised with books in the constructor.
        :param author: The author's name as string.
        :param book: The book title as string.
        :param text: The book content as string.
        """
        assert (len(self.tokens) == 0)
        self.authors.append(author)
        self.books.append(book)
        self.texts.append(text)

    def tokenise(self, tokenise):
        """
        Tokenise all books in the corpus using the custom tokenisation function.

        :param tokenise: a tokenise function which takes a str and returns a list of tokens. It is language-specific and should remove pronouns.
        """
        self.tokens = [tokenise(text) for text in self.texts]

    def split(self, segment_length: int = 1000):
        """
        Split the books in the corpus into smaller books of a specified maximum length. So if segment_length = 1000 and you have a book of length 1500,
        it will be split into two books of lengths 1000 and 500 respectively, with the same author and title attributes.
        This Corpus object is not modified but a new instance will be constructed.

        :param segment_length: The maximum length of books in the new corpus.
        :return: The new corpus with split books.
        """
        new_authors = []
        new_books = []
        new_tokens = []

        for author_id in range(len(self.authors)):
            for ctr, segment_start in enumerate(range(0, len(self.tokens[author_id]), segment_length)):
                segment_tokens = self.tokens[author_id][
                                 segment_start:min(segment_start + segment_length, len(self.tokens[author_id]))]
                new_authors.append(self.authors[author_id])
                new_books.append(self.books[author_id] + "_" + str(ctr))
                new_tokens.append(segment_tokens)

        return Corpus(new_authors, new_books, new_tokens)
