import sys

# import libraries
import nltk
nltk.download(['punkt', 'wordnet', 'averaged_perceptron_tagger'])

import re
import numpy as np
import pandas as pd
from nltk.tokenize import word_tokenize
from nltk.stem import WordNetLemmatizer
from sqlalchemy import create_engine

from sklearn.metrics import confusion_matrix
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


def load_data(messages_filepath, categories_filepath):
    '''
    reads message and category data from specified csv files to pandas dataframes. Then it converts and merges them to a single dataframe.
    
    INPUT:
    messages_filepath - string: path to the messages csv file
    categories_filepath - string: path to the categories csv file
            
    OUTPUT:
    df - dataframe: pandas dataframe containg the message and category data
    '''
    # load messages dataset
    messages = pd.read_csv(messages_filepath)
    
    # load categories dataset
    categories = pd.read_csv(categories_filepath)

    # merge datasets
    df = pd.merge(messages, categories, on='id')
    
    return df


def clean_data(df):
    '''
    cleans and transforms the dataframe to create the appropriate category columns.
    
    INPUT:
    df - dataframe: pandas dataframe containg the message and category data
            
    OUTPUT:
    df - dataframe: cleaned and transformed dataframe
    '''    
    # create a dataframe of the 36 individual category columns
    categories = df.categories.str.split(';', expand=True)
    
    # select the first row of the categories dataframe
    row = categories.iloc[0]

    # use this row to extract a list of new column names for categories.
    # one way is to apply a lambda function that takes everything 
    # up to the second to last character of each string with slicing
    category_colnames = [name[:-2] for name in row]
    
    # rename the columns of `categories`
    categories.columns = category_colnames
    
    # Convert category values to just numbers 0 or 1
    for column in categories:
        # set each value to be the last character of the string
        categories[column] = categories[column].apply(lambda x: x[-1])

        # convert column from string to numeric
        categories[column] = categories[column].astype(int)
    
    # keep only the rows with 0/1
    categories = categories[categories.related != 2]
    
    # drop the original categories column from `df`
    df = df.drop('categories', axis=1)
    
    # concatenate the original dataframe with the new `categories` dataframe
    df = pd.merge(df, categories, right_index=True, left_index=True)
    
    # drop duplicates
    df.drop_duplicates(subset='message', inplace=True) 

    return df


def save_data(df, database_filepath):
    '''
    save the created dataframe to a sql database.
    
    INPUT:
    df - dataframe: pandas dataframe containg the message and category data
            
    OUTPUT:
    sqlite table containg the data from the inputed dataframe
    '''        
    # connect to db
    path = 'sqlite:///' + database_filepath
    engine = create_engine(path)

    # save to db
    df.to_sql('messages', con=engine, index=False, if_exists='replace') 


def main():
    if len(sys.argv) == 4:

        messages_filepath, categories_filepath, database_filepath = sys.argv[1:]

        print('Loading data...\n    MESSAGES: {}\n    CATEGORIES: {}'
              .format(messages_filepath, categories_filepath))
        df = load_data(messages_filepath, categories_filepath)

        print('Cleaning data...')
        df = clean_data(df)
        
        print('Saving data...\n    DATABASE: {}'.format(database_filepath))
        save_data(df, database_filepath)
        
        print('Cleaned data saved to database!')
    
    else:
        print('Please provide the filepaths of the messages and categories '\
              'datasets as the first and second argument respectively, as '\
              'well as the filepath of the database to save the cleaned data '\
              'to as the third argument. \n\nExample: python process_data.py '\
              'disaster_messages.csv disaster_categories.csv '\
              'DisasterResponse.db')


if __name__ == '__main__':
    main()