"""
train_classifier.py is a ML pipeline that creates a classification model for the disaster response project from messages and categories data generated by process_data.py and stored in DisasterResponse.db.

To run the pipeline, run the code below in command prompt in the file directory:

python train_classifier.py DisasterResponse.db classifier.pkl
"""

import sys
import pandas as pd
import numpy as np
import pickle
from sklearn.externals import joblib
from sklearn.pipeline import Pipeline
from sklearn.metrics import confusion_matrix
from sklearn.metrics import classification_report
from sklearn.model_selection import GridSearchCV
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.svm import SVC
from sklearn.multioutput import MultiOutputClassifier
from sklearn.feature_extraction.text import CountVectorizer, TfidfTransformer
import sqlite3
from sqlalchemy import create_engine

import re
import nltk
nltk.download('punkt')
nltk.download('stopwords') # download for stop words
nltk.download('wordnet') # download for lemmatization
from nltk.stem.wordnet import WordNetLemmatizer
from nltk.corpus import stopwords
from nltk.tokenize import word_tokenize
from nltk.tokenize import sent_tokenize


def load_data(database_filepath):
    """
    load_data imports data from the database generated by process_data.py
    
    ARGS:
        database_filepath: relative path of the database file (e.g. ../data/DisasterResponse.db)
        
    RETURNS:
        X: feature variable (messages)
        Y: category variable (36 categories of disaster response)
        category_names: names of the 36 categories
    """
    
    engine = create_engine('sqlite:///' + database_filepath)
    df = pd.read_sql_table('DisasterMessages', engine)
    X = df.message
    Y = df.iloc[:, 4:40].values
    category_names = df.columns[4:40]
    return X, Y, category_names


def tokenize(text):
    """
    tokenize processes text input using nltk packages include normalizaation and lemmatization
    
    ARGS:
        text: string of words to be tokenized
        
    RETURNS:
        tokens - tokenized text to be used as features in the model
    """
    
    # Normalize text
    text = re.sub(r"[^a-zA-Z0-9]", " ", text.lower())
    # Tokenize text
    words = word_tokenize(text)
    # Remove stop words
    words = [w for w in words if w not in stopwords.words("english")]
    # Reduce words to their root form
    lemmed = [WordNetLemmatizer().lemmatize(w) for w in words]
    return lemmed


def build_model():
    """
    build_model 1) creates a pipeline including sklearn text feature extractions, multiple output random forest classifier, and 2) uses GridSearchCV to select the optimal pipeline parameters
    
    ARGS:
        none
    
    RETURNS:
        cv: a multi-output classifier pipeline that is tuned using the best parameters by GridSearchCV
    """
    
    # create pipeline
    pipeline = Pipeline([
        ('vect', CountVectorizer(tokenizer=tokenize)),
        ('tfidf', TfidfTransformer()),
        ('clf', MultiOutputClassifier(RandomForestClassifier()))
    ])
    # tune parameters
    parameters = {
        'vect__ngram_range': ((1, 1), (1, 2)),
        'vect__max_df': (0.5, 0.75, 1.0),
        'tfidf__use_idf': (True, False),
        'clf__estimator__min_samples_split':[2, 3],
    }
    cv = GridSearchCV(pipeline, param_grid=parameters)
    return cv


def evaluate_model(model, X_test, Y_test, category_names):
    """
    evaluate_model uses the model generated by build_model function to predict the test set and displays the classification report
    
    ARGS:
        model: the multi-output classification model created by build_model
        X_test: test feature variable (messages)
        Y_test: test categories variable
        category_names: 36 categories of the dataset
    
    RETURNS:
        print out classification report score for each category
    """
    y_pred = model.predict(X_test)
    # display classification report for each categories:
    # f1 score, precision, recall
    print("================================")
    for i in range(0, 36):
        print("Classification report for category:", category_names[i])
        cr = classification_report(Y_test[:,i], y_pred[:,i])
        print(cr)
    print("================================")
    print("accuracy values are:")
    accuracy = (y_pred == Y_test).mean()
    print(accuracy)


def save_model(model, model_filepath):
    """
    save_model saves the model as a pickle file to the model_filepath
    
    ARGS:
        model: multi-output classification model generated by build_model()
        model_filepath: path of the model.pkl file (e.g. ../model/classifier.pkl)
    
    RETURNS:
        none
    """
    
    joblib.dump(model, model_filepath)


def main():
    """
    main pulls the system arguments from command line and performs the following:
    1. loads data from database and creates feature and category variables
    2. builds a model using GridSearchCV and saves it
    3. fits the model by training data and evaluates the model by its classification report for each category.
    """
    
    if len(sys.argv) == 3:
        database_filepath, model_filepath = sys.argv[1:]
        print('Loading data...\n    DATABASE: {}'.format(database_filepath))
        X, Y, category_names = load_data(database_filepath)
        X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=0.2)
        
        print('Building model...')
        model = build_model()
        
        print('Training model...')
        model.fit(X_train, Y_train)
        
        print('Evaluating model...')
        evaluate_model(model, X_test, Y_test, category_names)

        print('Saving model...\n    MODEL: {}'.format(model_filepath))
        save_model(model, model_filepath)

        print('Trained model saved!')

    else:
        print('Please provide the filepath of the disaster messages database '\
              'as the first argument and the filepath of the pickle file to '\
              'save the model to as the second argument. \n\nExample: python '\
              'train_classifier.py ../data/DisasterResponse.db classifier.pkl')


if __name__ == '__main__':
    main()