# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License").
# You may not use this file except in compliance with the License.
# A copy of the License is located at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# or in the "license" file accompanying this file. This file is distributed
# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
# express or implied. See the License for the specific language governing
# permissions and limitations under the License.

from typing import Dict, List, Optional, Tuple

import mxnet as mx
import numpy as np

from gluonts_customized.core.component import Type, validated
from gluonts_customized.itertools import prod
from gluonts_customized.mx.model.deepar import DeepAREstimator
from gluonts_customized.mx.model.deepar._network import DeepARPredictionNetwork
from gluonts_customized.mx.model.deepvar_hierarchical._estimator import projection_mat
from gluonts_customized.mx.model.deepvar_hierarchical._network import coherency_error
from gluonts_customized.mx.distribution import Distribution, EmpiricalDistribution
from gluonts_customized.mx import Tensor
from gluonts_customized.mx.distribution import TransformedPiecewiseLinear

from gluonts_customized.nursery.temporal_hierarchical_forecasting.utils import utils
from gluonts_customized.nursery.temporal_hierarchical_forecasting.model.cop_deepar import (
    gluonts_fixes,
    gnn,
)


def reconcile_samples(
    reconciliation_mat: Tensor,
    samples: Tensor,
    non_negative: bool = False,
    num_iters: int = 10,
) -> Tensor:
    if not non_negative:
        return mx.nd.dot(samples, reconciliation_mat, transpose_b=True)
    else:
        # Dykstra's projection method: Projection onto the intersection of convex sets.
        x = samples
        p = mx.nd.zeros_like(x)
        q = mx.nd.zeros_like(x)
        for _ in range(num_iters):
            # Projection onto the non-negative orthant.
            y = mx.nd.relu(x + p)
            p = x + p - y

            # Projection onto the null space.
            x = mx.nd.dot(y + q, reconciliation_mat, transpose_b=True)
            q = y + q - x

        return x


class COPNetwork(mx.gluon.HybridBlock):
    @validated()
    def __init__(
        self,
        estimators: List[DeepAREstimator],
        prediction_length: int,
        temporal_hierarchy: utils.TemporalHierarchy,
        do_reconciliation: bool,
        dtype: Type,
        use_gnn: bool,
        use_mlp: bool,
        adj_mat_option: str,
        non_negative: bool = False,
        naive_reconciliation: bool = False,
        prediction: bool = False,
        loss_function: str = "crps_univariate",
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)

        self.prediction_length = prediction_length
        self.temporal_hierarchy = temporal_hierarchy
        self.use_gnn = use_gnn
        self.use_mlp = use_mlp
        self.adj_mat_option = adj_mat_option
        self.do_reconciliation = do_reconciliation
        self.non_negative = non_negative
        self.loss_function = loss_function
        self.dtype = dtype

        if naive_reconciliation:
            M = utils.naive_reconcilation_mat(
                self.temporal_hierarchy.agg_mat, self.temporal_hierarchy.nodes
            )
        else:
            M = projection_mat(S=self.temporal_hierarchy.agg_mat)
        self.M = mx.nd.array(M)

        self.estimators = estimators

        self.models = []
        with self.name_scope():
            for estimator in estimators:
                if not prediction:
                    self.network = estimator.create_training_network()
                else:
                    self.network = gluonts_fixes.create_prediction_network(
                        estimator
                    )

                self.register_child(self.network)
                self.models.append(self.network)

            if self.use_gnn:
                # GNN Layer: Do message passing for `L-1` times, where `L` is the number of levels of the hierarchy.
                self.gnn = gnn.GNN(
                    units=self.estimators[0].num_cells,
                    num_layers=len(self.temporal_hierarchy.agg_multiples) - 1,
                    adj_matrix=mx.nd.array(
                        self.temporal_hierarchy.adj_mat(
                            option=self.adj_mat_option
                        )
                    ),
                    use_mlp=self.use_mlp,
                )

    def get_target_related_feat_at_agg_level(
        self,
        agg_level: int,
        past_target: Tensor,
        past_observed_values: Tensor,
        past_is_pad: Tensor,
        future_target: Optional[Tensor] = None,
        future_observed_values: Optional[Tensor] = None,
    ) -> Dict:
        """
        Aggregate target at the given aggregate level along with updating
        observed value and pad indicators.

        :param agg_level:
        :param past_target:
        :param past_observed_values:
        :param past_is_pad:
        :param future_target:
        :param future_observed_values:
        :return:
        """
        agg_multiple = self.temporal_hierarchy.agg_multiples[agg_level]

        # Truncating the history length of the base time series to the nearest multiple.
        base_history_length = (
            past_target.shape[1] // agg_multiple
        ) * agg_multiple

        past_target_agg = (
            utils.agg_series(
                past_target.slice_axis(
                    axis=1, begin=-base_history_length, end=None
                ),
                agg_multiple=agg_multiple,
            )
            .squeeze(axis=-1)
            .slice_axis(
                axis=1, begin=-self.models[agg_level].history_length, end=None
            )
        )

        past_is_pad_agg = (
            utils.agg_series(
                past_is_pad.slice_axis(
                    axis=1, begin=-base_history_length, end=None
                ),
                agg_multiple=agg_multiple,
            )
            .squeeze(axis=-1)
            .slice_axis(
                axis=1, begin=-self.models[agg_level].history_length, end=None
            )
        )

        past_is_pad_agg = mx.nd.where(
            past_is_pad_agg == 0.0,
            mx.nd.zeros_like(past_is_pad_agg),
            mx.nd.ones_like(past_is_pad_agg),
        )

        past_observed_values_agg = (
            utils.agg_series(
                past_observed_values.slice_axis(
                    axis=1, begin=-base_history_length, end=None
                ),
                agg_multiple=agg_multiple,
            )
            .squeeze(axis=-1)
            .slice_axis(
                axis=1, begin=-self.models[agg_level].history_length, end=None
            )
        )

        past_observed_values_agg = mx.nd.where(
            # We sum observed values of base time series at `agg_multiple` time steps;
            # if all of them are 1, then the observed value for the aggregated time series is 1 and 0 otherwise.
            # We could redefine agg_series to actually compute mean, but overloading that term might cause other
            # problems later.
            past_observed_values_agg == agg_multiple,
            mx.nd.ones_like(past_observed_values_agg),
            mx.nd.zeros_like(past_observed_values_agg),
        )

        target_related_feat_agg = {
            "past_target": past_target_agg,
            "past_is_pad": past_is_pad_agg,
            "past_observed_values": past_observed_values_agg,
        }
        if future_target is not None:
            future_target_agg = utils.agg_series(
                future_target, agg_multiple=agg_multiple
            ).squeeze(axis=-1)

            future_observed_values_agg = utils.agg_series(
                future_observed_values, agg_multiple=agg_multiple
            ).squeeze(axis=-1)

            future_observed_values_agg = mx.nd.where(
                future_observed_values_agg == agg_multiple,
                mx.nd.ones_like(future_observed_values_agg),
                mx.nd.zeros_like(future_observed_values_agg),
            )

            target_related_feat_agg.update(
                {
                    "future_target": future_target_agg,
                    "future_observed_values": future_observed_values_agg,
                }
            )

        return target_related_feat_agg

    def _embeddings_to_distr(
        self,
        F,
        embeddings_at_all_levels: Tensor,
        scales: List,
    ) -> Distribution:
        distr_output = self.models[0].distr_output
        distr_args_at_all_levels: Dict = {
            arg_name: [] for arg_name in distr_output.args_dim.keys()
        }
        scales_ls = []

        start_ix = 0
        for i, num_nodes in enumerate(
            self.temporal_hierarchy.num_nodes_per_level
        ):
            end_ix = start_ix + num_nodes
            distr_args = self.models[i].proj_distr_args(
                embeddings_at_all_levels[..., start_ix:end_ix, :]
            )

            for j, arg_ls in enumerate(distr_args_at_all_levels.values()):
                arg_ls.append(distr_args[j])

            scales_ls.append(scales[i].broadcast_like(distr_args[0]))

            start_ix = end_ix

        # Last dimension contains parameters at all time-levels and aggregation can be done on it.
        distr_args_at_all_levels = {
            arg_name: F.concat(*arg_ls, dim=-1)
            for arg_name, arg_ls in distr_args_at_all_levels.items()
        }

        scale_at_all_levels = F.concat(*scales_ls, dim=-1)

        distr_at_all_levels = distr_output.distribution(
            distr_args=distr_args_at_all_levels.values(),
            scale=scale_at_all_levels,
        )

        if isinstance(distr_at_all_levels, TransformedPiecewiseLinear):
            distr_at_all_levels = TransformedPiecewiseLinear(
                base_distribution=gluonts_fixes.PiecewiseLinearWithSampling(
                    gamma=distr_at_all_levels.base_distribution.gamma,
                    slopes=distr_at_all_levels.base_distribution.slopes,
                    knot_spacings=distr_at_all_levels.base_distribution.knot_spacings,
                ),
                transforms=distr_at_all_levels.transforms,
            )

        return distr_at_all_levels

    def _distr_to_samples(
        self,
        distr_at_all_levels: Distribution,
        num_samples: int,
    ):
        if num_samples == 1:
            samples_at_all_levels = distr_at_all_levels.sample(
                num_samples=num_samples, dtype=self.dtype
            )

            # get rid of the redundant axis introduced by `sample`.
            samples_at_all_levels = samples_at_all_levels.squeeze(axis=0)
        else:
            samples_at_all_levels = distr_at_all_levels.sample_rep(
                num_samples=num_samples, dtype=self.dtype
            )

        return samples_at_all_levels


class COPDeepARTrainingNetwork(COPNetwork):
    @validated()
    def __init__(
        self,
        num_batches_per_epoch: int,
        epochs: int,
        warmstart_epoch_frac: float,
        num_samples_for_loss: int = 200,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.warmstart_epoch_frac = warmstart_epoch_frac
        self.epochs = epochs
        self.num_batches_per_epoch = num_batches_per_epoch
        self.batch_no = 0
        self.num_samples_for_loss = num_samples_for_loss

    # noinspection PyMethodOverriding,PyPep8Naming
    def hybrid_forward(
        self,
        F,
        feat_static_cat: Tensor,
        feat_static_real: Tensor,
        past_time_feat: Tensor,
        past_target: Tensor,
        past_observed_values: Tensor,
        past_is_pad: Optional[Tensor],
        future_time_feat: Tensor,
        future_target: Tensor,
        future_observed_values: Tensor,
        agg_features_dict: Dict,
    ) -> Tensor:
        """
        Computes the loss for training COPDeepAR, all inputs tensors
        representing time series have NTC layout.

        Parameters
        ----------
        F
        feat_static_cat : (batch_size, num_features)
        feat_static_real : (batch_size, num_features)
        past_time_feat : (batch_size, history_length, num_features)
        past_target : (batch_size, history_length, *target_shape)
        past_observed_values : (batch_size, history_length, *target_shape, seq_len)
        future_time_feat : (batch_size, prediction_length, num_features)
        future_target : (batch_size, prediction_length, *target_shape)
        future_observed_values : (batch_size, prediction_length, *target_shape)
        agg_features_dict: Dictionary of features for aggregated levels

        Returns loss with shape (batch_size, context + prediction_length, 1)
        -------
        """

        embeddings_at_all_levels_ls = []
        target_at_all_levels_ls = []
        scale_ls = []

        for i, agg_multiple in enumerate(
            self.temporal_hierarchy.agg_multiples
        ):
            if agg_multiple != 1:
                past_time_feat_agg = agg_features_dict[f"level_{i}"][
                    "past_time_feat_agg"
                ]
                future_time_feat_agg = agg_features_dict[f"level_{i}"][
                    "future_time_feat_agg"
                ]
            else:
                past_time_feat_agg = past_time_feat
                future_time_feat_agg = future_time_feat

            target_related_feat_agg = (
                self.get_target_related_feat_at_agg_level(
                    agg_level=i,
                    past_target=past_target,
                    past_is_pad=past_is_pad,
                    past_observed_values=past_observed_values,
                    future_target=future_target,
                    future_observed_values=future_observed_values,
                )
            )

            rnn_outputs, _, scale, _, _ = self.models[i].unroll_encoder(
                F=F,
                feat_static_cat=feat_static_cat,
                feat_static_real=feat_static_real,
                past_time_feat=past_time_feat_agg,
                future_time_feat=future_time_feat_agg,
                **target_related_feat_agg,
            )
            scale_ls.append(scale.expand_dims(axis=-1))

            # put together target sequence
            # (batch_size, seq_len, *target_shape)
            target = F.concat(
                target_related_feat_agg["past_target"].slice_axis(
                    axis=1,
                    begin=self.models[i].history_length
                    - self.models[i].context_length,
                    end=None,
                ),
                target_related_feat_agg["future_target"],
                dim=1,
            )

            # We reconcile blocks/windows of time steps: e.g., if we have 28 values of daily data, then we
            # reconcile 4 windows where each window has a length of 7 if number of leaves in the hierarchy is 7.
            window_size = self.temporal_hierarchy.num_leaves // agg_multiple
            num_windows = (
                self.models[i].context_length
                + self.models[i].prediction_length
            ) // window_size

            embeddings_at_all_levels_ls.append(
                rnn_outputs.reshape(
                    (
                        rnn_outputs.shape[0],
                        num_windows,
                        -1,
                        rnn_outputs.shape[-1],
                    )
                )
            )

            target_at_all_levels_ls.append(
                target.reshape((target.shape[0], num_windows, -1))
            )

        # Last dimension contains embeddings at all time-levels and message passing/aggregation can be done on it.
        # Shape: (bs, num_windows, total_num_time_steps_of_hierarchy, embedding_dim)
        embeddings_at_all_levels = F.concat(
            *embeddings_at_all_levels_ls, dim=-2
        )

        if self.use_gnn:
            embeddings_at_all_levels = self.gnn(embeddings_at_all_levels)

        distr_at_all_levels = self._embeddings_to_distr(
            F,
            embeddings_at_all_levels,
            scale_ls,
        )

        target_at_all_levels = F.concat(*target_at_all_levels_ls, dim=-1)

        if self.loss_function == "nll":
            loss = distr_at_all_levels.loss(x=target_at_all_levels)

            # Determine which epoch we are currently in.
            self.batch_no += 1
            epoch_no = self.batch_no // self.num_batches_per_epoch + 1
            epoch_frac = epoch_no / self.epochs

            if epoch_frac > self.warmstart_epoch_frac:
                print(
                    f"epoch_frac: {epoch_frac}. Switching the loss function to CRPS"
                )
                self.loss_function = "crps_univariate"

        else:
            samples_at_all_levels = self._distr_to_samples(
                distr_at_all_levels,
                num_samples=self.num_samples_for_loss,
            )

            if self.do_reconciliation:
                reconciled_samples_at_all_levels = reconcile_samples(
                    reconciliation_mat=self.M,
                    samples=samples_at_all_levels,
                    non_negative=self.non_negative,
                )
            else:
                reconciled_samples_at_all_levels = samples_at_all_levels

            loss = (
                EmpiricalDistribution(
                    samples=reconciled_samples_at_all_levels, event_dim=1
                )
                .loss(x=target_at_all_levels)
                .expand_dims(axis=-1)
            )

        return loss


class COPDeepARPredictionNetwork(COPNetwork):
    @validated()
    def __init__(
        self,
        return_forecasts_at_all_levels: bool = False,
        num_parallel_samples: int = 100,
        **kwargs,
    ) -> None:
        super().__init__(prediction=True, **kwargs)
        self.return_forecasts_at_all_levels = return_forecasts_at_all_levels
        self.num_parallel_samples = num_parallel_samples

    def _decode_one_window(
        self,
        F,
        model: DeepARPredictionNetwork,
        window_size: int,
        offset: int,
        static_feat: Tensor,
        past_target: Tensor,
        time_feat: Tensor,
        scale: Tensor,
        begin_states: List,
    ) -> Tuple[Tensor, Tensor]:
        """
        Computes RNN outputs by unrolling the LSTM starting with a initial
        input and state.

        Parameters
        ----------
        static_feat : Tensor
            static features. Shape: (batch_size, num_static_features).
        past_target : Tensor
            target history. Shape: (batch_size, history_length).
        time_feat : Tensor
            time features. Shape: (batch_size, prediction_length,
            num_time_features).
            Note: They still need to be for all `prediction_length` time steps.
            This function will slice the features it needs.
        scale : Tensor
            tensor containing the scale of each element in the batch.
            Shape: (batch_size, 1, 1).
        begin_states : List
            list of initial states for the LSTM layers. The shape of each
            tensor of the list should be (batch_size, num_cells)
        Returns
        --------
        Tensor
            A tensor containing sampled paths.
            Shape: (batch_size, num_sample_paths, window_size).
        """

        rnn_outputs_ls = []

        # for each future time-units we draw new samples for this time-unit and
        # update the state
        for k in range(offset, offset + window_size):
            # (batch_size * num_samples, 1, *target_shape, num_lags)
            lags = model.get_lagged_subsequences(
                F=F,
                sequence=past_target,
                sequence_length=model.history_length + k,
                indices=model.shifted_lags,
                subsequences_length=1,
            )

            # (batch_size * num_samples, 1, *target_shape, num_lags)
            lags_scaled = F.broadcast_div(lags, scale.expand_dims(axis=-1))

            # from (batch_size * num_samples, 1, *target_shape, num_lags)
            # to (batch_size * num_samples, 1, prod(target_shape) * num_lags)
            input_lags = F.reshape(
                data=lags_scaled,
                shape=(-1, 1, prod(model.target_shape) * len(model.lags_seq)),
            )

            # (batch_size * num_samples, 1, prod(target_shape) * num_lags +
            # num_time_features + num_static_features)
            decoder_input = F.concat(
                input_lags,
                time_feat.slice_axis(axis=1, begin=k, end=k + 1),
                # observed_values.expand_dims(axis=1),
                static_feat,
                dim=-1,
            )

            # output shape: (batch_size * num_samples, 1, num_cells)
            # state shape: (batch_size * num_samples, num_cells)
            rnn_outputs, begin_states = model.rnn.unroll(
                inputs=decoder_input,
                length=1,
                begin_state=begin_states,
                layout="NTC",
                merge_outputs=True,
            )
            rnn_outputs_ls.append(rnn_outputs)

            distr_args = model.proj_distr_args(rnn_outputs)

            # compute likelihood of target given the predicted parameters
            distr = model.distr_output.distribution(distr_args, scale=scale)

            # (batch_size * num_samples, 1, *target_shape)

            new_samples = distr.sample(dtype=self.dtype)

            # (batch_size * num_samples, seq_len, *target_shape)
            past_target = F.concat(past_target, new_samples, dim=1)

        # (batch_size * num_samples, prediction_length, *target_shape)
        rnn_outputs = F.concat(*rnn_outputs_ls, dim=1)

        return rnn_outputs, begin_states

    def sampling_decoder(
        self,
        F,
        state_ls,
        scale_ls,
        static_feat_ls,
        past_target_ls,
        future_time_feat_agg_ls,
    ):
        num_windows = (
            self.prediction_length // self.temporal_hierarchy.num_leaves
        )
        num_nodes_per_level = self.temporal_hierarchy.num_nodes_per_level

        reconciled_samples_at_all_levels_ls = []
        for j in range(num_windows):
            embeddings_at_all_levels_ls = []
            for i, agg_multiple in enumerate(
                self.temporal_hierarchy.agg_multiples
            ):
                rnn_outputs, states = self._decode_one_window(
                    F=F,
                    model=self.models[i],
                    window_size=num_nodes_per_level[i],
                    offset=j * num_nodes_per_level[i],
                    past_target=past_target_ls[i],
                    time_feat=future_time_feat_agg_ls[i],
                    static_feat=static_feat_ls[i],
                    scale=scale_ls[i],
                    begin_states=state_ls[i],
                )

                state_ls[i] = states
                embeddings_at_all_levels_ls.append(
                    rnn_outputs.reshape(
                        (rnn_outputs.shape[0], -1, rnn_outputs.shape[-1])
                    )
                )

            # Last dimension contains embeddings at all time-levels and message passing/aggregation can be done on it.
            # Shape: (bs, total_num_time_steps_of_hierarchy, embedding_dim)
            embeddings_at_all_levels = F.concat(
                *embeddings_at_all_levels_ls, dim=-2
            )

            if self.use_gnn:
                embeddings_at_all_levels = self.gnn(embeddings_at_all_levels)

            distr_at_all_levels = self._embeddings_to_distr(
                F,
                embeddings_at_all_levels,
                scale_ls,
            )

            samples_at_all_levels = self._distr_to_samples(
                distr_at_all_levels,
                num_samples=1,
            )

            if self.do_reconciliation:
                reconciled_samples_at_all_levels = reconcile_samples(
                    reconciliation_mat=self.M,
                    samples=samples_at_all_levels,
                    non_negative=self.non_negative,
                )
            else:
                reconciled_samples_at_all_levels = samples_at_all_levels

            rec_err = coherency_error(
                S=self.temporal_hierarchy.agg_mat,
                samples=reconciled_samples_at_all_levels.asnumpy(),
            )
            print(f"Reconciliation error: {rec_err}")

            cumsum_nodes_per_level = np.cumsum([0] + num_nodes_per_level)
            for i in range(len(self.temporal_hierarchy.agg_multiples)):
                # (batch_size * num_samples, seq_len, *target_shape)
                reconciled_samples = (
                    reconciled_samples_at_all_levels.slice_axis(
                        axis=-1,
                        begin=cumsum_nodes_per_level[i],
                        end=cumsum_nodes_per_level[i + 1],
                    )
                )
                past_target_ls[i] = F.concat(
                    past_target_ls[i], reconciled_samples, dim=1
                )

            reconciled_samples_at_all_levels_ls.append(
                reconciled_samples_at_all_levels.reshape(
                    shape=(
                        -1,
                        self.num_parallel_samples,
                        reconciled_samples_at_all_levels.shape[-1],
                    )
                ).expand_dims(axis=-2)
            )

        reconciled_samples_at_all_levels = F.concat(
            *reconciled_samples_at_all_levels_ls, dim=-2
        )
        print(reconciled_samples_at_all_levels.shape)
        return reconciled_samples_at_all_levels

    # noinspection PyMethodOverriding,PyPep8Naming
    def hybrid_forward(
        self,
        F,
        feat_static_cat: Tensor,  # (batch_size, num_features)
        feat_static_real: Tensor,  # (batch_size, num_features)
        past_time_feat: Tensor,  # (batch_size, history_length, num_features)
        past_target: Tensor,  # (batch_size, history_length, *target_shape)
        past_observed_values: Tensor,  # (batch_size, history_length, *target_shape)
        future_time_feat: Tensor,  # (batch_size, prediction_length, num_features)
        past_is_pad: Tensor,
        agg_features_dict: Dict,
    ) -> Tensor:
        """
        Predicts samples, all tensors should have NTC layout.
        Parameters
        ----------
        F
        feat_static_cat : (batch_size, num_features)
        feat_static_real : (batch_size, num_features)
        past_time_feat : (batch_size, history_length, num_features)
        past_target : (batch_size, history_length, *target_shape)
        past_observed_values : (batch_size, history_length, *target_shape)
        future_time_feat : (batch_size, prediction_length, num_features)
        agg_features_dict: Dictionary of features for aggregated levels

        Returns
        -------
        Tensor
            Predicted samples
        """

        (
            state_ls,
            scale_ls,
            static_feat_ls,
            past_target_ls,
            future_time_feat_agg_ls,
        ) = ([], [], [], [], [])

        for i, agg_multiple in enumerate(
            self.temporal_hierarchy.agg_multiples
        ):
            if agg_multiple != 1:
                past_time_feat_agg = agg_features_dict[f"level_{i}"][
                    "past_time_feat_agg"
                ]
                future_time_feat_agg = agg_features_dict[f"level_{i}"][
                    "future_time_feat_agg"
                ]
            else:
                past_time_feat_agg = past_time_feat
                future_time_feat_agg = future_time_feat

            target_related_feat_agg = (
                self.get_target_related_feat_at_agg_level(
                    agg_level=i,
                    past_target=past_target,
                    past_is_pad=past_is_pad,
                    past_observed_values=past_observed_values,
                )
            )

            # unroll the decoder in "prediction mode", i.e. with past data only
            _, states, scale, static_feat, imputed_sequence = self.models[
                i
            ].unroll_encoder(
                F=F,
                feat_static_cat=feat_static_cat,
                feat_static_real=feat_static_real,
                past_time_feat=past_time_feat_agg,
                future_observed_values=None,
                future_time_feat=None,
                future_target=None,
                **target_related_feat_agg,
            )

            # blows-up the dimension of each tensor to batch_size *
            # self.num_parallel_samples for increasing parallelism
            repeated_past_target = imputed_sequence.repeat(
                repeats=self.num_parallel_samples, axis=0
            )
            repeated_states = [
                s.repeat(repeats=self.num_parallel_samples, axis=0)
                for s in states
            ]
            repeated_time_feat = future_time_feat_agg.repeat(
                repeats=self.num_parallel_samples, axis=0
            )
            repeated_static_feat = static_feat.repeat(
                repeats=self.num_parallel_samples, axis=0
            ).expand_dims(axis=1)
            repeated_scale = scale.repeat(
                repeats=self.num_parallel_samples, axis=0
            )

            state_ls.append(repeated_states)
            scale_ls.append(repeated_scale)
            static_feat_ls.append(repeated_static_feat)
            past_target_ls.append(repeated_past_target)
            future_time_feat_agg_ls.append(repeated_time_feat)

        reconciled_samples_at_all_levels = self.sampling_decoder(
            F,
            state_ls=state_ls,
            scale_ls=scale_ls,
            static_feat_ls=static_feat_ls,
            past_target_ls=past_target_ls,
            future_time_feat_agg_ls=future_time_feat_agg_ls,
        )

        if self.return_forecasts_at_all_levels:
            return reconciled_samples_at_all_levels
        else:
            reconciled_samples_at_bottom_level = (
                reconciled_samples_at_all_levels.slice_axis(
                    axis=-1,
                    begin=-self.temporal_hierarchy.num_leaves,
                    end=None,
                )
            )

            reconciled_samples_at_bottom_level = (
                reconciled_samples_at_bottom_level.reshape(
                    (
                        reconciled_samples_at_bottom_level.shape[0],
                        reconciled_samples_at_bottom_level.shape[1],
                        -1,
                    )
                )
            )

            return reconciled_samples_at_bottom_level
