"""
This is the helper functions for evaluation purposes

"""
import numpy as np
from sklearn.metrics import confusion_matrix
import seaborn as sns
import matplotlib.pyplot as plt
import os
import pandas as pd


def compare_truth_pred(pred_file, truth_file, cut_off_outlier_thres=None, quiet_mode=False):
    """
    Read truth and pred from csv files, compute their mean-absolute-error and the mean-squared-error
    :param pred_file: full path to pred file
    :param truth_file: full path to truth file
    :return: mae and mse
    """
    if isinstance(pred_file, str):      # If input is a file name (original set up)
        pred = pd.read_csv(pred_file, header=None, sep=' ').values
        print(np.shape(pred))
        if np.shape(pred)[1] == 1:
            pred = pd.read_csv(pred_file, header=None, sep=',').values
        truth = pd.read_csv(truth_file, header=None, sep=' ').values
        print(np.shape(truth))
        if np.shape(truth)[1] == 1:
            truth = pd.read_csv(truth_file, header=None, sep=',').values
    elif isinstance(pred_file, np.ndarray):
        pred = pred_file
        truth = truth_file
    else:
        print('In the compare_truth_pred function, your input pred and truth is neither a file nor a numpy array')
    if not quiet_mode:
        print("in compare truth pred function in eval_help package, your shape of pred file is", np.shape(pred))
    
    # Getting the mean absolute error and the mean squared error
    mae = np.mean(np.abs(pred-truth), axis=1)
    mse = np.mean(np.square(pred-truth), axis=1)
    
    # When necessary you can choose to cut off the outliers here
    if cut_off_outlier_thres is not None:
        mse = mse[mse < cut_off_outlier_thres]
        mae = mae[mae < cut_off_outlier_thres]

    return mae, mse


def plotMSELossDistrib(pred_file, truth_file, save_dir='data/'):
    """
    Function to plot the MSE distribution histogram
    :param: pred_file: The Y prediction file
    :param: truth_file: The Y truth file
    :param: flags: The flags of the model/evaluation
    :param: save_dir: The directory to save the plot
    """
    mae, mse = compare_truth_pred(pred_file, truth_file)
    plt.figure(figsize=(12, 6))
    plt.hist(mse, bins=100)
    plt.xlabel('Mean Squared Error')
    plt.ylabel('cnt')
    plt.suptitle('(Avg MSE={:.4e}, 25%={:.3e}, 75%={:.3e})'.format(np.mean(mse), np.percentile(mse, 25), np.percentile(mse, 75)))
    if isinstance(pred_file, str):
        eval_model_str = pred_file.split('Ypred')[-1].split('.')[0]
    else:
        eval_model_str = 'MSE_unknon_name'
    plt.savefig(os.path.join(save_dir,
                            '{}.png'.format(eval_model_str)))
    print('(Avg MSE={:.4e})'.format(np.mean(mse)))
    return np.mean(mse)
