import numpy as np


def MSE(y_true: np.ndarray, y_pred: np.ndarray):
    """Mean squared error"""
    return np.mean((y_true - y_pred) ** 2)


def MAE(y_true: np.ndarray, y_pred: np.ndarray):
    """Mean absolute error"""
    return np.mean(np.abs(y_true - y_pred))


def MASE(y_true: np.ndarray, y_pred: np.ndarray, freq: str = "h"):
    """Mean absolute scaled error"""
    DEFAULT_SEASONALITIES = {
        "S": 3600,  # 1 hour
        "s": 3600,  # 1 hour
        "T": 1440,  # 1 day
        "min": 1440,  # 1 day
        "H": 24,  # 1 day
        "h": 24,  # 1 day
        "D": 1,  # 1 day
        "W": 1,  # 1 week
        "M": 12,
        "ME": 12,
        "B": 5,
        "Q": 4,
        "QE": 4,
    }
    # seasonality = DEFAULT_SEASONALITIES[freq]
    if len(y_true.shape) == 3:  # num_batch, bs, seq_len
        y_t = y_true[:, :, 1:] - y_true[:, :, :-1]
    else:  # num_seq, seq_len
        y_t = y_true[:, 1:] - y_true[:, :-1]
    return np.mean(np.abs(y_true - y_pred) / (np.mean(np.abs(y_t)) + 1e-5))


def MAPE(y_true: np.ndarray, y_pred: np.ndarray):
    """Mean absolute percentage error"""
    return np.mean(np.abs(y_true - y_pred) / (y_true + 1e-5))


def RMSE(y_true: np.ndarray, y_pred: np.ndarray):
    """Root mean squared error"""
    return np.sqrt(MSE(y_true, y_pred))


def NRMSE(y_true: np.ndarray, y_pred: np.ndarray):
    """Normalized root mean squared error"""
    return RMSE(y_true, y_pred) / (np.max(y_true) - np.min(y_true) + 1e-5)


def SMAPE(y_true: np.ndarray, y_pred: np.ndarray):
    """Symmetric mean absolute percentage error"""
    return np.mean(
        2.0 * np.abs(y_true - y_pred) / (np.abs(y_true) + np.abs(y_pred) + 1e-5)
    )


def MSIS(y_true: np.ndarray, y_pred: np.ndarray, alpha: float = 0.05):
    """Mean scaled interval score"""
    q1 = np.percentile(y_true, 100 * alpha / 2)
    q2 = np.percentile(y_true, 100 * (1 - alpha / 2))
    denominator = q2 - q1
    penalties = 2 * ((y_true < q1) * (q1 - y_pred) + (y_true > q2) * (y_pred - q2))
    return np.mean(np.abs(y_true - y_pred) / (denominator + 1e-5)) + np.mean(
        penalties / (denominator + 1e-5)
    )


def ND(y_true: np.ndarray, y_pred: np.ndarray):
    """Normalized deviation"""
    return np.mean(np.abs(y_true - y_pred)) / (np.mean(y_true) + 1e-5)


def MWSQ(y_true: np.ndarray, y_pred: np.ndarray, quantiles: np.ndarray):
    """
    Mean weighted squared quantile loss
    y_true: (num_seq, n_var, seq_len)
    y_pred: (q, num_seq, n_var, seq_len)
    quantiles: (q, )
    """

    y_true = np.expand_dims(y_true, axis=0)  # (1, num_seq, n_var, seq_len)
    diff = y_true - y_pred  # (q, num_seq, n_var, seq_len)
    quantiles = np.expand_dims(
        np.expand_dims(np.expand_dims(quantiles, axis=-1), axis=-1), axis=-1
    )  # (q, 1, 1, 1)
    pinball = np.maximum(quantiles * diff, (quantiles - 1) * diff)  # (num_seq, n_var, seq_len, q)
    mwsq = np.mean(pinball ** 2)  # (num_seq, n_var, seq_len)
    return mwsq


def CRPS(y_true: np.ndarray, y_pred: np.ndarray, quantiles: np.ndarray):
    """
    Continuous ranked probability score
    y_true: (num_seq, n_var, seq_len)
    y_pred: (q, num_seq, n_var, seq_len)
    quantiles: (q, )
    """
    y_true = np.expand_dims(y_true, axis=0)  # (1, num_seq, n_var, seq_len)
    diff = y_true - y_pred  # (q, num_seq, n_var, seq_len)
    quantiles = np.expand_dims(
        np.expand_dims(np.expand_dims(quantiles, axis=-1), axis=-1), axis=-1
    )  # (q, 1, 1, 1)
    pinball = np.maximum(quantiles * diff, (quantiles - 1) * diff)  # (num_seq, n_var, seq_len, q)
    return np.mean(pinball)
