import sys
import nltk
nltk.download(['punkt', 'wordnet', 'averaged_perceptron_tagger', 'stopwords'])
# import libraries
import numpy as np
import pandas as pd
import pickle
import re
from nltk.tokenize import word_tokenize, sent_tokenize
from nltk.corpus import stopwords
from nltk.stem import WordNetLemmatizer
from nltk import pos_tag, ne_chunk
from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline
from sklearn.pipeline import FeatureUnion
from sklearn.metrics import confusion_matrix
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import GridSearchCV
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import CountVectorizer, TfidfTransformer
from sklearn.multioutput import MultiOutputClassifier
from sklearn.neighbors import KNeighborsClassifier
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics import classification_report
from sqlalchemy import create_engine

def load_data(database_filepath):
    
    '''Function to load data in the project. The database_filepath is ../Directory/Database.db
       It outputs the data sliced in X (independent variable), Y (target) and category names of the targets.
    '''
    
    engine = create_engine('sqlite:///'+str(database_filepath))
    df = pd.read_sql_table('messages_etl', con=engine)
    
    X = df.message
    Y = df[df.columns[4:]]
    category_names = Y.columns
    
    return X, Y, category_names


def tokenize(text):
    
    '''tokenize gets the text input and return it in words already lemmatized, stripped, without dots or commas
       and with no stopwords.    
    '''
    
    stop_words = stopwords.words('english')
    tokens = word_tokenize(re.sub('[.:,;-]', ' ', text))
    lemmatizer = WordNetLemmatizer()

    clean_tokens = []
    for tok in tokens:
        if tok not in stop_words:
            clean_tok = lemmatizer.lemmatize(tok).lower().strip()
            clean_tokens.append(clean_tok)

    return clean_tokens

def build_model():
    
    '''Function to return a model pipeline with RandomForestClassifier passed through GridSearch 
    
    '''
    
    pipeline = Pipeline([('vect', CountVectorizer(tokenizer=tokenize)),
                         ('tfidf', TfidfTransformer()),
                         ('clf', MultiOutputClassifier(estimator=RandomForestClassifier()))
                        ])
    
    parameters = {'tfidf__norm': ('l1', 'l2'),
                  'vect__ngram_range': ((1,1), (1,2)),
                  'clf__estimator__n_estimators': (35, 50)}

    model = GridSearchCV(pipeline, param_grid=parameters)
    
    return model

def evaluate_model(model, X_test, Y_test, category_names):
    
    '''Function to evaluate the model in a classification report.
       Input model, X_test (real input values sliced before), Y_test (real target values sliced before)
       and category_names to loop for in.
    '''
    
    y_pred = model.predict(X_test)
    
    for i in range(len(category_names)):
        
        print(category_names[i], classification_report(Y_test.T.values[i], y_pred.T[i]))
    
    

def save_model(model, model_filepath):
    
    '''Function to save the model in a pickle file.
    '''
    
    with open(str(model_filepath), 'wb') as f:
        pickle.dump(model, f)


def main():
    
    '''Main function to call the train_classifier
    '''
    
    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()
