import numpy as np
import math
from sklearn.metrics import r2_score


def compute_error(ground_truth, prediction):
    rmse = r_m_s_e_error(ground_truth, prediction)
    r2 = r_2_error(ground_truth, prediction)
    mape = m_a_p_e_error(ground_truth, prediction)
    return rmse, mape, r2


def r_2_error(ground_truth, prediction):
    return r2_score(ground_truth[~np.isnan(ground_truth)], prediction[~np.isnan(ground_truth)])


def m_a_p_e_error(ground_truth, prediction):
    return np.nanmean(np.abs((ground_truth - prediction) / ground_truth)) * 100


def r_m_s_e_error(ground_truth, prediction):
    mse = np.nanmean((ground_truth - prediction) ** 2)
    return math.sqrt(mse)


# class Norm:
#
#     def __init__(self, x):
#         self.mean = np.nanmean(x, axis=0)
#         self.std = np.nanstd(x, axis=0)
#
#     def transform(self, x):
#         return np.divide(x - self.mean, self.std, out=np.zeros_like(x - self.mean), where=self.std != 0.0)
#
#
# def normalize_mat(input_mat, if_retain_last_dim=True):
#     """ normalize the input feature matrix
#
#     Params:
#         feature_mat: (n_times, n_features, n_rows, n_cols)
#         if_retain_last_dim: if keeping the last dimension as original
#     Return:
#         norm_mat: (n_times, n_features, n_rows, n_cols)
#     """
#
#     n_times, n_features, n_rows, n_cols = input_mat.shape
#     x_2d = np.moveaxis(input_mat, 1, -1)  # => (n_times, n_rows, n_cols, n_features)
#     x_2d = x_2d.reshape(-1, n_features)  # => (n_samples, n_features)
#     norm = Norm(x_2d)
#     norm_mat = norm.transform(x_2d)
#
#     norm_mat = norm_mat.reshape(n_times, n_rows, n_cols, n_features)  # => (n_times, n_rows, n_cols, n_features)
#     if if_retain_last_dim:
#         norm_mat = np.moveaxis(norm_mat, -1, 1)  # => (n_times, n_features, n_rows, n_cols)
#
#     # print('Shape of the normalized matrix = {}'.format(norm_mat.shape))
#     return norm_mat
