import numpy as np
from sklearn.metrics import mean_squared_error


def mase(n, seas, h, y, f):
    return np.mean(((n - seas) / h
                    * (np.sum(np.abs(y[n:n + h, :] - f), axis=0)
                       / np.sum(np.abs(y[seas:n, :] - y[:n - seas, :]), axis=0))))


def calculate_metrics(pred_samples,
                      groups):
    pred_s0 = pred_samples.shape[0]
    pred_s1 = pred_samples.shape[1]
    pred_s2 = pred_samples.shape[2]

    pred_samples = pred_samples.reshape(pred_s0, pred_s1 * pred_s2, order='F')

    seasonality = groups['seasonality']
    h = groups['h']

    n = groups['predict']['n']
    s = groups['predict']['s']
    y_f = groups['predict']['data'].reshape(s, n).T

    y_all_g = {}
    f_all_g = {}

    mase_ = {}
    rmse_ = {}

    # Bottom
    y_all_g['bottom'] = y_f
    f_all_g['bottom'] = np.mean(pred_samples, axis=0).reshape(s, n).T[n - h:n, :]

    mase_['bottom'] = np.round(mase(n=n - h,
                                    seas=seasonality,
                                    h=h,
                                    y=y_f,
                                    f=np.mean(pred_samples, axis=0).reshape(s, n).T[n - h:n, :]), 3)
    rmse_['bottom'] = np.round(
        mean_squared_error(y_f[n - h:n, :], np.mean(pred_samples, axis=0).reshape(s, n).T[n - h:n, :], squared=False),
        3)

    # Total
    y_all_g['total'] = np.sum(y_f, axis=1).reshape(-1, 1)
    f_all_g['total'] = np.sum(np.mean(pred_samples, axis=0).reshape(s, n).T[n - h:n, :], axis=1).reshape(-1, 1)

    mase_['total'] = np.round(mase(n=n - h,
                                   seas=seasonality,
                                   h=h,
                                   y=np.sum(y_f, axis=1).reshape(-1, 1),
                                   f=np.sum(np.mean(pred_samples, axis=0).reshape(s, n).T[n - h:n, :], axis=1).reshape(
                                       -1, 1))
                              , 3)
    rmse_['total'] = np.round(mean_squared_error(np.sum(y_f, axis=1).reshape(-1, 1)[n - h:n, :],
                                                 np.sum(np.mean(pred_samples, axis=0).reshape(s, n).T[n - h:n, :],
                                                        axis=1).reshape(-1, 1),
                                                 squared=False), 3)

    # Groups
    idx_dict_new = {}
    for group in list(groups['predict']['groups_names'].keys()):
        y_g = np.zeros((groups['predict']['n'], groups['predict']['groups_names'][group].shape[0]))
        f_g = np.zeros((h, groups['predict']['groups_names'][group].shape[0]))

        for idx, name in enumerate(groups['predict']['groups_names'][group]):
            g_n = groups['predict']['groups_n'][group]

            idx_dict_new[name] = np.where(groups['predict']['groups_idx'][group] == idx, 1, 0)

            y_g[:, idx] = np.sum(idx_dict_new[name] * y_f, axis=1)
            f_g[:, idx] = np.sum(idx_dict_new[name] * np.mean(pred_samples, axis=0).reshape(s, n).T, axis=1)[n - h:n]

        y_all_g[group] = np.sum(y_g, axis=1).reshape(-1, 1)
        f_all_g[group] = np.sum(f_g, axis=1).reshape(-1, 1)

        mase_[group] = np.round(mase(n=n - h,
                                     seas=seasonality,
                                     h=h,
                                     y=y_g,
                                     f=f_g)
                                , 3)

        rmse_[group] = np.round(mean_squared_error(y_g[n - h:n, :], f_g, squared=False), 3)

    # All
    y_all = np.concatenate([y_all_g[x] for x in y_all_g], 1)
    f_all = np.concatenate([f_all_g[x] for x in f_all_g], 1)

    mase_['all'] = np.round(mase(n=n - h,
                                 seas=seasonality,
                                 h=h,
                                 y=y_all,
                                 f=f_all), 3)
    rmse_['all'] = np.round(mean_squared_error(y_all[n - h:n, :], f_all, squared=False), 3)

    results = {}
    results['mase'] = mase_
    results['rmse'] = rmse_
    return results


