import os
import tarfile
import urllib
import email
import email.policy
from collections import Counter

from bs4 import BeautifulSoup
from html import unescape
import re
import nltk
from urlextract import URLExtract

import numpy as np

from TAD_tools_v00 import *

from sklearn.model_selection import train_test_split
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.pipeline import Pipeline
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import cross_val_score
from sklearn.metrics import precision_score, recall_score
from sklearn.model_selection import GridSearchCV

# file locations
ROOT_SITE = 'http://spamassassin.apache.org/old/publiccorpus/'
HAM_URL = ROOT_SITE + "20030228_easy_ham.tar.bz2"   # set of regular emails
SPAM_URL = ROOT_SITE + "20030228_spam.tar.bz2"      # set of spam emails
FOLDER_LOC = os.path.join('dataset')

def download_data(data_path=FOLDER_LOC, spam_url=SPAM_URL, ham_url=HAM_URL):
    '''
    Helper function:
        - create folders
        - load data from urls
    '''
    print('**LOADING EMAILS AND SPAMS**')
    print('   loading data...')
    
    # create folder if necessary
    if not os.path.isdir(data_path):
        os.mkdir(data_path)
    
    # download spam and ham data
    for filename, url in [('ham.tar.bz2', ham_url),('spam.tar.bz2', spam_url)]:
        # create full path and check if file has been dowloaded
        path = os.path.join(data_path, filename)
        if not os.path.isfile(path):
            # download file
            urllib.request.urlretrieve(url, filename=path)
        
        # open tar file and extract content
        with tarfile.open(path, mode='r') as tar_bz2_file:
            tar_bz2_file.extractall(path=data_path)
    print('   import completed...')

# download emails (spams and regular) and extract content of tar files
download_data()

# folder paths
HAM_DIR = os.path.join(FOLDER_LOC, "easy_ham")
SPAM_DIR = os.path.join(FOLDER_LOC, "spam")

# collect spam and ham files (a cmds appears when the content of the tar file are retrieved, this file is not included in collection)
ham_filenames = [filename for filename in sorted(os.listdir(HAM_DIR)) if filename!='cmds']
spam_filenames = [filename for filename in sorted(os.listdir(SPAM_DIR)) if filename!='cmds']

# count emails and spams
print('Dataset contains {} emails.'.format(len(ham_filenames)))
print('Dataset contains {} spams.'.format(len(spam_filenames)))

# inspect a file to understand its structure
sample_filename = os.path.join(HAM_DIR, ham_filenames[2])
sample_file_content = open(sample_filename, 'r').read()
print(sample_file_content)

def load_email(filename, is_spam, data_path=FOLDER_LOC):
    '''
    Helper function:
        - read email file and convert it into email object
    '''
    subdirectory = 'spam' if is_spam else 'easy_ham'
    
    # open file (read binary)
    with open(os.path.join(data_path, subdirectory, filename), 'rb') as mail:
        return email.parser.BytesParser(policy=email.policy.default).parse(mail)

# load all files
ham_files = [load_email(filename, False) for filename in ham_filenames]
spam_files = [load_email(filename, True) for filename in spam_filenames]

# verify that the extraction worked
print('Dataset contains {} emails.'.format(len(ham_files)))
print('Dataset contains {} spams.'.format(len(spam_files)))

# check sender information
print(spam_files[10].get_unixfrom().strip())
print('-'*80)
print()
# check file content
print(spam_files[10].get_content().strip())

# check sender information
print(ham_files[10].get_unixfrom().strip())
print('-'*80)
print()
# check file content
print(ham_files[10].get_content().strip())

multipart_ham, multipart_spam = 0, 0

for email in ham_files:
    multipart_ham+=email.is_multipart()
    
for email in spam_files:
    multipart_spam+=email.is_multipart()
    
print('{} hams contain multiparts.'.format(multipart_ham))
print('{} spams contain multiparts.'.format(multipart_spam))

def get_email_structure(mail):
    '''
    Recursive function used to retrieve email content
    '''
    # for multipart email, recursion on each part
    if mail.is_multipart():
        payload = mail.get_payload()
        return "multipart({})".format(', '.join([get_email_structure(sub_mail) for sub_mail in payload]))
    else:
        return mail.get_content_type()
    
def structures_counter(emails):
    '''
    Return counter containing key: email structure, value: count
    '''
    # use counter instead of dictionary to handle key initialization
    structures = Counter()
    
    # loop over email -> get structure -> update counter
    for email in emails:
        structure = get_email_structure(email)
        structures[structure] += 1
    return structures

structures_counter(ham_files).most_common()

structures_counter(spam_files).most_common()

# create y array containing binary classification
X = np.array(ham_files + spam_files)
y = np.array([0] * len(ham_files) + [1] * len(spam_files))

# split dataset
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, shuffle=True)

print('Training set size: \t{}'.format(X_train.shape))
print('Test set size: \t\t{}'.format(X_test.shape))
print('Target train set size: \t{}'.format(y_train.shape))
print('Target test set size: \t{}'.format(y_test.shape))

# helper function: retrieve text content from html content

def html_to_plain_text(htlm):
    '''
    Return plain text from html content.
    '''
    
    # replace hyperlink
    # flags
    # re.I: ignore case
    # re.M: multiline
    # re.S: make / match newline
    htlm = re.sub('<a\s.*?>', ' HYPERLINK ', htlm, flags=re.M | re.S | re.I)
    
    soup = BeautifulSoup(htlm, 'html.parser')
    
    # kill all script and style elements
    for script in soup(["script", "style"]):
        script.extract()    # rip it out

    # get text
    text = soup.get_text()
    
    # break into lines and remove leading and trailing space on each
    lines = (line.strip() for line in text.splitlines())
    
    # break multi-headlines into a line each
    chunks = (phrase.strip() for line in lines for phrase in line.split("  "))
    
    # drop blank lines
    text = '\n'.join(chunk for chunk in chunks if chunk)
    
    return unescape(text)

# find spam with the text/html structure
html_spam_emails = [email for email in X_train[y_train==1]
                    if get_email_structure(email) == "text/html"]

# select a candidate
sample_html_spam = html_spam_emails[24]
print(sample_html_spam.get_content().strip()[0:5000],'.....')

print(html_to_plain_text(sample_html_spam.get_content())[:5000], "...")

def email_to_text(email):
    '''
    Return text/plain email structure or converted html structure.
    '''
    
    # loop over email content to find text/plain or text/html
    for part in email.walk():
        
        # retrieve content type
        content_type = part.get_content_type()
        
        # skip content if type is not text/plain or text/html
        if not content_type in ('text/plain','text/html'):
            continue
        
        # try to extract content (error with encoding)
        try:
            content = part.get_content()
        except:
            content = str(part.get_payload())
        
        # return plain text or converted html
        if content_type == 'text/plain':
            return content
        else:
            html = content
            return html_to_plain_text(html)

print(email_to_text(sample_html_spam)[:200], "...")

stemmer = nltk.PorterStemmer()

for w in ('association', 'associated', 'associates', 'company', 'corporation', 'incorporated'):
    print ('{0}: {1}'.format(w,stemmer.stem(w)))

from nltk.corpus import stopwords 
from nltk.tokenize import word_tokenize 
nltk.download('stopwords')

class emailToWordCount(BaseEstimator, TransformerMixin):
    
    def __init__(self, replace_urls=True, stemming=True, replace_numbers=True, lower_case=True, remove_punctuation=True, remove_stop_words=True):
        self.replace_urls = replace_urls
        self.stemming = stemming
        self.replace_numbers = replace_numbers
        self.lower_case = lower_case
        self.remove_punctuation = remove_punctuation
        self.remove_stop_words = remove_stop_words
        
    def fit(self, X, y=None):
        return self
        
    def transform(self, X, y=None):
        X_transformed = []
        
        for email in X:
            
            # retrieve text content
            text = email_to_text(email) or ""
            
            # convert url to token
            if self.replace_urls:
                extractor = URLExtract()
                for url in list(set(extractor.find_urls(text))):
                    text = text.replace(url, ' URL ')
            
            # replace float, integers, scientific notation, and percent
            if self.replace_numbers:
                text = re.sub(r'([-+]?\d*\.?\d+[eE]?[+-]?\d*[%]?)', 'NUMBER', text)
                
            # switch content to lower case
            if self.lower_case:
                text = text.lower()
                
            # remove punctuation
            if self.remove_punctuation:
                text = re.sub(r'\W+', ' ', text, flags=re.M)
                
            # split email
            email_spilt = text.split()
            
            # remove stop words (english)
            if self.remove_stop_words:
                stop_words = set(stopwords.words('english')) 
                email_split_clean = []
                for word in email_spilt:
                    if not word in stop_words:
                        email_split_clean.append(word)
                email_spilt = email_split_clean
            
            word_counts = Counter(email_spilt)
            
            # replace words using stemming
            if self.stemming:
                stemmed_word_count = Counter()
                for word, count in word_counts.items():
                    stemmed_word = stemmer.stem(word)
                    stemmed_word_count[stemmed_word] += count
                
                word_counts = stemmed_word_count
            
            X_transformed.append(word_counts)
        return np.array(X_transformed)

X_sample = X_train[0:2]
X_sample_wordcounts = emailToWordCount().fit_transform(X_sample)
X_sample_wordcounts

from scipy.sparse import csr_matrix

class WordCountToVector(BaseEstimator, TransformerMixin):
    
    def __init__(self, vocabulary_size=10000, cap_size=10):
        self.vocabulary_size = vocabulary_size
        self.cap_size = cap_size
        
    def fit(self, X, y=None):
        total_count = Counter()
        for word_count in X:
            for word, count in word_count.items():
                total_count[word] += min(count, self.cap_size)
                
        most_common = total_count.most_common()[:self.vocabulary_size]
        self.most_common = most_common
        self.vocabulary_ = {word:index+1 for index, (word, count) in enumerate(most_common)}
        return self
        
    def transform(self, X, y=None):
        # array mapping: rows, cols, data
        # then create array using the mapping
        rows = []   # one row per email
        cols = []   # one col per word
        data = []   # word counter
        for row, word_count in enumerate(X):
            for word, count in word_count.items():
                rows.append(row)
                cols.append(self.vocabulary_.get(word,0))  # <OOV> is index 0 (key is not found)
                data.append(count)
        
        # csr_matrix((data, ij), [shape=(M, N)]
        return csr_matrix((data, (rows, cols)), shape=(len(X), self.vocabulary_size+1))

vocab_transformer = WordCountToVector(vocabulary_size=10, cap_size=20)
X_sample_vectors = vocab_transformer.fit_transform(X_sample_wordcounts)
X_sample_vectors

X_sample_vectors.toarray()

vocab_transformer.vocabulary_

# Create full pipeline email -> wordCount -> countVector
preprocess_pipeline = Pipeline([
    ("email_to_wordcount", emailToWordCount()),
    ("wordcount_to_vector", WordCountToVector()),
])

# transform the data
X_train_transformed = preprocess_pipeline.fit_transform(X_train)

# create model and perform cross-validation
log_clf = LogisticRegression(solver="liblinear", random_state=42)
score = cross_val_score(log_clf, X_train_transformed, y_train, cv=3, verbose=1, n_jobs=4)

print('Mean accuracy over {0} folds: {1:.2f}%'.format(len(score),score.mean()*100))
print('Standard deviation accuracy over {0} folds: {1:.2f}%'.format(len(score),score.std()*100))

# create full pipeline
full_pipeline = Pipeline([
    ("email_to_wordcount", emailToWordCount()),
    ("wordcount_to_vector", WordCountToVector()),
    ("clf_log_reg",LogisticRegression())
])

# create candidates value for hyperparameters
params = {
    "email_to_wordcount__stemming":[True, False],
    "email_to_wordcount__remove_stop_words":[True, False],
    "wordcount_to_vector__cap_size":[10, 25, 50],
    "clf_log_reg__penalty":['l1', 'l2'],
    "clf_log_reg__C":np.logspace(0, 4, 5)
}

# create grid-search
clf = GridSearchCV(full_pipeline, params, cv=3, verbose=3, n_jobs=4)

# fit grid-search
best_model = clf.fit(X_train, y_train)

print('Grid Search best accuracy: {:.2f}%\n'.format(clf.best_score_*100))

print('Grid Search best parameters:')
for param, value in clf.best_params_.items():
    print('  ',param,'=', value)

# create a model based on the best combination of parameters
best_clf = LogisticRegression(solver="liblinear", random_state=42, C=1.0, penalty='l2')
best_clf.fit(X_train_transformed, y_train)

# make prediction on training set
y_train_pred = best_clf.predict(X_train_transformed)
print('Results on training set\n')
plot_confusion_matrix(y_train, y_train_pred, np.array([0,1]), title='Results on train set')
print('Precision: {:.2f}%'.format(precision_score(y_train, y_train_pred)*100))
print('Recall: {:.2f}%'.format(recall_score(y_train, y_train_pred)*100))

# make prediction on test set
X_test_transformed = preprocess_pipeline.transform(X_test)
y_test_pred = best_clf.predict(X_test_transformed)
print('Results on test set\n')
plot_confusion_matrix(y_test, y_test_pred, np.array([0,1]), title='Results on test set')
print('Precision: {:.2f}%'.format(precision_score(y_test, y_test_pred)*100))
print('Recall: {:.2f}%'.format(recall_score(y_test, y_test_pred)*100))

# Extract vocabulary from pipeline
word_to_index_vocab = preprocess_pipeline.named_steps['wordcount_to_vector'].vocabulary_
index_to_word_vocab = dict(zip(word_to_index_vocab.values(), word_to_index_vocab.keys()))

sorted_weights = np.argsort(best_clf.coef_)

# print the 20 words with the largest weights leading to the 'ham' classification
for loc in sorted_weights[0,-20:]:
    print('index: {0}, word: {1}, coeff: {2:.3f}'.format(loc, index_to_word_vocab[loc], best_clf.coef_[0,loc]))

# print the 20 words with the largest weights leading to the 'spam' classification
for loc in sorted_weights[0,0:20]:
    print('index: {0}, word: {1}, coeff: {2:.3f}'.format(loc, index_to_word_vocab[loc], best_clf.coef_[0,loc]))