# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/hierarchical.ipynb (unless otherwise specified).

__all__ = ['bottom_up', 'HierarchicalStatsForecast']

# Cell
from typing import Callable, List, Optional, Tuple

import numpy as np
import pandas as pd
from sklearn.preprocessing import OneHotEncoder

from .core import StatsForecast

# Cell
def _aggregate_key(df: pd.DataFrame, keys: List[List[str]], agg_fn: Callable = np.sum):
    """Aggregates `df` according to `keys` using `agg_fn`."""
    df = df.copy()
    max_len_idx = np.argmax([len(key) for key in keys])
    bottom_comb = keys[max_len_idx]
    orig_cols = df.drop(labels=['ds', 'y'], axis=1).columns.to_list()
    df_keys = []
    for key in keys:
        df_key = df.groupby(key + ['ds'])['y'].apply(agg_fn).reset_index()
        df_key['unique_id'] = df_key[key].agg('_'.join, axis=1)
        if key == bottom_comb:
            bottom_keys = df_key['unique_id'].unique()
        df_keys.append(df_key)
    df_keys = pd.concat(df_keys)
    s_df = df_keys[['unique_id'] + bottom_comb].drop_duplicates().reset_index(drop=True)
    s_df = s_df.set_index('unique_id')
    s_df = s_df.fillna('agg')
    keys_cols = []
    for key in keys:
        key_col = '_'.join(key)
        s_df[key_col] = s_df[key].agg('_'.join, axis=1)
        keys_cols.append(key_col)
    y_df = df_keys[['unique_id', 'ds', 'y']].set_index('unique_id')
    #s_mat definition
    s_mat = _to_summing_matrix(s_df.loc[bottom_keys, keys_cols])
    return s_df[keys_cols], s_mat, bottom_keys, y_df

# Cell
def bottom_up(y_hat: np.ndarray, s_mat: np.ndarray):
    # size of y_hat = n_levels, horizon
    _, h = y_hat.shape
    n_levels, n_bottom = s_mat.shape
    g_mat = np.eye(n_bottom, n_levels, k=(n_levels - n_bottom), dtype=np.float32)
    s_g_mat = s_mat @ g_mat
    return np.matmul(s_g_mat, y_hat)

# Cell
class HierarchicalStatsForecast:

    def __init__(self, df: pd.DataFrame, keys: List[List[str]],
                 models: List, freq: str,
                 n_jobs: int = 1, ray_address: Optional[str] = None):
        self.s_df, self.s_mat, _,  y_df = _aggregate_key(df, keys=keys)
        self.fcst = StatsForecast(df=y_df, models=models, freq=freq,
                                  n_jobs=n_jobs, ray_address=ray_address,
                                  sort_df=False)

    def forecast(self, h: int, reconcile_fns: List[Callable] = [bottom_up],
                 xreg: Optional[pd.DataFrame] = None,
                 level: Optional[Tuple] = None):
        fcsts = self.fcst.forecast(h=h, xreg=xreg, level=level)
        model_names = fcsts.drop(columns=['ds'], axis=1).columns.to_list()
        for model_name in model_names:
            fcsts_model = fcsts[model_name].values.reshape(-1, h)
            for reconcile_fn in reconcile_fns:
                reconcile_fn_name = reconcile_fn.__name__
                fcsts_model = reconcile_fn(fcsts_model, self.s_mat)
                fcsts[f'{reconcile_fn_name}_{model_name}'] = fcsts_model.flatten()
        return fcsts

    def cross_validation(self, h: int, test_size: int,
                         input_size: Optional[int] = None,
                         reconcile_fns: List[Callable] = [bottom_up]):
        fcsts = self.fcst.cross_validation(h=h, test_size=test_size, input_size=input_size)
        model_names = fcsts.drop(columns=['ds', 'cutoff', 'y'], axis=1).columns.to_list()
        cutoffs = fcsts['cutoff'].unique()
        for model_name in model_names:
            for cutoff in cutoffs:
                cutoff_idx = fcsts['cutoff'] == cutoff
                fcsts_model = fcsts.loc[cutoff_idx, model_name].values.reshape(-1, h)
                for reconcile_fn in reconcile_fns:
                    reconcile_fn_name = reconcile_fn.__name__
                    fcsts_model = reconcile_fn(fcsts_model, self.s_mat)
                    fcsts.loc[cutoff_idx, f'{reconcile_fn_name}_{model_name}'] = fcsts_model.flatten()
        return fcsts
add missing hierarchical module