'''Viral Tweets
Tweet Classifier'''
#Imports
from nltk.tokenize import TweetTokenizer
ttok = TweetTokenizer()
import nltk
from nltk import ne_chunk, pos_tag, word_tokenize
from nltk.tree import Tree
from nltk.util import ngrams
from textblob import TextBlob as TB
import csv, random
import pickle, numpy

#Inputs tweet string
#Outputs feature dictionary
def get_features(tweet):
    tokens=ttok.tokenize(tweet)
    tags = TB(tweet).tags
    
    features = {}
    features['Length'] = len(tweet)
    features['# of Tokens'] = len(tokens)
    for gram in all_1gram_tokens:
        features['Has "%s"' % gram] = (gram in tweet.lower())
    for gram in all_2gram_tokens:
        features['Has "%s"' % gram] = (gram in tweet.lower())
    for gram in all_3gram_tokens:
        features['Has "%s"' % gram] = (gram in tweet.lower())
    features['Sentiment Polarity'] = TB(tweet).sentiment.polarity
    for i, tok in enumerate(tokens):
        if tok == 'flu':
            if i > 0:
                features['Before "flu"'] = tokens[i - 1]
            if i < (i - 1):
                features['After "flu"'] = tokens[i + 1]
    features['% Personal Pronouns'] = len([w for w in tags if w[1]=="PRP"])/len(tags)
    features['# of Named Entities'] = name_entities(tweet)
    return features

#Makes list of all uni-gram tokens in the data set
#If the data set is a list of tuples of tweets and classes
all_1gram_tokens=[]
def get_all_1gram_tokens(data_set):
    for tweet in data_set:
        toks = [t.lower() for t in ttok.tokenize(tweet[0])]
        grams = ngrams(toks, 1)
        for gram in grams:
            if " ".join(gram) not in all_1gram_tokens:
                all_1gram_tokens.append(" ".join(gram))
                
#Makes list of all bi-gram tokens in the data set
#If the data set is a list of tuples of tweets and classes
all_2gram_tokens=[]
def get_all_2gram_tokens(data_set):
    for tweet in data_set:
        toks = [t.lower() for t in ttok.tokenize(tweet[0])]
        grams = ngrams(toks, 2)
        for gram in grams:
            if " ".join(gram) not in all_2gram_tokens:
                all_2gram_tokens.append(" ".join(gram))
                
#Makes list of all tri-gram tokens in the data set
#If the data set is a list of tuples of tweets and classes
all_3gram_tokens=[]
def get_all_3gram_tokens(data_set):
    for tweet in data_set:
        toks = [t.lower() for t in ttok.tokenize(tweet[0])]
        grams = ngrams(toks, 3)
        for gram in grams:
            if " ".join(gram) not in all_3gram_tokens:
                all_3gram_tokens.append(" ".join(gram))
                
#Returns number of Named Entities
#http://stackoverflow.com/questions/31836058/nltk-named-entity-recognition-to-a-python-list
def name_entities(s):
     chunked = ne_chunk(pos_tag(word_tokenize(s)))
     continuous_chunk = []
     current_chunk = []
     for i in chunked:
             if type(i) == Tree:
                     current_chunk.append(" ".join([token for token, pos in i.leaves()]))
             elif current_chunk:
                     named_entity = " ".join(current_chunk)
                     if named_entity not in continuous_chunk:
                             continuous_chunk.append(named_entity)
                             current_chunk = []
             else:
                     continue
     return len(continuous_chunk) 

def main():
    #Reads in annotated tweet data (tweet, binary)
    annotated=csv.reader(open('Annotated.csv', encoding='utf8'))
    annotated_data=[(row[0], row[1]) for row in annotated][1:]
    print("Corpus size " + str(len(annotated_data)))    
    
    #Gets ngrams from labeled data
    print("Creating ngram lists")
    get_all_1gram_tokens(annotated_data)
    get_all_2gram_tokens(annotated_data)
    get_all_3gram_tokens(annotated_data)
    
    class_accs=[]
    baseline_accs=[]
    
    #Development trains and tests classifier
    for i in range (10):
        print("Dev iteration " + str(i))
        random.shuffle(annotated_data)
        dev_train = annotated_data[:int((len(annotated_data)+1)*.80)]
        dev_test = annotated_data[int(len(annotated_data)*.80+1):]
        dev_trainer = [(get_features(tweet), binary) for (tweet, binary) in dev_train]
        dev_classifier = nltk.NaiveBayesClassifier.train(dev_trainer)
        #dev_classifier.show_most_informative_features(6)
        class_accs.append(nltk.classify.accuracy(dev_classifier, [(get_features(tweet), binary) for (tweet, binary) in dev_test]))
        baseline_accs.append(len([c for c in dev_test if c[1]=='0'])/len(dev_test))
    
    print("Classifier Average Accuracy " + str(numpy.mean(class_accs)))
    print("Baseline Average Accuracy " + str(numpy.mean(baseline_accs)))
    
    #Trains classifier on all annotated data
    print("Training Final Classifier")
    train = [(get_features(tweet), binary) for (tweet, binary) in annotated_data]
    classifier = nltk.NaiveBayesClassifier.train(train)
    
    pickle.dump(classifier, open("classifier.pickle", "wb"))
    print("Classifier Pickled")

if __name__ == "__main__": 
    main()
