# 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 mxnet as mx

from gluonts_customized.core.component import validated
from gluonts_customized.mx import Tensor
from gluonts_customized.mx.block.feature import FeatureEmbedder
from gluonts_customized.mx.block.scaler import MeanScaler, NOPScaler
from gluonts_customized.mx.distribution.lds import LDS, LDSArgsProj, ParameterBounds
from gluonts_customized.mx.util import make_nd_diag, weighted_average

from .issm import ISSM


class DeepStateNetwork(mx.gluon.HybridBlock):
    @validated()
    def __init__(
        self,
        num_layers: int,
        num_cells: int,
        cell_type: str,
        past_length: int,
        prediction_length: int,
        issm: ISSM,
        dropout_rate: float,
        cardinality: List[int],
        embedding_dimension: List[int],
        scaling: bool = True,
        noise_std_bounds: ParameterBounds = ParameterBounds(1e-6, 1.0),
        prior_cov_bounds: ParameterBounds = ParameterBounds(1e-6, 1.0),
        innovation_bounds: ParameterBounds = ParameterBounds(1e-6, 0.01),
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.num_layers = num_layers
        self.num_cells = num_cells
        self.cell_type = cell_type
        self.past_length = past_length
        self.prediction_length = prediction_length
        self.issm = issm
        self.dropout_rate = dropout_rate
        self.cardinality = cardinality
        self.embedding_dimension = embedding_dimension
        self.num_cat = len(cardinality)
        self.scaling = scaling

        assert len(cardinality) == len(embedding_dimension), (
            "embedding_dimension should be a list with the same size as"
            " cardinality"
        )
        self.univariate = self.issm.output_dim() == 1

        self.noise_std_bounds = noise_std_bounds
        self.prior_cov_bounds = prior_cov_bounds
        self.innovation_bounds = innovation_bounds

        with self.name_scope():
            self.prior_mean_model = mx.gluon.nn.Dense(
                units=self.issm.latent_dim(), flatten=False
            )
            self.prior_cov_diag_model = mx.gluon.nn.Dense(
                units=self.issm.latent_dim(),
                activation="sigmoid",
                flatten=False,
            )
            self.lstm = mx.gluon.rnn.HybridSequentialRNNCell()
            self.lds_proj = LDSArgsProj(
                output_dim=self.issm.output_dim(),
                noise_std_bounds=self.noise_std_bounds,
                innovation_bounds=self.innovation_bounds,
            )
            for k in range(num_layers):
                cell = mx.gluon.rnn.LSTMCell(hidden_size=num_cells)
                cell = mx.gluon.rnn.ResidualCell(cell) if k > 0 else cell
                cell = (
                    mx.gluon.rnn.ZoneoutCell(cell, zoneout_states=dropout_rate)
                    if dropout_rate > 0.0
                    else cell
                )
                self.lstm.add(cell)
            self.embedder = FeatureEmbedder(
                cardinalities=cardinality, embedding_dims=embedding_dimension
            )
            if scaling:
                self.scaler = MeanScaler(keepdims=False)
            else:
                self.scaler = NOPScaler(keepdims=False)

    def compute_lds(
        self,
        F,
        feat_static_cat: Tensor,
        seasonal_indicators: Tensor,
        time_feat: Tensor,
        length: int,
        prior_mean: Optional[Tensor] = None,
        prior_cov: Optional[Tensor] = None,
        lstm_begin_state: Optional[List[Tensor]] = None,
    ):
        # embed categorical features and expand along time axis
        embedded_cat = self.embedder(feat_static_cat)
        repeated_static_features = embedded_cat.expand_dims(axis=1).repeat(
            axis=1, repeats=length
        )

        # construct big features tensor (context)
        features = F.concat(time_feat, repeated_static_features, dim=2)

        output, lstm_final_state = self.lstm.unroll(
            inputs=features,
            begin_state=lstm_begin_state,
            length=length,
            merge_outputs=True,
        )

        if prior_mean is None:
            prior_input = F.slice_axis(output, axis=1, begin=0, end=1).squeeze(
                axis=1
            )

            prior_mean = self.prior_mean_model(prior_input)
            prior_cov_diag = (
                self.prior_cov_diag_model(prior_input)
                * (self.prior_cov_bounds.upper - self.prior_cov_bounds.lower)
                + self.prior_cov_bounds.lower
            )
            prior_cov = make_nd_diag(F, prior_cov_diag, self.issm.latent_dim())

        (
            emission_coeff,
            transition_coeff,
            innovation_coeff,
        ) = self.issm.get_issm_coeff(seasonal_indicators)

        noise_std, innovation, residuals = self.lds_proj(output)

        lds = LDS(
            emission_coeff=emission_coeff,
            transition_coeff=transition_coeff,
            innovation_coeff=F.broadcast_mul(innovation, innovation_coeff),
            noise_std=noise_std,
            residuals=residuals,
            prior_mean=prior_mean,
            prior_cov=prior_cov,
            latent_dim=self.issm.latent_dim(),
            output_dim=self.issm.output_dim(),
            seq_length=length,
        )

        return lds, lstm_final_state


class DeepStateTrainingNetwork(DeepStateNetwork):
    def hybrid_forward(
        self,
        F,
        feat_static_cat: Tensor,
        past_observed_values: Tensor,
        past_seasonal_indicators: Tensor,
        past_time_feat: Tensor,
        past_target: Tensor,
    ) -> Tensor:
        lds, _ = self.compute_lds(
            F,
            feat_static_cat=feat_static_cat,
            seasonal_indicators=past_seasonal_indicators.slice_axis(
                axis=1, begin=-self.past_length, end=None
            ),
            time_feat=past_time_feat.slice_axis(
                axis=1, begin=-self.past_length, end=None
            ),
            length=self.past_length,
        )

        _, scale = self.scaler(past_target, past_observed_values)

        observed_context = past_observed_values.slice_axis(
            axis=1, begin=-self.past_length, end=None
        )

        ll, _, _ = lds.log_prob(
            x=past_target.slice_axis(
                axis=1, begin=-self.past_length, end=None
            ),
            observed=observed_context.min(axis=-1, keepdims=False),
            scale=scale,
        )

        return weighted_average(
            F=F, x=-ll, axis=1, weights=observed_context.squeeze(axis=-1)
        )


class DeepStatePredictionNetwork(DeepStateNetwork):
    @validated()
    def __init__(self, num_parallel_samples: int, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.num_parallel_samples = num_parallel_samples

    def hybrid_forward(
        self,
        F,
        feat_static_cat: Tensor,
        past_observed_values: Tensor,
        past_seasonal_indicators: Tensor,
        past_time_feat: Tensor,
        past_target: Tensor,
        future_seasonal_indicators: Tensor,
        future_time_feat: Tensor,
    ) -> Tensor:
        lds, lstm_state = self.compute_lds(
            F,
            feat_static_cat=feat_static_cat,
            seasonal_indicators=past_seasonal_indicators.slice_axis(
                axis=1, begin=-self.past_length, end=None
            ),
            time_feat=past_time_feat.slice_axis(
                axis=1, begin=-self.past_length, end=None
            ),
            length=self.past_length,
        )

        _, scale = self.scaler(past_target, past_observed_values)

        observed_context = past_observed_values.slice_axis(
            axis=1, begin=-self.past_length, end=None
        )

        _, final_mean, final_cov = lds.log_prob(
            x=past_target.slice_axis(
                axis=1, begin=-self.past_length, end=None
            ),
            observed=observed_context.min(axis=-1, keepdims=False),
            scale=scale,
        )

        lds_prediction, _ = self.compute_lds(
            F,
            feat_static_cat=feat_static_cat,
            seasonal_indicators=future_seasonal_indicators,
            time_feat=future_time_feat,
            length=self.prediction_length,
            lstm_begin_state=lstm_state,
            prior_mean=final_mean,
            prior_cov=final_cov,
        )

        samples = lds_prediction.sample(
            num_samples=self.num_parallel_samples, scale=scale
        )

        # convert samples from
        # (num_samples, batch_size, prediction_length, target_dim)
        # to
        # (batch_size, num_samples, prediction_length, target_dim)
        # and squeeze last axis in the univariate case
        if self.univariate:
            return samples.transpose(axes=(1, 0, 2, 3)).squeeze(axis=3)
        else:
            return samples.transpose(axes=(1, 0, 2, 3))
