# 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 List, Optional

import torch

from gluonts.core.component import validated
from gluonts.torch.model.deepar.module import DeepARModel
from gluonts.torch.distributions import (
    DistributionOutput,
    MQF2DistributionOutput,
)

from cpflows.flows import ActNorm
from cpflows.icnn import PICNN
from .icnn_utils import DeepConvexNet, SequentialNet


class MQF2MultiHorizonModel(DeepARModel):
    @validated()
    def __init__(
        self,
        freq: str,
        context_length: int,
        prediction_length: int,
        num_feat_dynamic_real: int,
        num_feat_static_real: int,
        num_feat_static_cat: int,
        cardinality: List[int],
        distr_output: Optional[DistributionOutput] = None,
        embedding_dimension: Optional[List[int]] = None,
        num_layers: int = 2,
        hidden_size: int = 40,
        dropout_rate: float = 0.1,
        lags_seq: Optional[List[int]] = None,
        scaling: bool = True,
        num_parallel_samples: int = 100,
        icnn_hidden_size: int = 20,
        icnn_num_layers: int = 2,
        is_energy_score: bool = True,
        threshold_input: float = 100,
        es_num_samples: int = 50,
        estimate_logdet: bool = False,
    ) -> None:
        r"""
        Model class for the model MQF2 proposed in the paper
        ``Multivariate Quantile Function Forecaster``
        by Kan, Aubet, Januschowski, Park, Benidis, Ruthotto, Gasthaus

        This is the multi-horizon (multivariate in time step) variant of MQF2

        This class is based on gluonts.torch.model.deepar.module.DeepARModel

        Refer to MQF2MultiHorizonEstimator for the description of parameters
        """

        super().__init__(
            freq=freq,
            context_length=context_length,
            prediction_length=prediction_length,
            num_feat_dynamic_real=num_feat_dynamic_real,
            num_feat_static_real=num_feat_static_real,
            num_feat_static_cat=num_feat_static_cat,
            cardinality=cardinality,
            embedding_dimension=embedding_dimension,
            num_layers=num_layers,
            hidden_size=hidden_size,
            dropout_rate=dropout_rate,
            distr_output=(
                distr_output
                if distr_output is not None
                else MQF2DistributionOutput(prediction_length)
            ),
            lags_seq=lags_seq,
            scaling=scaling,
            num_parallel_samples=num_parallel_samples,
        )

        self.threshold_input = threshold_input
        self.es_num_samples = es_num_samples

        convexnet = PICNN(
            dim=prediction_length,
            dimh=icnn_hidden_size,
            dimc=hidden_size,
            num_hidden_layers=icnn_num_layers,
            symm_act_first=True,
        )
        deepconvexnet = DeepConvexNet(
            convexnet,
            prediction_length,
            is_energy_score=is_energy_score,
            estimate_logdet=estimate_logdet,
        )

        if is_energy_score:
            networks = [deepconvexnet]
        else:
            networks = [
                ActNorm(prediction_length),
                deepconvexnet,
                ActNorm(prediction_length),
            ]

        self.picnn = SequentialNet(networks)

    @torch.jit.ignore
    def output_distribution(
        self,
        picnn: SequentialNet,
        hidden_state: torch.Tensor,
        scale: Optional[torch.Tensor] = None,
        inference: bool = False,
    ) -> torch.distributions.Distribution:
        """
        Returns the MQF2Distribution instance.

        Parameters
        ----------
        picnn
            A SequentialNet instance of a
            partially input convex neural network (picnn)
        hidden_state
            RNN hidden state (batch_size, context_length, hidden_size)
        scale
            scaling of the data (batch_size, 1)
        inference
            If True, pass only the last hidden state
            to the forecaster for prediction
            Otherwise, pass all the hidden states to train the forecaster

        Returns
        -------
        MQF2Distribution instance
            MQF2 parametrized by hidden_state
        """

        if inference:
            hidden_state = hidden_state[:, -1]

        return self.distr_output.distribution(picnn, hidden_state, scale=scale)

    def forward(
        self,
        feat_static_cat: torch.Tensor,
        feat_static_real: torch.Tensor,
        past_time_feat: torch.Tensor,
        past_target: torch.Tensor,
        past_observed_values: torch.Tensor,
        future_time_feat: torch.Tensor,
        num_parallel_samples: Optional[int] = None,
    ) -> torch.Tensor:
        """
        Generates the predicted sample paths.

        Parameters
        ----------
        feat_static_cat
            Static categorical features (batch_size, num_feat_static_cat)
        feat_static_real
            Static real-valued features (batch_size, num_feat_static_real)
        past_time_feat
            Past time features (batch_size, history_length, num_features)
        past_target
            Past target values (batch_size, history_length)
        past_observed_values
            Indicator whether or not the values were observed
            (batch_size, history_length)
        future_time_feat
            Future time features (batch_size, prediction_length, num_features)
        num_parallel_samples
            Number of parallel sample paths generated for each time series

        Returns
        -------
        sample_paths
            Sample paths (batch_size, num_parallel_samples, prediction_length)
        """

        if num_parallel_samples is None:
            num_parallel_samples = self.num_parallel_samples

        _, scale, hidden_state, _, _ = self.unroll_lagged_rnn(
            feat_static_cat,
            feat_static_real,
            past_time_feat,
            past_target,
            past_observed_values,
            future_time_feat[:, :1],
        )

        hidden_state = hidden_state[:, : self.context_length]

        distr = self.output_distribution(
            self.picnn, hidden_state, inference=True
        )

        unscaled_future_samples = distr.sample(
            sample_shape=(num_parallel_samples,)
        )

        return unscaled_future_samples * scale.unsqueeze(-1)
