import argparse
import os
import pandas as pd
import pickle
from sklearn import metrics
from sklearn.feature_selection import RFECV
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier, AdaBoostClassifier
from sklearn_porter import Porter


def dt_model(parameters, X, y, score_function):
    """Returns the decision tree model"""
    clf = DecisionTreeClassifier(min_samples_leaf=parameters['min_samples_leaf'],
                                 min_samples_split=parameters['min_samples_split'],
                                 criterion=parameters['criterion'],
                                 splitter='best')
    selector = RFECV(clf, step=1, cv=10, scoring=metrics.make_scorer(score_function)).fit(X, y)
    return selector

def rf_model(parameters, X, y, score_function):
    """Returns the random forest model"""
    clf = RandomForestClassifier(n_estimators=parameters['n_estimators'],
                                 min_samples_leaf=parameters['min_samples_leaf'],
                                 min_samples_split=parameters['min_samples_split'],
                                 criterion=parameters['criterion'],
                                 random_state=0)
    selector = RFECV(clf, step=1, cv=10, scoring=metrics.make_scorer(score_function)).fit(X, y)
    return selector

def adaboost_model(parameters, X, y):
    """Returns the AdaBoost model"""
    estimator = DecisionTreeClassifier(min_samples_leaf=parameters['min_samples_leaf'],
                                       min_samples_split=parameters['min_samples_split'],
                                       criterion=parameters['criterion'],
                                       splitter='best')
    clf = AdaBoostClassifier(base_estimator=estimator, n_estimators=100, random_state=0).fit(X, y)
    return clf

def export_model_to_js(estimator, filename):
    """Exports the trained model to JavaScript"""
    porter = Porter(estimator, language='js')
    output = porter.export(embed_data=True)
    with open(os.path.join(args.out_dir, filename), mode='w+', encoding='utf8') as f:
        f.write('export ' + output)

def serialize_model(model, filename):
    """Serialize the trained model into pickle format"""
    file_path = '/'.join(os.path.abspath(__file__).split('/')[:-5]) + '/assets/' + filename
    pickle.dump(model, open(file_path, 'wb'))


parser = argparse.ArgumentParser(description='Train a random forest model to recognize headings and a decision tree model to compute their levels.')
parser.add_argument('dataset_dir', help='folder containing the .csv files generated by build_dataset.py')
parser.add_argument('out_dir', help='folder in which to save the two trained models')
args = parser.parse_args()

dataset_dir = args.dataset_dir
paths = os.listdir(dataset_dir)
dataset = []

# retrieving the dataset
for path in paths:
    df = pd.read_csv(os.path.join(dataset_dir, path), header=0)
    df['font_size'] = df['font_size'].apply(lambda x: round(x, 2))
    df['label'] = df['label'].apply(lambda x: 1 if x == 'heading' else 0)

    dataset.append(df[['is_different_style', 'is_font_bigger', 'is_font_unique',
                       'text_case', 'word_count', 'different_color',
                       'is_number', 'font_size', 'is_bold',
                       'font_ratio', 'level', 'label']])

df_dataset = pd.concat(dataset)
X_heading = df_dataset[['is_different_style', 'is_font_bigger', 'is_font_unique',
                        'text_case', 'word_count', 'different_color',
                        'is_number', 'font_ratio']].to_numpy()
y_heading = list(df_dataset['label'])

# only taking the headings into account for the levels
df_level = df_dataset.loc[df_dataset['label']==1].reset_index()
X_level = df_level[['font_size', 'is_bold', 'text_case',
                    'is_font_bigger', 'different_color', 'word_count',
                    'font_ratio']].to_numpy()
y_level = list(df_level['level'])

# splitting the dataset into training and test sets
X_train1, X_test1, y_train1, y_test1 = train_test_split(X_heading, y_heading, test_size=0.2)
X_train2, X_test2, y_train2, y_test2 = train_test_split(X_level, y_level, test_size=0.2)

# these parameters are found through grid search
# n_estimators is only used for the Random Forest Classifier
parameters_heading = {'n_estimators': 50, 'min_samples_leaf':1, 'min_samples_split': 2, 'criterion': 'gini'}
parameters_level = {'n_estimators': 70, 'min_samples_leaf': 1, 'min_samples_split': 2, 'criterion': 'entropy'}              

# computing the models
selector_heading = rf_model(parameters_heading, X_heading, y_heading, metrics.f1_score)
selector_level = dt_model(parameters_level, X_level, y_level, metrics.accuracy_score)

# evaluating the performance
y_pred_heading = selector_heading.predict(X_test1)
y_pred_level = selector_level.predict(X_test2)
print('precision (heading detection):', metrics.precision_score(y_test1, y_pred_heading))
print('recall (heading detection):', metrics.recall_score(y_test1, y_pred_heading))
print('f1-score (heading detection):', metrics.f1_score(y_test1, y_pred_heading))
print('accuracy (headings level):', metrics.accuracy_score(y_test2, y_pred_level))

# transpiling the models to js
export_model_to_js(selector_heading.estimator_, 'model.js')
export_model_to_js(selector_level.estimator_, 'model_level.js')

# serializing the models with pickle
serialize_model(selector_heading, 'headings_model.pkl')
serialize_model(selector_level, 'levels_model.pkl')
