import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

from sklearn import svm
from sklearn.feature_extraction.text import CountVectorizer

from sklearn.naive_bayes import MultinomialNB, BernoulliNB
from sklearn.feature_extraction.text import TfidfTransformer
from sklearn.feature_extraction.text import TfidfVectorizer

def concat(*args):
    """ Concatenate any numer of input streams """
    items = []
    
    for item in args:
        if isinstance(item, str):
            items.append(item)
            
    return ' '.join(items)

#-----------------------------------------------------------------------------------------

def precision(predicted, truth, verbose=True):
    """ Calculate precision of predicted values """
    total = 0
    correct = 0
    
    for p, t in zip(predicted, truth):
        if p == 1:
            total += 1
            
            if t == 1:
                correct += 1

    if verbose:
        print("Precision {}/{} = {}".format(correct, total, correct/total))
    
    return correct/total
                
#-----------------------------------------------------------------------------------------
    
def recall(predicted, truth, verbose=True):
    """ Calculate recal of predicted values """
    total = 0
    correct = 0
    
    for p, t in zip(predicted, truth):
        if t == 1:
            total += 1
            
            if p == 1:
                correct += 1
    if verbose:
        print("Recal {}/{} = {}".format(correct, total, correct/total))
    
    return correct/total

#-----------------------------------------------------------------------------------------
        
def f1(predicted, truth):
    """ Calculate F1 score """
    p = precision(predicted, truth, False)
    r = recall(predicted, truth, False)
    
    return 2*p*r/(p+r)

#-----------------------------------------------------------------------------------------

"""
Explanation: Programming Assignment 4
605.744 Information Retrieval
Justin Ely
End of explanation
"""


train = pd.read_csv('phase1.train.shuf.tsv',
                   sep='\t',
                   header=None,
                   names=["assessment", "docid", "title", "authors", "journal", "issn", "year", "language", "abstract", "keywords"])

dev = pd.read_csv('phase1.dev.shuf.tsv',
                   sep='\t',
                   header=None,
                   names=["assessment", "docid", "title", "authors", "journal", "issn", "year", "language", "abstract", "keywords"])

test = pd.read_csv('phase1.test.shuf.tsv',
                   sep='\t',
                   header=None,
                   names=["assessment", "docid", "title", "authors", "journal", "issn", "year", "language", "abstract", "keywords"])

#vectorizer = TfidfVectorizer(sublinear_tf=True, max_df=0.5, stop_words='english')
#all_abs = [item if isinstance(item, str) else '' for item in train['title']]
#train_counts = vectorizer.fit_transform(all_abs)

dev_text = vectorizer.transform(dev['title'])
print(dev_text[0])

clf = svm.SVC().fit(train_counts, train['assessment'])
print(clf)

preds = clf.predict(dev_text)
preds[preds>0]

"""
Explanation: Data loading
Here we load the supplied data for initial training, development, and final testing.  We use the pandas library to parse the TSV files and load them each into separate dataframes
End of explanation
"""


vectorizer = CountVectorizer()
train_counts = vectorizer.fit_transform(train['title'])

#clf = MultinomialNB().fit(train_counts, train['assessment'])
clf = BernoulliNB(alpha=.001).fit(train_counts, train['assessment'])

preds = clf.predict(train_counts)

recall(preds, train['assessment'])
precision(preds, train['assessment'])
f1(preds, train['assessment'])

dev_text = vectorizer.transform(dev['title'])
preds = clf.predict(dev_text)

recall(preds, dev['assessment'])
precision(preds, dev['assessment'])
f1(preds, dev['assessment'])

clf.predict(vectorizer.transform([train.iloc[133]['abstract']]))

"""
Explanation: Bayes
Title Data
End of explanation
"""


train_data = [concat(*item) for item in zip(train['title'], train['abstract'], train['keywords'])]

vectorizer = CountVectorizer()
train_counts = vectorizer.fit_transform(train_data)

clf = BernoulliNB(alpha=.001).fit(train_counts, train['assessment'])

train_data = [concat(*item) for item in zip(dev['title'] + dev['abstract'] + dev['keywords'])]
dev_text = vectorizer.transform(train_data)
preds = clf.predict(dev_text)

recall(preds, dev['assessment'])
precision(preds, dev['assessment'])
f1(preds, dev['assessment'])

"""
Explanation: Title + Abstract + Keywords
End of explanation
"""


vectorizer = CountVectorizer()
train_counts = vectorizer.fit_transform(train['title'])

#vectorizer = TfidfVectorizer(sublinear_tf=True, max_df=0.5, stop_words='english')
#all_abs = [item if isinstance(item, str) else '' for item in train['title']]
#train_counts = vectorizer.fit_transform(train['title'])

clf = svm.LinearSVC(C=1).fit(train_counts, train['assessment'])
print(clf)

preds = clf.predict(train_counts)

recall(preds, train['assessment'])
precision(preds, train['assessment'])
f1(preds, train['assessment'])

dev_text = vectorizer.transform(dev['title'])
preds = clf.predict(dev_text)

recall(preds, dev['assessment'])
precision(preds, dev['assessment'])
f1(preds, dev['assessment'])

"""
Explanation: SVM
End of explanation
"""


