#!/usr/bin/env python3
"""
Trigrams Lab
"""

import random


def import_text(file):
    """
    Import text from file and return a list of words.
    :param file: input file containing starting text
    :return: list of words
    """

    # Only use alpha-numeric words from file
    with open(file=file, mode='r') as text:
        word_list = [word for word in text.read().split() if word.isalnum()]
    return word_list


def build_trigram_dict(word_list):
    """
    Generate trigram dictionary from a list of words
    :param word_list: list of words
    :return: trigram dictionary
    """

    trigram = {}
    for index in range(len(word_list) - 2):
        word1 = word_list[index]
        word2 = word_list[index + 1]
        word3 = word_list[index + 2]
        pair = (word1, word2)
        if pair not in trigram:
            trigram[pair] = [word3]
        else:
            trigram[pair].append(word3)
    return trigram


def generate_sentence(trigram):
    """
    Generate sentence based on trigram dictionary and return in list form
        1) Choose random starting key
        2) Create new sentence list and append random key
        3) Append key's value to list
        4) Loop n times, using last two elements of sentence list as trigram key, and append key's value to sentence
            * Continue to append to list until '.' is the final character in the last element or 15 words are generated
            * If key doesn't exist, break loop.
        5) Add '.' to end of sentence if not already there and capitalize first word.
    :param trigram: dictionary
    :return: new_sentence list
    """

    # Start of sentence
    sentence = []
    starting_key = random.choice(list(trigram.keys()))
    sentence.extend([word for word in starting_key])
    sentence.append(trigram[starting_key][random.choice(range(0, len(trigram[starting_key])))])

    # Append words to sentence until length is 15 words or KeyError occurs.
    while len(sentence) < 15:
        try:
            key1, key2 = sentence[-2:]
            sentence.append(trigram[(key1, key2)][random.choice(range(0, len(trigram[(key1, key2)])))])
        except KeyError:
            break

    # Cleanup sentence punctuation
    sentence[0] = sentence[0].capitalize()
    if sentence[-1][-1] != '.':
        sentence[-1] = f'{sentence[-1]}.'
    return sentence


def generate_paragraph(trigram):
    """
    Combine sentences generated by generate_sentence() into paragraph
    :param trigram: dictionary
    :return: list of sentences
    """

    new_paragraph = []
    while len(new_paragraph) <= 10:
        new_paragraph.append(generate_sentence(trigram))
    return new_paragraph


def generate_new_book(text):
    """
    Iterate through generated text lists display new book
    :param text: list of sentences
    :return: None
    """

    for paragraph in text:
        for sentence in paragraph:
            for word in sentence:
                print(word, end=' ')
            print()
        print()


def main():
    # Import text and generate trigram dictionary
    initial_text = import_text('sherlock.txt')
    trigram = (build_trigram_dict(initial_text))

    # Generate 10 paragraphs of text based on trigram dictionary and store in list
    new_text = []
    for _ in range(10):
        new_text.append(generate_paragraph(trigram))

    # Generate new book
    generate_new_book(new_text)


if __name__ == '__main__':
    main()
