# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/src/methods.ipynb.

# %% auto 0
__all__ = ['BottomUp', 'BottomUpSparse', 'TopDown', 'TopDownSparse', 'MiddleOut', 'MiddleOutSparse', 'MinTrace', 'MinTraceSparse',
           'OptimalCombination', 'ERM']

# %% ../nbs/src/methods.ipynb 3
import warnings
from collections import OrderedDict
from concurrent.futures import ThreadPoolExecutor
from copy import deepcopy
from typing import Optional, Union

import numpy as np
from quadprog import solve_qp
from scipy import sparse

# %% ../nbs/src/methods.ipynb 4
from hierarchicalforecast.utils import (
    is_strictly_hierarchical,
    _ma_cov,
    _shrunk_covariance_schaferstrimmer_no_nans,
    _shrunk_covariance_schaferstrimmer_with_nans,
    _lasso,
)
from .probabilistic_methods import Normality, Bootstrap, PERMBU

# %% ../nbs/src/methods.ipynb 6
class HReconciler:
    fitted = False
    is_sparse_method = False
    insample = False
    P = None
    sampler = None

    def _get_sampler(
        self,
        intervals_method,
        S,
        P,
        y_hat,
        y_insample,
        y_hat_insample,
        W,
        sigmah,
        num_samples,
        seed,
        tags,
    ):
        if intervals_method == "normality":
            sampler = Normality(S=S, P=P, y_hat=y_hat, W=W, sigmah=sigmah, seed=seed)
        elif intervals_method == "permbu":
            sampler = PERMBU(
                S=S,
                P=P,
                y_hat=(S @ (P @ y_hat)),
                tags=tags,
                y_insample=y_insample,
                y_hat_insample=y_hat_insample,
                sigmah=sigmah,
                num_samples=num_samples,
                seed=seed,
            )
        elif intervals_method == "bootstrap":
            sampler = Bootstrap(
                S=S,
                P=P,
                y_hat=y_hat,
                y_insample=y_insample,
                y_hat_insample=y_hat_insample,
                num_samples=num_samples,
                seed=seed,
            )
        else:
            sampler = None
        return sampler

    def _reconcile(
        self,
        S: np.ndarray,
        P: np.ndarray,
        y_hat: np.ndarray,
        SP: np.ndarray = None,
        level: Optional[list[int]] = None,
        sampler: Optional[Union[Normality, PERMBU, Bootstrap]] = None,
    ):

        # Mean reconciliation
        res = {"mean": (S @ (P @ y_hat))}

        # Probabilistic reconciliation
        if (level is not None) and (sampler is not None):
            # Update results dictionary within
            # Vectorized quantiles
            quantiles = np.concatenate(
                [[(100 - lv) / 200, ((100 - lv) / 200) + lv / 100] for lv in level]
            )
            quantiles = np.sort(quantiles)
            res = sampler.get_prediction_quantiles(res, quantiles)

        return res

    def predict(
        self, S: np.ndarray, y_hat: np.ndarray, level: Optional[list[int]] = None
    ):
        """Predict using reconciler.

        Predict using fitted mean and probabilistic reconcilers.

        **Parameters:**<br>
        `S`: Summing matrix of size (`base`, `bottom`).<br>
        `y_hat`: Forecast values of size (`base`, `horizon`).<br>
        `level`: float list 0-100, confidence levels for prediction intervals.<br>

        **Returns:**<br>
        `y_tilde`: Reconciliated predictions.
        """
        if not self.fitted:
            raise Exception("This model instance is not fitted yet, Call fit method.")

        return self._reconcile(
            S=S, P=self.P, y_hat=y_hat, sampler=self.sampler, level=level
        )

    def sample(self, num_samples: int):
        """Sample probabilistic coherent distribution.

        Generates n samples from a probabilistic coherent distribution.
        The method uses fitted mean and probabilistic reconcilers, defined by
        the `intervals_method` selected during the reconciler's
        instantiation. Currently available: `normality`, `bootstrap`, `permbu`.

        **Parameters:**<br>
        `num_samples`: int, number of samples generated from coherent distribution.<br>

        **Returns:**<br>
        `samples`: Coherent samples of size (`num_series`, `horizon`, `num_samples`).
        """
        if not self.fitted:
            raise Exception("This model instance is not fitted yet, Call fit method.")
        if self.sampler is None:
            raise Exception(
                "This model instance does not have sampler. Call fit with `intervals_method`."
            )

        samples = self.sampler.get_samples(num_samples=num_samples)
        return samples

    def fit(self, *args, **kwargs):

        raise NotImplementedError("This method is not implemented yet.")

    def fit_predict(self, *args, **kwargs):

        raise NotImplementedError("This method is not implemented yet.")

    __call__ = fit_predict

# %% ../nbs/src/methods.ipynb 8
class BottomUp(HReconciler):
    """Bottom Up Reconciliation Class.
    The most basic hierarchical reconciliation is performed using an Bottom-Up strategy. It was proposed for
    the first time by Orcutt in 1968.
    The corresponding hierarchical \"projection\" matrix is defined as:
    $$\mathbf{P}_{\\text{BU}} = [\mathbf{0}_{\mathrm{[b],[a]}}\;|\;\mathbf{I}_{\mathrm{[b][b]}}]$$

    **Parameters:**<br>
    None

    **References:**<br>
    - [Orcutt, G.H., Watts, H.W., & Edwards, J.B.(1968). \"Data aggregation and information loss\". The American
    Economic Review, 58 , 773(787)](http://www.jstor.org/stable/1815532).
    """

    insample = False

    def _get_PW_matrices(self, S, idx_bottom):
        n_hiers, n_bottom = S.shape
        P = np.eye(n_bottom, n_hiers, n_hiers - n_bottom, np.float64)
        if getattr(self, "intervals_method", False) is None:
            W = None
        else:
            W = np.eye(n_hiers, dtype=np.float64)
        return P, W

    def fit(
        self,
        S: np.ndarray,
        y_hat: np.ndarray,
        idx_bottom: np.ndarray,
        y_insample: Optional[np.ndarray] = None,
        y_hat_insample: Optional[np.ndarray] = None,
        sigmah: Optional[np.ndarray] = None,
        intervals_method: Optional[str] = None,
        num_samples: Optional[int] = None,
        seed: Optional[int] = None,
        tags: Optional[dict[str, np.ndarray]] = None,
    ):
        """Bottom Up Fit Method.

        **Parameters:**<br>
        `S`: Summing matrix of size (`base`, `bottom`).<br>
        `y_hat`: Forecast values of size (`base`, `horizon`).<br>
        `idx_bottom`: Indices corresponding to the bottom level of `S`, size (`bottom`).<br>
        `level`: float list 0-100, confidence levels for prediction intervals.<br>
        `intervals_method`: Sampler for prediction intevals, one of `normality`, `bootstrap`, `permbu`.<br>
        `**sampler_kwargs`: Coherent sampler instantiation arguments.<br>

        **Returns:**<br>
        `self`: object, fitted reconciler.
        """
        self.intervals_method = intervals_method
        self.P, self.W = self._get_PW_matrices(S=S, idx_bottom=idx_bottom)
        self.sampler = self._get_sampler(
            S=S,
            P=self.P,
            W=self.W,
            y_hat=y_hat,
            y_insample=y_insample,
            y_hat_insample=y_hat_insample,
            sigmah=sigmah,
            intervals_method=intervals_method,
            num_samples=num_samples,
            seed=seed,
            tags=tags,
        )
        self.fitted = True
        return self

    def fit_predict(
        self,
        S: np.ndarray,
        y_hat: np.ndarray,
        idx_bottom: np.ndarray,
        y_insample: Optional[np.ndarray] = None,
        y_hat_insample: Optional[np.ndarray] = None,
        sigmah: Optional[np.ndarray] = None,
        level: Optional[list[int]] = None,
        intervals_method: Optional[str] = None,
        num_samples: Optional[int] = None,
        seed: Optional[int] = None,
        tags: Optional[dict[str, np.ndarray]] = None,
    ):
        """BottomUp Reconciliation Method.

        **Parameters:**<br>
        `S`: Summing matrix of size (`base`, `bottom`).<br>
        `y_hat`: Forecast values of size (`base`, `horizon`).<br>
        `idx_bottom`: Indices corresponding to the bottom level of `S`, size (`bottom`).<br>
        `level`: float list 0-100, confidence levels for prediction intervals.<br>
        `intervals_method`: Sampler for prediction intevals, one of `normality`, `bootstrap`, `permbu`.<br>
        `**sampler_kwargs`: Coherent sampler instantiation arguments.<br>

        **Returns:**<br>
        `y_tilde`: Reconciliated y_hat using the Bottom Up approach.
        """
        # Fit creates P, W and sampler attributes
        self.fit(
            S=S,
            y_hat=y_hat,
            y_insample=y_insample,
            y_hat_insample=y_hat_insample,
            sigmah=sigmah,
            intervals_method=intervals_method,
            num_samples=num_samples,
            seed=seed,
            tags=tags,
            idx_bottom=idx_bottom,
        )

        return self._reconcile(
            S=S, P=self.P, y_hat=y_hat, sampler=self.sampler, level=level
        )

    __call__ = fit_predict

# %% ../nbs/src/methods.ipynb 14
class BottomUpSparse(BottomUp):
    """BottomUpSparse Reconciliation Class.

    This is the implementation of a Bottom Up reconciliation using the sparse
    matrix approach. It works much more efficient on datasets with many time series.
    [makoren: At least I hope so, I only checked up until ~20k time series, and
    there's no real improvement, it would be great to check for smth like 1M time
    series, where the dense S matrix really stops fitting in memory]

    See the parent class for more details.
    """

    is_sparse_method = True

    def _get_PW_matrices(self, S, idx_bottom):
        n_hiers, n_bottom = S.shape
        P = sparse.eye(n_bottom, n_hiers, n_hiers - n_bottom, np.float64, "csr")
        if getattr(self, "intervals_method", False) is None:
            W = None
        else:
            W = sparse.eye(n_hiers, dtype=np.float64, format="csr")
        return P, W

# %% ../nbs/src/methods.ipynb 27
def _get_child_nodes(
    S: Union[np.ndarray, sparse.csr_matrix], tags: dict[str, np.ndarray]
):
    if isinstance(S, sparse.spmatrix):
        S = S.toarray()
    level_names = list(tags.keys())
    nodes = OrderedDict()
    for i_level, level in enumerate(level_names[:-1]):
        parent = tags[level]
        child = np.zeros_like(S)
        idx_child = tags[level_names[i_level + 1]]
        child[idx_child] = S[idx_child]
        nodes_level = {}
        for idx_parent_node in parent:
            parent_node = S[idx_parent_node]
            idx_node = child * parent_node.astype(bool)
            (idx_node,) = np.where(idx_node.sum(axis=1) > 0)
            nodes_level[idx_parent_node] = [idx for idx in idx_child if idx in idx_node]
        nodes[level] = nodes_level
    return nodes

# %% ../nbs/src/methods.ipynb 28
def _reconcile_fcst_proportions(
    S: np.ndarray,
    y_hat: np.ndarray,
    tags: dict[str, np.ndarray],
    nodes: dict[str, dict[int, np.ndarray]],
    idx_top: int,
):
    reconciled = np.zeros_like(y_hat)
    reconciled[idx_top] = y_hat[idx_top]
    level_names = list(tags.keys())
    for i_level, level in enumerate(level_names[:-1]):
        nodes_level = nodes[level]
        for idx_parent, idx_childs in nodes_level.items():
            fcst_parent = reconciled[idx_parent]
            childs_sum = y_hat[idx_childs].sum()
            for idx_child in idx_childs:
                reconciled[idx_child] = y_hat[idx_child] * fcst_parent / childs_sum
    return reconciled

# %% ../nbs/src/methods.ipynb 29
class TopDown(HReconciler):
    """Top Down Reconciliation Class.

    The Top Down hierarchical reconciliation method, distributes the total aggregate predictions and decomposes
    it down the hierarchy using proportions $\mathbf{p}_{\mathrm{[b]}}$ that can be actual historical values
    or estimated.

    $$\mathbf{P}=[\mathbf{p}_{\mathrm{[b]}}\;|\;\mathbf{0}_{\mathrm{[b][a,b\;-1]}}]$$
    **Parameters:**<br>
    `method`: One of `forecast_proportions`, `average_proportions` and `proportion_averages`.<br>

    **References:**<br>
    - [CW. Gross (1990). \"Disaggregation methods to expedite product line forecasting\". Journal of Forecasting, 9 , 233–254.
    doi:10.1002/for.3980090304](https://onlinelibrary.wiley.com/doi/abs/10.1002/for.3980090304).<br>
    - [G. Fliedner (1999). \"An investigation of aggregate variable time series forecast strategies with specific subaggregate
    time series statistical correlation\". Computers and Operations Research, 26 , 1133–1149.
    doi:10.1016/S0305-0548(99)00017-9](https://doi.org/10.1016/S0305-0548(99)00017-9).
    """

    def __init__(self, method: str):
        self.method = method
        self.insample = method in ["average_proportions", "proportion_averages"]

    def _get_PW_matrices(
        self,
        S: np.ndarray,
        y_hat: np.ndarray,
        y_insample: np.ndarray,
        tags: Optional[dict[str, np.ndarray]] = None,
    ):

        n_hiers, n_bottom = S.shape

        # Check if the data structure is strictly hierarchical.
        if tags is not None:
            if not is_strictly_hierarchical(S, tags):
                raise ValueError(
                    "Top-down reconciliation requires strictly hierarchical structures."
                )
            idx_top = int(S.sum(axis=1).argmax())
            levels_ = dict(sorted(tags.items(), key=lambda x: len(x[1])))
            idx_bottom = levels_[list(levels_)[-1]]
            y_btm = y_insample[idx_bottom]
        else:
            idx_top = 0
            y_btm = y_insample[(n_hiers - n_bottom) :]

        y_top = y_insample[idx_top]

        if self.method == "average_proportions":
            prop = np.mean(y_btm / y_top, axis=1)
        elif self.method == "proportion_averages":
            prop = np.mean(y_btm, axis=1) / np.mean(y_top)
        elif self.method == "forecast_proportions":
            raise Exception(f"Fit method not implemented for {self.method} yet")
        else:
            raise Exception(f"Unknown method {self.method}")

        P = np.zeros_like(
            S, np.float64
        ).T  # float 64 if prop is too small, happens with wiki2
        P[:, idx_top] = prop
        W = np.eye(n_hiers, dtype=np.float64)
        return P, W

    def fit(
        self,
        S,
        y_hat,
        y_insample: np.ndarray,
        y_hat_insample: Optional[np.ndarray] = None,
        sigmah: Optional[np.ndarray] = None,
        intervals_method: Optional[str] = None,
        num_samples: Optional[int] = None,
        seed: Optional[int] = None,
        tags: Optional[dict[str, np.ndarray]] = None,
        idx_bottom: Optional[np.ndarray] = None,
    ):
        """TopDown Fit Method.

        **Parameters:**<br>
        `S`: Summing matrix of size (`base`, `bottom`).<br>
        `y_hat`: Forecast values of size (`base`, `horizon`).<br>
        `tags`: Each key is a level and each value its `S` indices.<br>
        `y_insample`: Insample values of size (`base`, `insample_size`). Optional for `forecast_proportions` method.<br>
        `idx_bottom`: Indices corresponding to the bottom level of `S`, size (`bottom`).<br>
        `level`: float list 0-100, confidence levels for prediction intervals.<br>
        `intervals_method`: Sampler for prediction intevals, one of `normality`, `bootstrap`, `permbu`.<br>
        `**sampler_kwargs`: Coherent sampler instantiation arguments.<br>

        **Returns:**<br>
        `self`: object, fitted reconciler.
        """
        self.intervals_method = intervals_method
        self.P, self.W = self._get_PW_matrices(
            S=S, y_hat=y_hat, tags=tags, y_insample=y_insample
        )
        self.sampler = self._get_sampler(
            S=S,
            P=self.P,
            W=self.W,
            y_hat=y_hat,
            y_insample=y_insample,
            y_hat_insample=y_hat_insample,
            sigmah=sigmah,
            intervals_method=intervals_method,
            num_samples=num_samples,
            seed=seed,
            tags=tags,
        )
        self.fitted = True
        return self

    def fit_predict(
        self,
        S: np.ndarray,
        y_hat: np.ndarray,
        tags: dict[str, np.ndarray],
        idx_bottom: np.ndarray = None,
        y_insample: Optional[np.ndarray] = None,
        y_hat_insample: Optional[np.ndarray] = None,
        sigmah: Optional[np.ndarray] = None,
        level: Optional[list[int]] = None,
        intervals_method: Optional[str] = None,
        num_samples: Optional[int] = None,
        seed: Optional[int] = None,
    ):
        """Top Down Reconciliation Method.

        **Parameters:**<br>
        `S`: Summing matrix of size (`base`, `bottom`).<br>
        `y_hat`: Forecast values of size (`base`, `horizon`).<br>
        `tags`: Each key is a level and each value its `S` indices.<br>
        `y_insample`: Insample values of size (`base`, `insample_size`). Optional for `forecast_proportions` method.<br>
        `idx_bottom`: Indices corresponding to the bottom level of `S`, size (`bottom`).<br>
        `level`: float list 0-100, confidence levels for prediction intervals.<br>
        `intervals_method`: Sampler for prediction intevals, one of `normality`, `bootstrap`, `permbu`.<br>
        `**sampler_kwargs`: Coherent sampler instantiation arguments.<br>

        **Returns:**<br>
        `y_tilde`: Reconciliated y_hat using the Top Down approach.
        """
        if self.method == "forecast_proportions":
            idx_top = int(S.sum(axis=1).argmax())
            levels_ = dict(sorted(tags.items(), key=lambda x: len(x[1])))
            if level is not None:
                warnings.warn(
                    "Prediction intervals not implement for `forecast_proportions`"
                )
            nodes = _get_child_nodes(S=S, tags=levels_)
            reconciled = [
                _reconcile_fcst_proportions(
                    S=S,
                    y_hat=y_hat_[:, None],
                    tags=levels_,
                    nodes=nodes,
                    idx_top=idx_top,
                )
                for y_hat_ in y_hat.T
            ]
            reconciled = np.hstack(reconciled)
            return {"mean": reconciled}
        else:
            # Fit creates P, W and sampler attributes
            self.fit(
                S=S,
                y_hat=y_hat,
                y_insample=y_insample,
                y_hat_insample=y_hat_insample,
                sigmah=sigmah,
                intervals_method=intervals_method,
                num_samples=num_samples,
                seed=seed,
                tags=tags,
                idx_bottom=idx_bottom,
            )
            return self._reconcile(
                S=S, P=self.P, y_hat=y_hat, level=level, sampler=self.sampler
            )

    __call__ = fit_predict

# %% ../nbs/src/methods.ipynb 35
class TopDownSparse(TopDown):
    """TopDownSparse Reconciliation Class.

    This is an implementation of top-down reconciliation using the sparse matrix
    approach. It works much more efficiently on data sets with many time series.

    See the parent class for more details.
    """

    is_sparse_method = True

    def _get_PW_matrices(
        self,
        S: sparse.csr_matrix,
        y_hat: np.ndarray,
        y_insample: np.ndarray,
        tags: Optional[dict[str, np.ndarray]] = None,
    ):
        # Check if the data structure is strictly hierarchical.
        if tags is not None and not is_strictly_hierarchical(S, tags):
            raise ValueError(
                "Top-down reconciliation requires strictly hierarchical structures."
            )

        # Get the dimensions of the "summing" matrix.
        n_hiers, n_bottom = S.shape

        # Get the in-sample values of the top node and bottom nodes.
        y_top = y_insample[0]
        y_btm = y_insample[(n_hiers - n_bottom) :]

        # Calculate the disaggregation proportions.
        if self.method == "average_proportions":
            prop = np.mean(y_btm / y_top, 1)
        elif self.method == "proportion_averages":
            prop = np.mean(y_btm, 1) / np.mean(y_top)
        elif self.method == "forecast_proportions":
            raise Exception(f"Fit method not yet implemented for {self.method}.")
        else:
            raise Exception(f"{self.method} is an unknown disaggregation method.")

        # Instantiate and allocate the "projection" matrix to distribute the
        # disaggregated base forecast of the top node to the bottom nodes.
        P = sparse.csr_matrix(
            (
                prop,
                np.zeros_like(prop, np.uint8),
                np.arange(len(prop) + 1, dtype=np.min_scalar_type(n_bottom)),
            ),
            shape=(n_bottom, n_hiers),
            dtype=np.float64,
        )

        # Instantiate and allocate the "weight" matrix.
        if getattr(self, "intervals_method", False) is None:
            W = None
        else:
            W = sparse.eye(n_hiers, dtype=np.float64, format="csr")

        return P, W

# %% ../nbs/src/methods.ipynb 45
class MiddleOut(HReconciler):
    """Middle Out Reconciliation Class.

    This method is only available for **strictly hierarchical structures**. It anchors the base predictions
    in a middle level. The levels above the base predictions use the Bottom-Up approach, while the levels
    below use a Top-Down.

    **Parameters:**<br>
    `middle_level`: Middle level.<br>
    `top_down_method`: One of `forecast_proportions`, `average_proportions` and `proportion_averages`.<br>

    **References:**<br>
    - [Hyndman, R.J., & Athanasopoulos, G. (2021). \"Forecasting: principles and practice, 3rd edition:
    Chapter 11: Forecasting hierarchical and grouped series.\". OTexts: Melbourne, Australia. OTexts.com/fpp3
    Accessed on July 2022.](https://otexts.com/fpp3/hierarchical.html)

    """

    def __init__(self, middle_level: str, top_down_method: str):
        self.middle_level = middle_level
        self.top_down_method = top_down_method
        self.insample = top_down_method in [
            "average_proportions",
            "proportion_averages",
        ]

    def _get_PW_matrices(self, **kwargs):
        raise Exception("Not implemented")

    def fit(self, **kwargs):
        raise Exception("Not implemented")

    def predict(self, **kwargs):
        raise Exception("Not implemented")

    def fit_predict(
        self,
        S: np.ndarray,
        y_hat: np.ndarray,
        tags: dict[str, np.ndarray],
        y_insample: Optional[np.ndarray] = None,
        level: Optional[list[int]] = None,
        intervals_method: Optional[str] = None,
    ):
        """Middle Out Reconciliation Method.

        **Parameters:**<br>
        `S`: Summing matrix of size (`base`, `bottom`).<br>
        `y_hat`: Forecast values of size (`base`, `horizon`).<br>
        `tags`: Each key is a level and each value its `S` indices.<br>
        `y_insample`: Insample values of size (`base`, `insample_size`). Only used for `forecast_proportions`<br>

        **Returns:**<br>
        `y_tilde`: Reconciliated y_hat using the Middle Out approach.
        """
        if not is_strictly_hierarchical(S, tags):
            raise ValueError(
                "Middle out reconciliation requires strictly hierarchical structures."
            )
        if self.middle_level not in tags.keys():
            raise ValueError("You have to provide a `middle_level` in `tags`.")

        levels_ = dict(sorted(tags.items(), key=lambda x: len(x[1])))
        reconciled = np.full_like(y_hat, fill_value=np.nan)
        cut_nodes = levels_[self.middle_level]
        # bottom up reconciliation
        idxs_bu = []
        for node, idx_node in levels_.items():
            idxs_bu.append(idx_node)
            if node == self.middle_level:
                break
        idxs_bu = np.hstack(idxs_bu)
        # bottom up forecasts
        bu = BottomUp().fit_predict(
            S=np.fliplr(np.unique(S[idxs_bu], axis=1)),
            y_hat=y_hat[idxs_bu],
            idx_bottom=np.arange(len(idxs_bu))[-len(cut_nodes) :],
        )
        reconciled[idxs_bu] = bu["mean"]

        # top down
        child_nodes = _get_child_nodes(S, levels_)
        # parents contains each node in the middle out level
        # as key. The values of each node are the levels that
        # are conected to that node.
        parents = {node: {self.middle_level: np.array([node])} for node in cut_nodes}
        level_names = list(levels_.keys())
        for lv, lv_child in zip(level_names[:-1], level_names[1:]):
            # if lv is not part of the middle out to bottom
            # structure we continue
            if lv not in list(parents.values())[0].keys():
                continue
            for idx_middle_out in parents.keys():
                idxs_parents = parents[idx_middle_out].values()
                complete_idxs_child = []
                for idx_parent, idxs_child in child_nodes[lv].items():
                    if any(idx_parent in val for val in idxs_parents):
                        complete_idxs_child.append(idxs_child)
                parents[idx_middle_out][lv_child] = np.hstack(complete_idxs_child)

        for node, levels_node in parents.items():
            idxs_node = np.hstack(list(levels_node.values()))
            S_node = S[idxs_node]
            S_node = S_node[:, ~np.all(S_node == 0, axis=0)]
            counter = 0
            levels_node_ = deepcopy(levels_node)
            for lv_name, idxs_level in levels_node_.items():
                idxs_len = len(idxs_level)
                levels_node_[lv_name] = np.arange(counter, idxs_len + counter)
                counter += idxs_len
            td = TopDown(self.top_down_method).fit_predict(
                S=S_node,
                y_hat=y_hat[idxs_node],
                y_insample=y_insample[idxs_node] if y_insample is not None else None,
                tags=levels_node_,
            )
            reconciled[idxs_node] = td["mean"]
        return {"mean": reconciled}

    __call__ = fit_predict

# %% ../nbs/src/methods.ipynb 51
class MiddleOutSparse(MiddleOut):
    """MiddleOutSparse Reconciliation Class.

    This is an implementation of middle-out reconciliation using the sparse matrix
    approach. It works much more efficiently on data sets with many time series.

    See the parent class for more details.
    """

    # Although this is a sparse method, as we need to a dense representation of the
    # "summing" matrix for the required transformations in the fit_predict method
    # prior to bottom-up and top-down reconciliation, we can avoid a redundant
    # conversion.
    is_sparse_method = False

    def fit_predict(
        self,
        S: np.ndarray,
        y_hat: np.ndarray,
        tags: dict[str, np.ndarray],
        y_insample: Optional[np.ndarray] = None,
        level: Optional[list[int]] = None,
        intervals_method: Optional[str] = None,
    ) -> dict[str, np.ndarray]:
        # Check if the data structure is strictly hierarchical.
        if not is_strictly_hierarchical(S, tags):
            raise ValueError(
                "Middle-out reconciliation requires strictly hierarchical structures."
            )
        # Check if the middle level exists in the level to nodes mapping.
        if self.middle_level not in tags.keys():
            raise KeyError(f"{self.middle_level} is not a key in `tags`.")

        # Sort the levels by the number of nodes.
        levels = dict(sorted(tags.items(), key=lambda x: len(x[1])))
        # Allocate an array to store the reconciled point forecasts.
        y_tilde = np.full_like(y_hat, np.nan)
        # Find the nodes that constitute the middle level.
        cut_nodes = levels[self.middle_level]

        # Calculate the cut that separates the middle level from the lower levels.
        cut_idx = max(cut_nodes) + 1

        # Perform sparse bottom-up reconciliation from the middle level.
        y_tilde[:cut_idx, :] = BottomUpSparse().fit_predict(
            S=sparse.csr_matrix(np.fliplr(np.unique(S[:cut_idx, :], axis=1))),
            y_hat=y_hat[:cut_idx, :],
            idx_bottom=None,
        )["mean"]

        # Perform sparse top-down reconciliation from the middle level.
        for cut_node in cut_nodes:
            # Find the leaf nodes of the subgraph for the cut node.
            leaf_idx = np.flatnonzero(S[cut_node, :])
            # Find all the nodes in the subgraph for the cut node.
            sub_idx = np.hstack(
                (cut_node, cut_idx + np.flatnonzero((np.any(S[cut_idx:, leaf_idx], 1))))
            )

            # Construct the "tags" argument for the cut node.
            if self.insample:
                # It is not required for in-sample disaggregation methods.
                sub_tags = None
            else:
                # Disaggregating using forecast proportions requires the "tags" for
                # the subgraph.
                sub_tags = {}
                acc = 0
                for level_, nodes in levels.items():
                    # Find all the nodes in the subgraph for the level.
                    nodes = np.intersect1d(nodes, sub_idx, True)
                    # Get the number of nodes in the level.
                    n = len(nodes)
                    # Exclude any levels above the cut node or empty ones below.
                    if len(nodes) > 0:
                        sub_tags[level_] = np.arange(acc, n + acc)
                        acc += n

            # Perform sparse top-down reconciliation from the cut node.
            y_tilde[sub_idx, :] = TopDownSparse(self.top_down_method).fit_predict(
                S=sparse.csr_matrix(S[sub_idx[:, None], leaf_idx]),
                y_hat=y_hat[sub_idx, :],
                y_insample=y_insample[sub_idx, :] if y_insample is not None else None,
                tags=sub_tags,
            )["mean"]

        return {"mean": y_tilde}

    __call__ = fit_predict

# %% ../nbs/src/methods.ipynb 61
class MinTrace(HReconciler):
    """MinTrace Reconciliation Class.

    This reconciliation algorithm proposed by Wickramasuriya et al. depends on a generalized least squares estimator
    and an estimator of the covariance matrix of the coherency errors $\mathbf{W}_{h}$. The Min Trace algorithm
    minimizes the squared errors for the coherent forecasts under an unbiasedness assumption; the solution has a
    closed form.<br>

    $$
    \mathbf{P}_{\\text{MinT}}=\\left(\mathbf{S}^{\intercal}\mathbf{W}_{h}\mathbf{S}\\right)^{-1}
    \mathbf{S}^{\intercal}\mathbf{W}^{-1}_{h}
    $$

    **Parameters:**<br>
    `method`: str, one of `ols`, `wls_struct`, `wls_var`, `mint_shrink`, `mint_cov`.<br>
    `nonnegative`: bool, reconciled forecasts should be nonnegative?<br>
    `mint_shr_ridge`: float=2e-8, ridge numeric protection to MinTrace-shr covariance estimator.<br>
    `num_threads`: int=1, number of threads to use for solving the optimization problems (when nonnegative=True).

    **References:**<br>
    - [Wickramasuriya, S. L., Athanasopoulos, G., & Hyndman, R. J. (2019). \"Optimal forecast reconciliation for
    hierarchical and grouped time series through trace minimization\". Journal of the American Statistical Association,
    114 , 804–819. doi:10.1080/01621459.2018.1448825.](https://robjhyndman.com/publications/mint/).
    - [Wickramasuriya, S.L., Turlach, B.A. & Hyndman, R.J. (2020). \"Optimal non-negative
    forecast reconciliation". Stat Comput 30, 1167–1182,
    https://doi.org/10.1007/s11222-020-09930-0](https://robjhyndman.com/publications/nnmint/).
    """

    def __init__(
        self,
        method: str,
        nonnegative: bool = False,
        mint_shr_ridge: Optional[float] = 2e-8,
        num_threads: int = 1,
    ):
        self.method = method
        self.nonnegative = nonnegative
        self.insample = method in ["wls_var", "mint_cov", "mint_shrink"]
        if method == "mint_shrink":
            self.mint_shr_ridge = mint_shr_ridge
        self.num_threads = num_threads
        if not self.nonnegative and self.num_threads > 1:
            warnings.warn("`num_threads` is only used when `nonnegative=True`")

    def _get_PW_matrices(
        self,
        S: np.ndarray,
        y_hat: np.ndarray,
        y_insample: Optional[np.ndarray] = None,
        y_hat_insample: Optional[np.ndarray] = None,
        idx_bottom: Optional[list[int]] = None,
    ):
        # shape residuals_insample (n_hiers, obs)
        res_methods = ["wls_var", "mint_cov", "mint_shrink"]
        if self.method in res_methods and y_insample is None and y_hat_insample is None:
            raise ValueError(
                f"For methods {', '.join(res_methods)} you need to pass residuals"
            )
        n_hiers, n_bottom = S.shape
        n_aggs = n_hiers - n_bottom
        # Construct J and U.T
        J = np.concatenate(
            (np.zeros((n_bottom, n_aggs), dtype=np.float64), S[n_aggs:]), axis=1
        )
        Ut = np.concatenate((np.eye(n_aggs, dtype=np.float64), -S[:n_aggs]), axis=1)
        if self.method == "ols":
            W = np.eye(n_hiers)
            UtW = Ut
        elif self.method == "wls_struct":
            Wdiag = np.sum(S, axis=1, dtype=np.float64)
            UtW = Ut * Wdiag
            W = np.diag(Wdiag)
        elif (
            self.method in res_methods
            and y_insample is not None
            and y_hat_insample is not None
        ):
            # Residuals with shape (obs, n_hiers)
            residuals = (y_insample - y_hat_insample).T
            n, _ = residuals.shape

            # Protection: against overfitted model
            residuals_sum = np.sum(residuals, axis=0)
            zero_residual_prc = np.abs(residuals_sum) < 1e-4
            zero_residual_prc = np.mean(zero_residual_prc)
            if zero_residual_prc > 0.98:
                raise Exception(
                    f"Insample residuals close to 0, zero_residual_prc={zero_residual_prc}. Check `Y_df`"
                )

            if self.method == "wls_var":
                Wdiag = (
                    np.nansum(residuals**2, axis=0, dtype=np.float64)
                    / residuals.shape[0]
                )
                Wdiag += np.full(n_hiers, 2e-8, dtype=np.float64)
                W = np.diag(Wdiag)
                UtW = Ut * Wdiag
            elif self.method == "mint_cov":
                # Compute nans
                nan_mask = np.isnan(residuals.T)
                if np.any(nan_mask):
                    W = _ma_cov(residuals.T, ~nan_mask)
                else:
                    W = np.cov(residuals.T)

                UtW = Ut @ W
            elif self.method == "mint_shrink":
                # Compute nans
                nan_mask = np.isnan(residuals.T)
                # Compute shrunk empirical covariance
                if np.any(nan_mask):
                    W = _shrunk_covariance_schaferstrimmer_with_nans(
                        residuals.T, ~nan_mask, self.mint_shr_ridge
                    )
                else:
                    W = _shrunk_covariance_schaferstrimmer_no_nans(
                        residuals.T, self.mint_shr_ridge
                    )

                UtW = Ut @ W
        else:
            raise ValueError(f"Unknown reconciliation method {self.method}")

        try:
            P = (
                J
                - np.linalg.solve(
                    UtW[:, n_aggs:] @ Ut.T[n_aggs:] + UtW[:, :n_aggs],
                    UtW[:, n_aggs:] @ J.T[n_aggs:],
                ).T
                @ Ut
            )
        except np.linalg.LinAlgError:
            if self.method == "mint_shrink":
                raise Exception(
                    f"min_trace ({self.method}) is ill-conditioned. Increase the value of parameter 'mint_shr_ridge' or use another reconciliation method."
                )
            else:
                raise Exception(
                    f"min_trace ({self.method}) is ill-conditioned. Please use another reconciliation method."
                )

        return P, W

    def fit(
        self,
        S,
        y_hat,
        y_insample: Optional[np.ndarray] = None,
        y_hat_insample: Optional[np.ndarray] = None,
        sigmah: Optional[np.ndarray] = None,
        intervals_method: Optional[str] = None,
        num_samples: Optional[int] = None,
        seed: Optional[int] = None,
        tags: Optional[dict[str, np.ndarray]] = None,
        idx_bottom: Optional[np.ndarray] = None,
    ):
        """MinTrace Fit Method.

        **Parameters:**<br>
        `S`: Summing matrix of size (`base`, `bottom`).<br>
        `y_hat`: Forecast values of size (`base`, `horizon`).<br>
        `tags`: Each key is a level and each value its `S` indices.<br>
        `y_insample`: Insample values of size (`base`, `insample_size`). Optional for `forecast_proportions` method.<br>
        `idx_bottom`: Indices corresponding to the bottom level of `S`, size (`bottom`).<br>
        `level`: float list 0-100, confidence levels for prediction intervals.<br>
        `intervals_method`: Sampler for prediction intevals, one of `normality`, `bootstrap`, `permbu`.<br>
        `**sampler_kwargs`: Coherent sampler instantiation arguments.<br>

        **Returns:**<br>
        `self`: object, fitted reconciler.
        """
        self.y_hat = y_hat
        self.P, self.W = self._get_PW_matrices(
            S=S,
            y_hat=y_hat,
            y_insample=y_insample,
            y_hat_insample=y_hat_insample,
            idx_bottom=idx_bottom,
        )

        if self.nonnegative:
            _, n_bottom = S.shape
            W_inv = np.linalg.pinv(self.W)
            negatives = y_hat < 0
            if negatives.any():
                warnings.warn("Replacing negative forecasts with zero.")
                y_hat = np.copy(y_hat)
                y_hat[negatives] = 0.0
            # Quadratic progamming formulation
            # here we are solving the quadratic programming problem
            # formulated in the origial paper
            # https://robjhyndman.com/publications/nnmint/
            # The library quadprog was chosen
            # based on these benchmarks:
            # https://scaron.info/blog/quadratic-programming-in-python.html
            a = S.T @ W_inv
            G = a @ S
            try:
                _ = np.linalg.cholesky(G)
            except np.linalg.LinAlgError:
                raise Exception(
                    f"min_trace ({self.method}) is ill-conditioned. Try setting nonnegative=False or use another reconciliation method."
                )
            C = np.eye(n_bottom)
            b = np.zeros(n_bottom)
            # the quadratic programming problem
            # returns the forecasts of the bottom series
            if self.num_threads == 1:
                bottom_fcts = np.apply_along_axis(
                    lambda y_hat: solve_qp(G=G, a=a @ y_hat, C=C, b=b)[0],
                    axis=0,
                    arr=y_hat,
                )
            else:
                futures = []
                with ThreadPoolExecutor(self.num_threads) as executor:
                    for j in range(y_hat.shape[1]):
                        future = executor.submit(
                            solve_qp, G=G, a=a @ y_hat[:, j], C=C, b=b
                        )
                        futures.append(future)
                    bottom_fcts = np.hstack([f.result()[0][:, None] for f in futures])
            if not np.all(bottom_fcts > -1e-8):
                raise Exception("nonnegative optimization failed")
            # remove negative values close to zero
            bottom_fcts = np.clip(np.float32(bottom_fcts), a_min=0, a_max=None)
            self.y_hat = S @ bottom_fcts  # Hack

            # Overwrite P, W and sampler attributes with BottomUp's
            self.P, self.W = BottomUp()._get_PW_matrices(S=S, idx_bottom=idx_bottom)

        self.sampler = self._get_sampler(
            S=S,
            P=self.P,
            W=self.W,
            y_hat=y_hat,
            y_insample=y_insample,
            y_hat_insample=y_hat_insample,
            sigmah=sigmah,
            intervals_method=intervals_method,
            num_samples=num_samples,
            seed=seed,
            tags=tags,
        )
        self.fitted = True
        return self

    def fit_predict(
        self,
        S: np.ndarray,
        y_hat: np.ndarray,
        idx_bottom: np.ndarray = None,
        y_insample: Optional[np.ndarray] = None,
        y_hat_insample: Optional[np.ndarray] = None,
        sigmah: Optional[np.ndarray] = None,
        level: Optional[list[int]] = None,
        intervals_method: Optional[str] = None,
        num_samples: Optional[int] = None,
        seed: Optional[int] = None,
        tags: Optional[dict[str, np.ndarray]] = None,
    ):
        """MinTrace Reconciliation Method.

        **Parameters:**<br>
        `S`: Summing matrix of size (`base`, `bottom`).<br>
        `y_hat`: Forecast values of size (`base`, `horizon`).<br>
        `y_insample`: Insample values of size (`base`, `insample_size`). Only used by `wls_var`, `mint_cov`, `mint_shrink`<br>
        `y_hat_insample`: Insample fitted values of size (`base`, `insample_size`). Only used by `wls_var`, `mint_cov`, `mint_shrink`<br>
        `idx_bottom`: Indices corresponding to the bottom level of `S`, size (`bottom`).<br>
        `level`: float list 0-100, confidence levels for prediction intervals.<br>
        `sampler`: Sampler for prediction intevals, one of `normality`, `bootstrap`, `permbu`.<br>

        **Returns:**<br>
        `y_tilde`: Reconciliated y_hat using the MinTrace approach.
        """
        if self.nonnegative:
            if (level is not None) and intervals_method in ["bootstrap", "permbu"]:
                raise Exception(
                    "nonnegative reconciliation is not compatible with bootstrap forecasts"
                )
            if idx_bottom is None:
                raise Exception("idx_bottom needed for nonnegative reconciliation")

        # Fit creates P, W and sampler attributes
        self.fit(
            S=S,
            y_hat=y_hat,
            y_insample=y_insample,
            y_hat_insample=y_hat_insample,
            sigmah=sigmah,
            intervals_method=intervals_method,
            num_samples=num_samples,
            seed=seed,
            tags=tags,
            idx_bottom=idx_bottom,
        )

        return self._reconcile(
            S=S, P=self.P, y_hat=self.y_hat, level=level, sampler=self.sampler
        )

    __call__ = fit_predict

# %% ../nbs/src/methods.ipynb 67
class MinTraceSparse(MinTrace):
    """MinTraceSparse Reconciliation Class.

    This is the implementation of a subset of MinTrace features using the sparse
    matrix approach. It works much more efficient on datasets with many time series.

    See the parent class for more details.

    Currently supported:
    * Methods using diagonal W matrix, i.e. "ols", "wls_struct", "wls_var",
    * The standard MinT version (non-negative is not supported).

    Note: due to the numerical instability of the matrix inversion when creating the
    P matrix, the method is NOT guaranteed to give identical results to the non-sparse
    version.
    """

    is_sparse_method = True

    def _get_PW_matrices(
        self,
        S: Union[np.ndarray, sparse.spmatrix],
        y_hat: np.ndarray,
        y_insample: Optional[np.ndarray] = None,
        y_hat_insample: Optional[np.ndarray] = None,
        idx_bottom: Optional[list[int]] = None,
    ):
        # shape residuals_insample (n_hiers, obs)
        res_methods = ["wls_var", "mint_cov", "mint_shrink"]
        diag_only_methods = ["ols", "wls_struct", "wls_var"]

        if self.method not in diag_only_methods:
            raise NotImplementedError(
                "Only the methods with diagonal W are supported as sparse operations"
            )

        S = sparse.csr_matrix(S)

        if self.method in res_methods and y_insample is None and y_hat_insample is None:
            raise ValueError(
                f"For methods {', '.join(res_methods)} you need to pass residuals"
            )
        n_hiers, n_bottom = S.shape

        if self.method == "ols":
            W_diag = np.ones(n_hiers)
        elif self.method == "wls_struct":
            W_diag = S @ np.ones((n_bottom,))
        elif (
            self.method == "wls_var"
            and y_insample is not None
            and y_hat_insample is not None
        ):
            # Residuals with shape (obs, n_hiers)
            residuals = (y_insample - y_hat_insample).T
            n, _ = residuals.shape

            # Protection: against overfitted model
            residuals_sum = np.sum(residuals, axis=0)
            zero_residual_prc = np.abs(residuals_sum) < 1e-4
            zero_residual_prc = np.mean(zero_residual_prc)
            if zero_residual_prc > 0.98:
                raise Exception(
                    f"Insample residuals close to 0, zero_residual_prc={zero_residual_prc}. Check `Y_df`"
                )

            # Protection: cases where data is unavailable/nan
            # makoren: this masking stuff causes more harm than good, I found the results in the presence
            # of nan-s can often be rubbish, I'd argue it's better to fail than give rubbish results, here
            # the code is simply failing if it encounters nan in the variance vector.
            # masked_res = np.ma.array(residuals, mask=np.isnan(residuals))
            # covm = np.ma.cov(masked_res, rowvar=False, allow_masked=True).data

            W_diag = np.var(residuals, axis=0, ddof=1)
        else:
            raise ValueError(f"Unknown reconciliation method {self.method}")

        if any(W_diag < 1e-8):
            raise Exception(
                f"min_trace ({self.method}) needs covariance matrix to be positive definite."
            )

        if any(np.isnan(W_diag)):
            raise Exception(
                f"min_trace ({self.method}) needs covariance matrix to be positive definite (not nan)."
            )

        M = sparse.spdiags(np.reciprocal(W_diag), 0, W_diag.size, W_diag.size)
        R = sparse.csr_matrix(S.T @ M)

        # The implementation of P acting on a vector:
        def get_P_action(y):
            b = R @ y

            A = sparse.linalg.LinearOperator(
                (b.size, b.size), matvec=lambda v: R @ (S @ v)
            )

            x_tilde, exit_code = sparse.linalg.bicgstab(A, b, atol=1e-5)

            return x_tilde

        P = sparse.linalg.LinearOperator(
            (S.shape[1], y_hat.shape[0]), matvec=get_P_action
        )
        W = sparse.spdiags(W_diag, 0, W_diag.size, W_diag.size)

        return P, W

    def fit(
        self,
        S: sparse.csr_matrix,
        y_hat: np.ndarray,
        y_insample: Optional[np.ndarray] = None,
        y_hat_insample: Optional[np.ndarray] = None,
        sigmah: Optional[np.ndarray] = None,
        intervals_method: Optional[str] = None,
        num_samples: Optional[int] = None,
        seed: Optional[int] = None,
        tags: Optional[dict[str, np.ndarray]] = None,
        idx_bottom: Optional[np.ndarray] = None,
    ):
        # Clip the base forecasts if required to align them with their use in practice.
        if self.nonnegative:
            self.y_hat = np.clip(y_hat, 0, None)
        else:
            self.y_hat = y_hat
        # Get the reconciliation matrices.
        self.P, self.W = self._get_PW_matrices(
            S=S,
            y_hat=self.y_hat,
            y_insample=y_insample,
            y_hat_insample=y_hat_insample,
            idx_bottom=idx_bottom,
        )

        if self.nonnegative:
            # Get the number of leaf nodes.
            _, n_bottom = S.shape
            # Although it is now sufficient to ensure that all of the entries in P are
            # positive, as it is implemented as a linear operator for the iterative
            # method to solve the sparse linear system, we need to reconcile to find
            # if any of the coherent bottom level point forecasts are negative.
            y_tilde = self._reconcile(
                S=S, P=self.P, y_hat=self.y_hat, level=None, sampler=None
            )["mean"][-n_bottom:]
            # Find if any of the forecasts are negative.
            if np.any(y_tilde < 0):
                # Clip the negative forecasts.
                y_tilde = np.clip(y_tilde, 0, None)
                # Force non-negative coherence by overwriting the base forecasts with
                # the aggregated, clipped bottom level forecasts.
                self.y_hat = S @ y_tilde
                # Overwrite the attributes for the P and W matrices with those for
                # bottom-up reconciliation to force projection onto the non-negative
                # coherent subspace.
                self.P, self.W = BottomUpSparse()._get_PW_matrices(S=S, idx_bottom=None)

        # Get the sampler for probabilistic reconciliation.
        self.sampler = self._get_sampler(
            S=S,
            P=self.P,
            W=self.W,
            y_hat=self.y_hat,
            y_insample=y_insample,
            y_hat_insample=y_hat_insample,
            sigmah=sigmah,
            intervals_method=intervals_method,
            num_samples=num_samples,
            seed=seed,
            tags=tags,
        )
        # Set the instance as fitted.
        self.fitted = True
        return self

# %% ../nbs/src/methods.ipynb 78
class OptimalCombination(MinTrace):
    """Optimal Combination Reconciliation Class.

    This reconciliation algorithm was proposed by Hyndman et al. 2011, the method uses generalized least squares
    estimator using the coherency errors covariance matrix. Consider the covariance of the base forecast
    $\\textrm{Var}(\epsilon_{h}) = \Sigma_{h}$, the $\mathbf{P}$ matrix of this method is defined by:
    $$ \mathbf{P} = \\left(\mathbf{S}^{\intercal}\Sigma_{h}^{\dagger}\mathbf{S}\\right)^{-1}\mathbf{S}^{\intercal}\Sigma^{\dagger}_{h}$$
    where $\Sigma_{h}^{\dagger}$ denotes the variance pseudo-inverse. The method was later proven equivalent to
    `MinTrace` variants.

    **Parameters:**<br>
    `method`: str, allowed optimal combination methods: 'ols', 'wls_struct'.<br>
    `nonnegative`: bool, reconciled forecasts should be nonnegative?<br>

    **References:**<br>
    - [Rob J. Hyndman, Roman A. Ahmed, George Athanasopoulos, Han Lin Shang (2010). \"Optimal Combination Forecasts for
    Hierarchical Time Series\".](https://robjhyndman.com/papers/Hierarchical6.pdf).<br>
    - [Shanika L. Wickramasuriya, George Athanasopoulos and Rob J. Hyndman (2010). \"Optimal Combination Forecasts for
    Hierarchical Time Series\".](https://robjhyndman.com/papers/MinT.pdf).
    - [Wickramasuriya, S.L., Turlach, B.A. & Hyndman, R.J. (2020). \"Optimal non-negative
    forecast reconciliation". Stat Comput 30, 1167–1182,
    https://doi.org/10.1007/s11222-020-09930-0](https://robjhyndman.com/publications/nnmint/).
    """

    def __init__(self, method: str, nonnegative: bool = False, num_threads: int = 1):
        comb_methods = ["ols", "wls_struct"]
        if method not in comb_methods:
            raise ValueError(
                f'Optimal Combination class does not support method: "{method}"'
            )
        super().__init__(
            method=method, nonnegative=nonnegative, num_threads=num_threads
        )
        self.insample = False

# %% ../nbs/src/methods.ipynb 86
class ERM(HReconciler):
    """Optimal Combination Reconciliation Class.

    The Empirical Risk Minimization reconciliation strategy relaxes the unbiasedness assumptions from
    previous reconciliation methods like MinT and optimizes square errors between the reconciled predictions
    and the validation data to obtain an optimal reconciliation matrix P.

    The exact solution for $\mathbf{P}$ (`method='closed'`) follows the expression:
    $$\mathbf{P}^{*} = \\left(\mathbf{S}^{\intercal}\mathbf{S}\\right)^{-1}\mathbf{Y}^{\intercal}\hat{\mathbf{Y}}\\left(\hat{\mathbf{Y}}\hat{\mathbf{Y}}\\right)^{-1}$$

    The alternative Lasso regularized $\mathbf{P}$ solution (`method='reg_bu'`) is useful when the observations
    of validation data is limited or the exact solution has low numerical stability.
    $$\mathbf{P}^{*} = \\text{argmin}_{\mathbf{P}} ||\mathbf{Y}-\mathbf{S} \mathbf{P} \hat{Y} ||^{2}_{2} + \lambda ||\mathbf{P}-\mathbf{P}_{\\text{BU}}||_{1}$$

    **Parameters:**<br>
    `method`: str, one of `closed`, `reg` and `reg_bu`.<br>
    `lambda_reg`: float, l1 regularizer for `reg` and `reg_bu`.<br>

    **References:**<br>
    - [Ben Taieb, S., & Koo, B. (2019). Regularized regression for hierarchical forecasting without
    unbiasedness conditions. In Proceedings of the 25th ACM SIGKDD International Conference on Knowledge
    Discovery & Data Mining KDD '19 (p. 1337-1347). New York, NY, USA: Association for Computing Machinery.](https://doi.org/10.1145/3292500.3330976).<br>
    """

    def __init__(self, method: str, lambda_reg: float = 1e-2):
        self.method = method
        self.lambda_reg = lambda_reg
        self.insample = True

    def _get_PW_matrices(
        self,
        S: np.ndarray,
        y_hat: np.ndarray,
        y_insample: np.ndarray,
        y_hat_insample: np.ndarray,
        idx_bottom: np.ndarray,
    ):
        n_hiers, n_bottom = S.shape
        # y_hat_insample shape (n_hiers, obs)
        # remove obs with nan values
        nan_idx = np.isnan(y_hat_insample).any(axis=0)
        y_insample = y_insample[:, ~nan_idx]
        y_hat_insample = y_hat_insample[:, ~nan_idx]
        # only using h validation steps to avoid
        # computational burden
        # print(y_hat.shape)
        h = min(y_hat.shape[1], y_hat_insample.shape[1])
        y_hat_insample = y_hat_insample[:, -h:]  # shape (h, n_hiers)
        y_insample = y_insample[:, -h:]
        if self.method == "closed":
            B = np.linalg.inv(S.T @ S) @ S.T @ y_insample
            B = B.T
            P = np.linalg.pinv(y_hat_insample.T) @ B
            P = P.T
        elif self.method == "reg":
            X = np.kron(S, y_hat_insample.T)
            z = y_hat_insample.reshape(-1)

            if self.lambda_reg is None:
                lambda_reg = np.max(np.abs(X.T.dot(z)))
            else:
                lambda_reg = self.lambda_reg

            beta = _lasso(X, z, lambda_reg, max_iters=1000, tol=1e-4)
            P = beta.reshape(S.shape).T
        elif self.method == "reg_bu":
            X = np.kron(S, y_hat_insample.T)
            Pbu = np.zeros_like(S)
            Pbu[idx_bottom] = S[idx_bottom]
            z = y_hat_insample.reshape(-1) - X @ Pbu.reshape(-1)

            if self.lambda_reg is None:
                lambda_reg = np.max(np.abs(X.T.dot(z)))
            else:
                lambda_reg = self.lambda_reg

            beta = _lasso(X, z, lambda_reg, max_iters=1000, tol=1e-4)
            P = beta + Pbu.reshape(-1)
            P = P.reshape(S.shape).T
        else:
            raise ValueError(f"Unknown reconciliation method {self.method}")

        W = np.eye(n_hiers, dtype=np.float64)

        return P, W

    def fit(
        self,
        S,
        y_hat,
        y_insample,
        y_hat_insample,
        sigmah: Optional[np.ndarray] = None,
        intervals_method: Optional[str] = None,
        num_samples: Optional[int] = None,
        seed: Optional[int] = None,
        tags: Optional[dict[str, np.ndarray]] = None,
        idx_bottom: Optional[np.ndarray] = None,
    ):
        """ERM Fit Method.

        **Parameters:**<br>
        `S`: Summing matrix of size (`base`, `bottom`).<br>
        `y_hat`: Forecast values of size (`base`, `horizon`).<br>
        `y_insample`: Train values of size (`base`, `insample_size`).<br>
        `y_hat_insample`: Insample train predictions of size (`base`, `insample_size`).<br>
        `idx_bottom`: Indices corresponding to the bottom level of `S`, size (`bottom`).<br>
        `level`: float list 0-100, confidence levels for prediction intervals.<br>
        `intervals_method`: Sampler for prediction intevals, one of `normality`, `bootstrap`, `permbu`.<br>
        `**sampler_kwargs`: Coherent sampler instantiation arguments.<br>

        **Returns:**<br>
        `self`: object, fitted reconciler.
        """
        self.P, self.W = self._get_PW_matrices(
            S=S,
            y_hat=y_hat,
            y_insample=y_insample,
            y_hat_insample=y_hat_insample,
            idx_bottom=idx_bottom,
        )
        self.sampler = self._get_sampler(
            S=S,
            P=self.P,
            W=self.W,
            y_hat=y_hat,
            y_insample=y_insample,
            y_hat_insample=y_hat_insample,
            sigmah=sigmah,
            intervals_method=intervals_method,
            num_samples=num_samples,
            seed=seed,
            tags=tags,
        )
        self.fitted = True
        return self

    def fit_predict(
        self,
        S: np.ndarray,
        y_hat: np.ndarray,
        idx_bottom: np.ndarray = None,
        y_insample: Optional[np.ndarray] = None,
        y_hat_insample: Optional[np.ndarray] = None,
        sigmah: Optional[np.ndarray] = None,
        level: Optional[list[int]] = None,
        intervals_method: Optional[str] = None,
        num_samples: Optional[int] = None,
        seed: Optional[int] = None,
        tags: Optional[dict[str, np.ndarray]] = None,
    ):
        """ERM Reconciliation Method.

        **Parameters:**<br>
        `S`: Summing matrix of size (`base`, `bottom`).<br>
        `y_hat`: Forecast values of size (`base`, `horizon`).<br>
        `y_insample`: Train values of size (`base`, `insample_size`).<br>
        `y_hat_insample`: Insample train predictions of size (`base`, `insample_size`).<br>
        `idx_bottom`: Indices corresponding to the bottom level of `S`, size (`bottom`).<br>
        `level`: float list 0-100, confidence levels for prediction intervals.<br>
        `intervals_method`: Sampler for prediction intevals, one of `normality`, `bootstrap`, `permbu`.<br>

        **Returns:**<br>
        `y_tilde`: Reconciliated y_hat using the ERM approach.
        """
        # Fit creates P, W and sampler attributes
        self.fit(
            S=S,
            y_hat=y_hat,
            y_insample=y_insample,
            y_hat_insample=y_hat_insample,
            sigmah=sigmah,
            intervals_method=intervals_method,
            num_samples=num_samples,
            seed=seed,
            tags=tags,
            idx_bottom=idx_bottom,
        )

        return self._reconcile(
            S=S, P=self.P, y_hat=y_hat, level=level, sampler=self.sampler
        )

    __call__ = fit_predict
