# general libraries
import time
import pickle
import nltk
import progressbar

# from nltk
from nltk.corpus import stopwords
from nltk.stem.wordnet import WordNetLemmatizer
lmtzr = WordNetLemmatizer()

# setting personal nyt api key
from nytimesarticle import articleAPI
api = articleAPI('5c134948feff4ec98d4386749ce2f9d4')

# article cleanup functions
# code from http://dlab.berkeley.edu/blog/scraping-new-york-times-articles-python-tutorial
def parse_articles(articles):
    '''
    This function takes in a response to the NYT api and parses
    the articles into a list of dictionaries
    '''
    news = []
    if 'response' not in articles.keys():
        print articles['message']
        return
    for i in articles['response']['docs']:
        dic = {}
        dic['id'] = i['_id']
        if i['abstract'] is not None:
            dic['abstract'] = i['abstract'].encode("utf8")
        dic['headline'] = i['headline']['main'].encode("utf8")
        dic['desk'] = i['news_desk']
        dic['date'] = i['pub_date'][0:10] # cutting time of day.
        dic['section'] = i['section_name']
        if i['snippet'] is not None:
            dic['snippet'] = i['snippet'].encode("utf8")
        dic['source'] = i['source']
        dic['type'] = i['type_of_material']
        dic['url'] = i['web_url']
        dic['word_count'] = i['word_count']
        # locations
        locations = []
        for x in range(0,len(i['keywords'])):
            if 'glocations' in i['keywords'][x]['name']:
                locations.append(i['keywords'][x]['value'])
        dic['locations'] = locations
        # subject
        subjects = []
        for x in range(0,len(i['keywords'])):
            if 'subject' in i['keywords'][x]['name']:
                subjects.append(i['keywords'][x]['value'])
        dic['subjects'] = subjects   
        news.append(dic)
    return(news)

def get_articles(date, query, pages=100, delay=1):
    '''
    This function accepts a year in string format (e.g.'1980')
    and a query (e.g.'Amnesty International') and it will 
    return a list of parsed articles (in dictionaries)
    for that year.
    '''
    all_articles = []
    bar = progressbar.ProgressBar()
    for i in bar(range(0,pages)): #NYT limits pager to first 100 pages. But rarely will you find over 100 pages of results anyway.
        time.sleep(delay)
        articles = api.search(q = query,
               fq = {'source':['Reuters','AP', 'The New York Times']}, # query should be in headline too 
               begin_date = date + '0101',
               end_date = date + '1231',
               sort='oldest',
               page = str(i))
        articles = parse_articles(articles)
        if articles is None:
            continue
        all_articles = all_articles + articles
    return(all_articles)

def get_headlines(news, display=1):
    # get ordered list of headlines
    headlines = []
    for article in news:
        headlines.append((article['headline'], article['date']))
    if display:
        print "number of headlines is : " + str(len(headlines))
        for headline in headlines:
            print headline
    return headlines

def get_headlines_noStopWords(news, display=1):
    raw_headlines = get_headlines(news)
    new_headlines = []
    for i in range(len(raw_headlines)):
        new_headlines.append((content_fraction(raw_headlines[i][0]), raw_headlines[i][1]))
    if display:
        print new_headlines
    return new_headlines

# filter non "key" words
def content_fraction(text):
    stopwords = nltk.corpus.stopwords.words('english')
    content = []
    for w in unicode(text, 'utf-8').split():
        if w.lower() not in stopwords:
            content.append(w)    
    return content

# get helper function
def getKey(item):
    return item[1]

# pick top key words
def generate_unique_key_words_list(headlines):
    word_count_dict = {}
    #pick words from headline, lemmatize, populate
    for headline in headlines:
        text = headline[0]
        for word in text:
            lemma_word = lmtzr.lemmatize(word)
            if lemma_word not in word_count_dict.keys():
                word_count_dict[lemma_word] = 1
            else:
                word_count_dict[lemma_word] += 1

    # sort by count
    word_count_list = word_count_dict.items()
    sorted_word_list = sorted(word_count_list, key=getKey, reverse=True)
    if len(sorted_word_list) > 54000:
        return sorted_word_list[:54000]
    return sorted_word_list

# save/load objects in pickle
def save_obj(obj, name):
    with open('obj/'+ name + '.pkl', 'wb') as f:
        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)

def load_obj(name ):
    with open('obj/' + name + '.pkl', 'rb') as f:
        return pickle.load(f)

# main method (for testing purposes)
def main():
    news = get_articles('2017','Volkswagen', pages=10)
    #save_obj(news, "test_news")

    headlines = get_headlines_noStopWords(news, display=1)
    #save_obj(headlines, "test_headlines")

    key_words = generate_unique_key_words_list(headlines)
    print key_words

if __name__ == "__main__": main()

