# import the neccessary functions

import sys
import nltk
import pandas as pd
import numpy as np
from sqlalchemy import create_engine
from nltk.tokenize import sent_tokenize, word_tokenize
from nltk.stem import WordNetLemmatizer
from sklearn.multioutput import MultiOutputClassifier
from sklearn.model_selection import GridSearchCV
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline, FeatureUnion
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.feature_extraction.text import CountVectorizer, TfidfTransformer
from sklearn.metrics import classification_report
import pickle

def load_data(database_filepath):
    """This function loads the cleaned data from a sqlite file into X & Y variables"""
    engine = create_engine('sqlite:///' + database_filepath)
    df = pd.read_sql('SELECT * FROM disaster_clean', engine)
    X = df['message']
    Y = df[df.columns[4:]]
    cols = df.columns[4:]
    return X, Y, cols
    
    
def tokenize(text):
    """This function tokenizes text to be ready for processing"""
    tokens = word_tokenize(text)
    
    # initiate lemmatizer
    lemmatizer = WordNetLemmatizer()

    # iterate through each token
    clean_tokens = []
    for tok in tokens:
        
        # lemmatize, normalize case, and remove leading/trailing white space
        clean_tok = lemmatizer.lemmatize(tok).lower().strip()
        clean_tokens.append(clean_tok)

    return clean_tokens


def build_model():
    """This function builds a pipeline with gridsearch to generate optimized model"""
    # build pipeline
    pipeline = Pipeline([
    ('text_process',Pipeline([('vector',CountVectorizer(tokenizer=tokenize)), # text pipeline
                              ('tfid',TfidfTransformer())])),
    ('clf',MultiOutputClassifier(RandomForestClassifier()))])  # multi classification rand forest pipeline
    
    # gridsearch parameters
    parameters = {'text_process__vector__ngram_range': ((1, 1), (1, 2)),
              'text_process__vector__max_df': (0.5, 0.75, 1.0)}

    cv = GridSearchCV(pipeline, param_grid=parameters, verbose=10)
    return cv


def evaluate_model(model, X_test, Y_test, category_names):
    """Evaluate the model and return the classification report for each predicted variable"""
    y_pred = model.predict(X_test)
    for  ytru,yprd,col in zip(Y_test.values ,y_pred, category_names):
        print(col, ":", "\n", classification_report(ytru, yprd))
    

def save_model(model, model_filepath):
    """This function Saves the model in a pickle file"""
    pickle.dump(model, open(model_filepath, 'wb'))


def main():
    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()