from os import listdir
import string
from pickle import dump, load


# load doc into memory
def load_doc(filename):
    # open the file as read only
    file = open(filename, "r", encoding="utf-8")
    # read all text
    text = file.read()
    # close the file
    file.close()
    return text


# split a document into news story and highlights
def split_story(doc):
    index = doc.find("@highlight")
    # split into story and highlights
    story, highlights = doc[:index], doc[index:].split("@highlight")
    # strip extra white space around each highlight
    highlights = [h.strip() for h in highlights if len(h) > 0]
    return story, highlights


# load all stories in a directory
def load_stories(directory):
    stories = list()
    for name in listdir(directory):
        filename = directory + "/" + name
        # load document
        doc = load_doc(filename)
        # split into story and highlights
        story, highlights = split_story(doc)
        # store
        stories.append({"story": story, "highlights": highlights})
    return stories


# clean a list of lines
def clean_lines(lines):
    cleaned = list()
    # prepare a translation table to remove punctuation
    table = str.maketrans("", "", string.punctuation)
    for line in lines:
        # strip source cnn office if it exists
        index = line.find("(CNN) -- ")
        if index > -1:
            line = line[index + len("(CNN)") :]
        # tokenize on white space
        line = line.split()
        # convert to lower case
        line = [word.lower() for word in line]
        # remove punctuation from each token
        line = [w.translate(table) for w in line]
        # remove tokens with numbers in them
        line = [word for word in line if word.isalpha()]
        # store as string
        cleaned.append(" ".join(line))
    # remove empty strings
    cleaned = [c for c in cleaned if len(c) > 0]
    return cleaned


# load all stories in a directory
def load_single_story(doc):
    stories = list()
    # doc = load_doc(filename)
    # split into story and highlights
    story, highlights = split_story(doc)
    # store
    stories.append({"story": story, "highlights": highlights})
    for example in stories:
        example["story"] = clean_lines(example["story"].split("\n"))
    print("stories--", stories)
    return stories
