# This file contains basic functions necessary to extract features,
# create datasets, train and predict use by machine-learning (ML)
# models. Also, contains several plotting functions.

import os
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
import matplotlib.pyplot as plt
import seaborn as sns
import warnings  
warnings.filterwarnings('ignore')


def get_features_df(fileName, type='all'):
    
    df = pd.read_csv(fileName)
    if 'taskspercpu' not in df:
        df['taskspercpu'] = df['numtasks'] / df['numcpus']
        df['taskspernode'] = df['numtasks'] / df['numnodes_x']
    featureNames = ['jobid', 'user', 'exename', 'start_y', 'numnodes_x', 'numcpus', 'numtasks', 'elapsedraw', 'avgpower_ldms', 'avgmemory_ldms','taskspercpu', 'taskspernode' ]
    if type == 'all':
        featureNames = ['jobid', 'user', 'exename', 'numnodes', 'numcpus', 'numtasks', 'executiontime', 'power', 'memory','taskspercpu', 'taskspernode', 'appname']
    elif type == 'hardware':
        featureNames = ['jobid', 'numnodes', 'numcpus', 'numtasks', 'executiontime', 'power', 'memory','taskspercpu', 'taskspernode', 'appname']
    elif type == 'jobsub':
        featureNames = ['jobid', 'user', 'exename', 'numnodes', 'numcpus', 'numtasks', 'power', 'taskspercpu', 'taskspernode', 'appname']
    elif type == 'custom':
        featureNames = ['jobid', 'user', 'exename', 'power', 'appname']
   
    features = df[featureNames]
    features = features.drop_duplicates()
    features = features.dropna()
    featuresAll = features
    features = features.drop('jobid', axis=1)
    features = features.drop('appname', axis=1)
    
    # One-hot encode categorical features
    features = pd.get_dummies(features)
    
    return features, featuresAll


def convert_feature_df_to_vars(features, predictionVariable = 'power'):
    # Labels are the values we want to predict
    labels = features[predictionVariable]

    # Remove the labels from the features
    features= features.drop(predictionVariable, axis = 1)

    # Save feature names
    feature_list = list(features.columns)
    
    return labels, feature_list, features
    
    
def create_train_test_sets(features, labels):
    # Using Skicit-learn to split data into training and testing sets
    train_features, test_features, train_labels, test_labels = train_test_split(features, labels, test_size = 0.25, random_state = 42)

    print('Training Features Shape:', train_features.shape)
    print('Training Labels Shape:', train_labels.shape)
    print('Testing Features Shape:', test_features.shape)
    print('Testing Labels Shape:', test_labels.shape)
    
    return train_features, test_features, train_labels, test_labels
    
    
def train_using_RandomForestRegressor(train_features, train_labels):
    # Instantiate model 
    rf = RandomForestRegressor(random_state = 42, n_estimators=1000)

    # Train the model on training data
    rf.fit(train_features, train_labels)
    
    return rf
    

def make_predictions(rf, test_features, test_labels):
    # Use the forest's predict method on the test data
    predictions = rf.predict(test_features)

    # Calculate the absolute errors
    errorsAbs = abs(predictions - test_labels)
    errors = (predictions - test_labels)

    # Print out the mean absolute error (mae)
    print('Mean Absolute Error:', round(np.mean(errorsAbs), 2), 'W.')

    # Calculate mean absolute percentage error (MAPE)
    mape = 100 * (errorsAbs / test_labels)

    # Calculate and display accuracy
    accuracy = 100 - np.mean(mape)
    print('Accuracy:', round(accuracy, 2), '%.')
    
    return predictions, errors, errorsAbs, accuracy
    
    
def plot_feature_importances(importances, feature_list, plotAll=True, outputFileName=None, threshold=0.02):
    # Set the style
    plt.style.use('default')
    plt.rc('font', family='serif')
    
    plot_data = pd.DataFrame(columns=['x_values','feature_list', 'importances'])
    plot_data['feature_list'] = feature_list
    plot_data['importances'] = importances
    plot_data = plot_data.sort_values('importances')
    
    if plotAll == False:
        plot_data = plot_data.loc[plot_data['importances'] > threshold]
    
    # list of x locations for plotting
    x_values = list(range(len(plot_data['importances'])))
    plot_data['x_values'] = x_values

    fig, axis = plt.subplots()
    # Make a bar chart
    plt.bar(plot_data['x_values'], plot_data['importances'], orientation = 'vertical', color='peru', edgecolor='black')

    # Tick labels for x axis
    plt.xticks(plot_data['x_values'], plot_data['feature_list'], rotation='vertical')

    # Axis labels and title
    plt.ylabel('Importance', fontweight='bold'); plt.xlabel('Feature', fontweight='bold');
    
    axis.set_axisbelow(True)
    plt.grid(linestyle='--')
    
    if outputFileName is not None:
        outputDir = '.'
        print('Output File Name:', outputFileName)
        plt.savefig(os.path.realpath(os.path.join(os.getcwd(), outputDir, outputFileName)), bbox_inches="tight", dpi=300)
        plt.close()
        

def get_plot_data(featuresAll, test_labels, predictions):
    df_prediction = test_labels.to_frame('actual')
    df_prediction.insert(0, 'prediction', predictions)
    df_prediction = featuresAll.join(df_prediction, how='inner')
    
    return df_prediction


def plot_boxplot_prediction_vs_measured(df_prediction_plot, dataMean, kind='HardwareVariables', dataset='OneWeek', appName='VASP', prefix='Figure5'):
    outputFileName = prefix+'_PredictionVsMeasured_'+appName+'_'+dataset+'_RFR_boxplot_'+kind+'.png'
    
    df_prediction_plot['error'] = abs(df_prediction_plot['prediction'] - df_prediction_plot['actual'])
    calculationText = '[Mean Absolute Error: '+str(round(np.mean(df_prediction_plot['error']), 1))+ ' W; Accuracy: '+str((1.0 - round(np.mean(df_prediction_plot['error']/df_prediction_plot['actual']), 2)) * 100)+'%]'

    df_prediction_plot['mean_test'] = df_prediction_plot['actual'].mean()
    df_prediction_plot['error_from_mean_test'] = abs(df_prediction_plot['prediction'] - df_prediction_plot['mean_test'])
    errorsRF_fromMean_test = round(np.mean(df_prediction_plot['error_from_mean_test']), 1)
    mape_fromMean_test = 100 * (df_prediction_plot['error_from_mean_test'] / df_prediction_plot['mean_test'])
    accuracyRF_fromMean_test = 100 - np.mean(mape_fromMean_test)
    calculationText = calculationText.split(']')[0] + '\nMean: '+str(round(df_prediction_plot['actual'].mean(),1))+'W; Std: '+str(round(df_prediction_plot['actual'].std(),1))+'W\nMean Absolute Error (wrt Mean Test): '+str(errorsRF_fromMean_test)+ 'W; Accuracy (wrt Mean Test): '+str(round(accuracyRF_fromMean_test,1))+' %]'
    
    df_prediction_plot['mean_all'] = dataMean
    df_prediction_plot['error_from_mean_all'] = abs(df_prediction_plot['prediction'] - df_prediction_plot['mean_all'])
    errorsRF_fromMean_all = round(np.mean(df_prediction_plot['error_from_mean_all']), 1)
    mape_fromMean_all = 100 * (df_prediction_plot['error_from_mean_all'] / df_prediction_plot['mean_all'])
    accuracyRF_fromMean_all = 100 - np.mean(mape_fromMean_all)
    calculationText = calculationText.split(']')[0] + '\nMean Absolute Error (wrt Mean All): '+str(errorsRF_fromMean_all)+ 'W; Accuracy (wrt Mean All): '+str(round(accuracyRF_fromMean_all,1))+' %]'

    df_prediction_plot['variation_from_mean_test'] = abs(df_prediction_plot['actual'] - df_prediction_plot['mean_test'])
    meanVariation_fromMean_test = round(np.mean(df_prediction_plot['variation_from_mean_test']), 1)
    mapVar_fromMean_test = 100 * (df_prediction_plot['variation_from_mean_test'] / df_prediction_plot['mean_test'])
    accuracyVar_fromMean_test = 100 - np.mean(mapVar_fromMean_test)
    calculationText = calculationText.split(']')[0] + '\nMean Absolute Variation (wrt Mean Test): '+str(meanVariation_fromMean_test)+ 'W; Mean Absolute Variation % (wrt Mean Test): '+str(round(np.mean(mapVar_fromMean_test),1))+' %]'

    df_prediction_plot['variation_from_mean_all'] = abs(df_prediction_plot['actual'] - df_prediction_plot['mean_all'])
    meanVariation_fromMean_all = round(np.mean(df_prediction_plot['variation_from_mean_all']), 1)
    mapVar_fromMean_all = 100 * (df_prediction_plot['variation_from_mean_all'] / df_prediction_plot['mean_all'])
    accuracyVar_fromMean_all = 100 - np.mean(mapVar_fromMean_all)
    calculationText = calculationText.split(']')[0] + '\nMean Absolute Variation (wrt Mean All): '+str(meanVariation_fromMean_all)+ 'W; Mean Absolute Variation % (wrt Mean All): '+str(round(np.mean(mapVar_fromMean_all),1))+' %]'
    print(calculationText)

    bins = list(range(90,310, 10))
    df_prediction_plot['actual_bin'] = pd.cut(df_prediction_plot['actual'], bins=bins, labels = bins[1:])
    df_prediction_plot['actual_bin'] = df_prediction_plot['actual_bin'].apply(str)
    
    plt.style.use('default')
    plt.rc('font', family='serif')
    fig, axis = plt.subplots()
    sns.boxplot(df_prediction_plot['actual_bin'], df_prediction_plot['prediction'], color='salmon')
    plt.xlabel('Measured Power (W)', fontweight='bold')
    plt.ylabel('Predicted Power (W)', fontweight='bold')
    plt.xticks(rotation=90)
    plt.ylim(90,300)
    plt.yticks(np.arange(100, 310, 10))
    plt.grid(linestyle='--')
    axis.set_axisbelow(True)
    outputFileName = outputFileName.split('.')[0] + '_bins100-300-10W.png'

    outputDir = '.'
    print('Output File Name:', outputFileName)
    plt.savefig(os.path.realpath(os.path.join(os.getcwd(), outputDir, outputFileName)), bbox_inches="tight", dpi=300)
    plt.close()