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

__all__ = ['HierarchicalReconciliation']

# Cell
from functools import partial
from inspect import signature
from typing import Callable, List

import numpy as np
import pandas as pd

# Internal Cell
def _build_fn_name(fn) -> str:
    fn_name = type(fn).__name__
    func_params = fn.__dict__
    func_params = [f'{name}-{value}' for name, value in func_params.items()]
    if func_params:
        fn_name += '_' + '_'.join(func_params)
    return fn_name

# Cell
class HierarchicalReconciliation:

    def __init__(self, reconcilers: List[Callable]):
        self.reconcilers = reconcilers

    def reconcile(self, Y_h: pd.DataFrame, Y_df: pd.DataFrame, S: pd.DataFrame):
        """Reconcile base forecasts.

            Parameters
            ----------
            Y_h: pd.DataFrame
                Base forecasts with columns ['ds']
                and models to reconcile indexed by 'unique_id'.
            Y_df: pd.DataFrame
                Training set of base time series with columns
                ['ds', 'y'] indexed by 'unique_id'
                If a function of `self.reconcile_fns` receives
                residuals, `Y_df` must include them as columns.
            S: pd.DataFrame
                Summing matrix of size (hierarchies, bottom).
        """
        drop_cols = ['ds', 'y'] if 'y' in Y_h.columns else ['ds']
        model_names = Y_h.drop(columns=drop_cols, axis=1).columns.to_list()
        common_vals = dict(
            y = Y_df.pivot(columns='ds', values='y').loc[S.index].values,
            S = S.values,
            idx_bottom = [S.index.get_loc(col) for col in S.columns]
        )
        fcsts = Y_h.copy()
        for reconcile_fn in self.reconcilers:
            reconcile_fn_name = _build_fn_name(reconcile_fn)
            has_res = 'residuals' in signature(reconcile_fn).parameters
            for model_name in model_names:
                y_hat_model = Y_h.pivot(columns='ds', values=model_name).loc[S.index].values
                if has_res:
                    if model_name in Y_df:
                        common_vals['residuals'] = Y_df.pivot(columns='ds', values=model_name).loc[S.index].values.T
                    else:
                        # some methods have the residuals argument
                        # but they don't need them
                        # ej MinTrace(method='ols')
                        common_vals['residuals'] = None
                kwargs = [key for key in signature(reconcile_fn).parameters if key in common_vals.keys()]
                kwargs = {key: common_vals[key] for key in kwargs}
                fcsts_model = reconcile_fn(y_hat=y_hat_model, **kwargs)
                fcsts[f'{model_name}/{reconcile_fn_name}'] = fcsts_model.flatten()
                if has_res:
                    del common_vals['residuals']
        return fcsts