from e2eml.classification import classification_blueprints as cb
import pandas as pd
from sklearn.metrics import matthews_corrcoef
from sklearn.metrics import classification_report
from sklearn.utils import shuffle
import re
#import pytest


def load_titanic_data():
    """
    Load & preprocess Titanic dataset. The feature engineering simulates the business knowledge part.
    The code has been taken from:
    https://towardsdatascience.com/predicting-the-survival-of-titanic-passengers-30870ccc7e8
    :return: Several dataframes and series to be processed by blueprint.
    """
    data = pd.read_csv("titanic_train.csv")
    print('Create additional features and modify existing ones.')
    deck = {"A": 1, "B": 2, "C": 3, "D": 4, "E": 5, "F": 6, "G": 7, "U": 8}
    data['Cabin'] = data['Cabin'].fillna("U0")
    data['Deck'] = data['Cabin'].map(lambda x: re.compile("([a-zA-Z]+)").search(x).group())
    data['Deck'] = data['Deck'].map(deck)
    data['Deck'] = data['Deck'].fillna(0)
    data['Deck'] = data['Deck'].astype(int)

    data.loc[data['Fare'] <= 7.91, 'Fare'] = 0
    data.loc[(data['Fare'] > 7.91) & (data['Fare'] <= 14.454), 'Fare'] = 1
    data.loc[(data['Fare'] > 14.454) & (data['Fare'] <= 31), 'Fare'] = 2
    data.loc[(data['Fare'] > 31) & (data['Fare'] <= 99), 'Fare'] = 3
    data.loc[(data['Fare'] > 99) & (data['Fare'] <= 250), 'Fare'] = 4
    data.loc[data['Fare'] > 250, 'Fare'] = 5

    data['Age'].fillna(0, inplace=True)
    data['Age'] = data['Age'].astype(int)
    data.loc[data['Age'] <= 11, 'Age'] = 0
    data.loc[(data['Age'] > 11) & (data['Age'] <= 18), 'Age'] = 1
    data.loc[(data['Age'] > 18) & (data['Age'] <= 22), 'Age'] = 2
    data.loc[(data['Age'] > 22) & (data['Age'] <= 27), 'Age'] = 3
    data.loc[(data['Age'] > 27) & (data['Age'] <= 33), 'Age'] = 4
    data.loc[(data['Age'] > 33) & (data['Age'] <= 40), 'Age'] = 5
    data.loc[(data['Age'] > 40) & (data['Age'] <= 66), 'Age'] = 6
    data.loc[data['Age'] > 66, 'Age'] = 6

    titles = {"Mr": 1, "Miss": 2, "Mrs": 3, "Master": 4, "Rare": 5}
    # extract titles
    data['Title'] = data.Name.str.extract(' ([A-Za-z]+)\.', expand=False)
    # replace titles with a more common title or as Rare
    data['Title'] = data['Title'].replace(['Lady', 'Countess','Capt', 'Col','Don', 'Dr',
                                                 'Major', 'Rev', 'Sir', 'Jonkheer', 'Dona'], 'Rare')
    data['Title'] = data['Title'].replace('Mlle', 'Miss')
    data['Title'] = data['Title'].replace('Ms', 'Miss')
    data['Title'] = data['Title'].replace('Mme', 'Mrs')
    # convert titles into numbers
    data['Title'] = data['Title'].map(titles)
    # filling NaN with 0, to get safe
    data['Title'] = data['Title'].fillna(0)

    data['relatives'] = data['SibSp'] + data['Parch']
    data.loc[data['relatives'] > 0, 'not_alone'] = 0
    data.loc[data['relatives'] == 0, 'not_alone'] = 1
    data['not_alone'] = data['not_alone'].astype(int)

    data['Fare'] = data['Fare'].astype(int)
    data['Age_Class'] = data['Age']*data['Pclass']
    data['Fare_Per_Person'] = data['Fare']/(data['relatives']+1)
    data['Fare_Per_Person'] = data['Fare_Per_Person'].astype(int)
    print('Do dataframe splits.')
    test_df = data.head(800).copy()
    val_df = data.tail(91).copy()
    val_df_target = val_df["Survived"].copy()
    del val_df["Survived"]
    test_target = "Survived"
    test_categorical_cols = ["Pclass", "Name", 'Sex', 'PassengerId']
    return test_df, test_target, val_df, val_df_target, test_categorical_cols


def steel_fault_multiclass_data():
    data = pd.read_csv("faults.csv")
    data = shuffle(data)
    #data["leakage"] = data["target"]
    print(data.info())
    #from sklearn.preprocessing import LabelEncoder
    #le = LabelEncoder()
    #X = data.drop('target', axis=1)
    #Y = le.fit_transform(data['target'])

    #X['target'] = Y
    test_df = data.head(1500).copy()
    val_df = data.tail(441).copy()
    val_df_target = val_df["target"].copy()
    del val_df["target"]
    test_target = "target"
    test_categorical_cols = None
    return test_df, test_target, val_df, val_df_target, test_categorical_cols


def nlp_multiclass_data():
    data = pd.read_csv("Corona_NLP_train.csv", encoding='latin-1')
    test_df = data.head(100).copy()
    print(test_df[["OriginalTweet"]])
    val_df = data.tail(499).copy()
    val_df_target = val_df["Sentiment"].copy()
    del val_df["Sentiment"]
    test_target = "Sentiment"
    test_categorical_cols = ["Location", "OriginalTweet"]
    return test_df, test_target, val_df, val_df_target, test_categorical_cols


def test_ml_special_multiclass_full_processing_multimodel_max_voting(dataset='titanic'):
    if dataset == 'titanic':
        test_df, test_target, val_df, val_df_target, test_categorical_cols = load_titanic_data()
        titanic_auto_ml = cb.ClassificationBluePrint(datasource=test_df,
                                       target_variable=test_target,
                                       categorical_columns=test_categorical_cols,
                                                     preferred_training_mode='auto',
                                                     tune_mode='accurate')
        titanic_auto_ml.hyperparameter_tuning_sample_size = 1100

    elif dataset == 'synthetic_multiclass':
        test_df, test_target, val_df, val_df_target, test_categorical_cols = steel_fault_multiclass_data()
        titanic_auto_ml = cb.ClassificationBluePrint(datasource=test_df,
                                                     target_variable=test_target,
                                                     categorical_columns=test_categorical_cols)
    elif dataset == 'corona_tweet':
        test_df, test_target, val_df, val_df_target, test_categorical_cols = nlp_multiclass_data()
        titanic_auto_ml = cb.ClassificationBluePrint(datasource=test_df,
                                                     target_variable=test_target,
                                                     categorical_columns=test_categorical_cols,
                                                     preferred_training_mode='auto',
                                                     tune_mode='accurate',
                                                     nlp_transformer_columns='OriginalTweet'
                                                     )

    titanic_auto_ml.hyperparameter_tuning_rounds = {"xgboost": 10,
                                                    "lgbm": 8,
                                                    "tabnet": 3,
                                                    "ngboost": 10,
                                                    "sklearn_ensemble": 3,
                                                    "catboost": 10,
                                                    "ridge": 3,
                                                    "bruteforce_random": 10}
    titanic_auto_ml.special_blueprint_algorithms = {"ridge": False, #titanic, #synthetic_multiclass
                                                    "xgboost": False, #titanic
                                                    "ngboost": False, #titanic, #synthetic_multiclass
                                                    "lgbm": True, #titanic
                                                    "tabnet": False, #titanic
                                                    "vowpal_wabbit": True,
                                                    "sklearn_ensemble": False, #titanic, #synthetic_multiclass
                                                    "catboost": False
                                                    }
    titanic_auto_ml.blueprint_step_selection_non_nlp["synthetic_data_augmentation"] = False


    titanic_auto_ml.ml_special_multiclass_full_processing_multimodel_max_voting()
    titanic_auto_ml.ml_special_multiclass_full_processing_multimodel_max_voting(val_df)
    val_y_hat = titanic_auto_ml.predicted_classes['max_voting']

    algorithms = [alg for alg, value in titanic_auto_ml.special_blueprint_algorithms.items() if value]

    def get_matthews(algorithm):
        # Assess prediction quality on holdout data
        print(classification_report(pd.Series(val_df_target).astype(bool),
                                    titanic_auto_ml.predicted_classes[algorithm]))
        try:
            matthews = matthews_corrcoef(pd.Series(val_df_target).astype(bool),
                                         titanic_auto_ml.predicted_classes[algorithm])
        except Exception:
            print("Matthew failed.")
            matthews = 0
        print(matthews)

    for i in algorithms:
        print(f"---------Start evaluating {i}----------")
        get_matthews(i)
    finished = True
    assert finished == True


if __name__ == '__main__':
    test_ml_special_multiclass_full_processing_multimodel_max_voting('synthetic_multiclass')
