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

import mxnet as mx
from gluonts_customized.core.component import validated
from gluonts_customized.mx import Tensor
from gluonts_customized.mx.distribution import Distribution, DistributionOutput
from gluonts_customized.mx.distribution.distribution import getF
from mxnet import gluon


class DeepRenewalNetwork(gluon.HybridBlock):
    def __init__(
        self,
        prediction_length: int,
        context_length: int,
        interval_distr_output: DistributionOutput,
        size_distr_output: DistributionOutput,
        num_cells: int,
        num_layers: int,
        dropout_rate: float = 0.1,
        **kwargs,
    ):
        super().__init__(**kwargs)

        self.prediction_length = prediction_length
        self.context_length = context_length
        self.interval_distr_output = interval_distr_output
        self.size_distr_output = size_distr_output
        self.num_cells = num_cells
        self.num_layers = num_layers

        with self.name_scope():
            self.rnn_cell = mx.gluon.rnn.LSTM(
                hidden_size=num_cells,
                num_layers=num_layers,
                layout="NTC",
                dropout=dropout_rate,
                h2h_weight_initializer=mx.init.Xavier(),
                h2r_weight_initializer=mx.init.Xavier(),
            )
            self.dense = mx.gluon.nn.Dense(
                units=2,
                activation="softrelu",
                flatten=False,
                weight_initializer=mx.init.Xavier(),
                bias_initializer=mx.init.Constant(5.0),
            )

        # add constant learnable parameters if the distributions require
        # more than the conditional mean
        with self.name_scope():
            if len(self.interval_distr_output.args_dim) > 1:
                self.interval_alpha_bias = self.params.get(
                    "interval_alpha_bias",
                    shape=(1,),
                    init=mx.init.Constant(2),
                )
            if len(self.size_distr_output.args_dim) > 1:
                self.size_alpha_bias = self.params.get(
                    "size_alpha_bias",
                    shape=(1,),
                    init=mx.init.Constant(2),
                )

    def begin_state(self, batch_size: int) -> Tuple[Tensor, Tensor]:
        return self.rnn_cell.begin_state(batch_size=batch_size)

    def distribution(
        self,
        cond_mean: Tensor,
        interval_alpha_bias: Optional[Tensor] = None,
        size_alpha_bias: Optional[Tensor] = None,
    ) -> Tuple[Distribution, ...]:
        F = getF(cond_mean)

        cond_interval, cond_size = F.split(cond_mean, num_outputs=2, axis=-1)

        alpha_biases = [
            (
                F.broadcast_mul(F.ones_like(cond_interval), bias)
                if bias is not None
                else None
            )
            for bias in [interval_alpha_bias, size_alpha_bias]
        ]

        distr_params = zip(
            [self.interval_distr_output, self.size_distr_output],
            [cond_interval, cond_size],
            alpha_biases,
        )

        return tuple(
            (
                do.distribution(mean)
                if len(do.args_dim) == 1
                else do.distribution(
                    [mean, F.Activation(alpha_bias, "softrelu") + 1e-5]
                )
            )
            for ix, (do, mean, alpha_bias) in enumerate(distr_params)
        )

    @staticmethod
    def forwardshift(A):
        """
        Shift an array's content forward by 1 time step along the first axis,
        keeping the shape identical by padding on the left with zeros.

        Parameters
        ----------
        A : nd.NDArray
            Shape (N, T, ...), the tensor in which the entries will be shifted
            forward by one
        """
        F = getF(A)
        A = F.Concat(
            F.zeros_like(F.slice_axis(A, axis=1, begin=0, end=1)), A, dim=1
        )
        return F.slice_axis(A, axis=1, begin=0, end=-1)

    def mu_map(
        self,
        data: Tensor,
        shift: bool = True,
        state=None,
    ) -> Tensor:
        """
        Map a given (N, T, 2) tensor to conditional interval and size means of
        the next time step.

        Parameters
        ----------
        data: Tensor
            tensor of shape (N, T, 2) containing the past target, with the
            first array along the last dimension containing intervals, and
            the second array containing demand sizes.
        shift: bool
            if True, the past data will be shifted forward by one time step
        state
            RNN cell state if available

        Returns
        -------
        conditional_means: Tensor
            tensor of shape (N, T, 2) containing conditional means for
            intervals and sizes respectively
        rnn_out_state
            output state of the RNN, returned only if the input `state` is not
            None
        """
        data = self.forwardshift(data) if shift else data

        if state is not None:
            rnn_out, rnn_out_state = self.rnn_cell(data, state)
        else:
            rnn_out = self.rnn_cell(data)

        dense_out = self.dense(rnn_out)
        model_mu = dense_out + 1e-5

        return model_mu if state is None else (model_mu, rnn_out_state)


class DeepRenewalTrainingNetwork(DeepRenewalNetwork):
    def hybrid_forward(
        self,
        F,
        past_target,
        valid_length,
        interval_alpha_bias=None,
        size_alpha_bias=None,
        **kwargs,
    ) -> Tensor:
        """
        Compute negative log likelihood losses for given data.

        Parameters
        ----------
        F
        past_target
            shape (batch_size, history_length, 2). Expects a right-aligned
            ragged tensor containing the past data in an interval-size
            format.
        valid_length
            number of valid data points in each array of the batch. i.e.,
            non-valid points will be masked out. shape (batch_size, 1)
        interval_alpha_bias
        size_alpha_bias

        Returns
        -------
        loss
            negative log likelihood with shape (batch_size, history_length)
        """
        cond_mean = self.mu_map(past_target, shift=True)
        dist_interval, dist_size = self.distribution(
            cond_mean, interval_alpha_bias, size_alpha_bias
        )
        data_interval, data_size = F.split(past_target, num_outputs=2, axis=-1)

        # TODO: on windows, operators below may produce NaN values

        log_prob = F.squeeze(
            dist_interval.log_prob(data_interval - 1)
            + dist_size.log_prob(data_size - 1),
            axis=-1,
        )
        reverse_lengths = F.broadcast_sub(
            F.maximum(1, F.max(valid_length)), valid_length
        ).squeeze(-1)
        mask = F.SequenceMask(
            F.ones_like(log_prob),
            reverse_lengths,
            use_sequence_length=True,
            axis=1,
        )
        loss = -F.where(1 - mask, log_prob, F.zeros_like(log_prob))

        return loss


class DeepRenewalPredictionNetwork(DeepRenewalNetwork):
    @validated()
    def __init__(self, num_parallel_samples: int = 100, **kwargs) -> None:
        super().__init__(**kwargs)
        self.num_parallel_samples = num_parallel_samples

    def sampling_decoder(
        self,
        F,
        past_target: Tensor,
        interval_alpha_bias: Optional[Tensor] = None,
        size_alpha_bias: Optional[Tensor] = None,
        time_remaining: Optional[Tensor] = None,
    ) -> Tensor:
        """
        Sample a trajectory of interval-size pairs from the model given past
        target.

        Parameters
        ----------
        past_target
            (batch_size, history_length, 2) shaped tensor containing the past
            time series in an interval-size representation.
        time_remaining
            (batch_size, 1) shaped tensor containing the number of time steps
            that were zero before the start of the forecast horizon

        Returns
        -------
        samples
            samples of shape (batch_size, num_parallel_samples, 2,
            sequence_length)
        """
        if time_remaining is None:
            time_remaining = F.broadcast_like(
                F.zeros((1,)), past_target, lhs_axes=(0,), rhs_axes=(0,)
            )

        # if no valid points in the past_target, override time remaining with
        # zero
        batch_sum = F.sum(past_target, axis=1).sum(-1).expand_dims(-1)
        time_remaining = F.where(
            batch_sum > 0, time_remaining, F.zeros_like(time_remaining)
        )

        repeated_past_target = past_target.repeat(
            repeats=self.num_parallel_samples, axis=0
        )  # (N * samples, T, 2)
        repeated_time_remaining = time_remaining.repeat(
            repeats=self.num_parallel_samples, axis=0
        )  # (N * samples, 1)

        interval_samples = []
        size_samples = []

        for t in range(self.prediction_length):
            cond_mean = self.mu_map(repeated_past_target, shift=False)
            cond_mean_last = F.slice_axis(
                cond_mean, axis=1, begin=-1, end=None
            ).squeeze(1)

            dist_interval, dist_size = self.distribution(
                cond_mean_last, interval_alpha_bias, size_alpha_bias
            )

            # initial samples for interval should be taken conditionally
            # we achieve this via (leaky) rejection sampling
            if t == 0:
                interval_sample = F.zeros_like(repeated_time_remaining)
                for j in range(50):
                    interval_sample = F.where(
                        interval_sample > repeated_time_remaining,
                        interval_sample,
                        dist_interval.sample(),
                    )
                interval_sample = (
                    F.where(
                        interval_sample <= repeated_time_remaining,
                        repeated_time_remaining,
                        interval_sample,
                    )
                    + 1
                )
            else:
                interval_sample = dist_interval.sample() + 1
            size_sample = dist_size.sample() + 1

            interval_samples.append(interval_sample)
            size_samples.append(size_sample)

            repeated_past_target = F.concat(
                repeated_past_target,
                F.concat(interval_sample, size_sample, dim=-1).expand_dims(1),
                dim=1,
            )

        interval_samples[0] = interval_samples[0] - repeated_time_remaining
        samples = F.concat(
            *[
                F.concat(x, y, dim=-1).expand_dims(1)
                for x, y in zip(interval_samples, size_samples)
            ],
            dim=1,
        )

        return samples.reshape(
            shape=(-1, self.num_parallel_samples) + samples.shape[1:]
        ).swapaxes(2, 3)

    def hybrid_forward(
        self,
        F,
        past_target,
        time_remaining,
        interval_alpha_bias=None,
        size_alpha_bias=None,
        **kwargs,
    ) -> Tensor:
        return self.sampling_decoder(
            F,
            past_target=past_target,
            interval_alpha_bias=interval_alpha_bias,
            size_alpha_bias=size_alpha_bias,
            time_remaining=time_remaining,
        )
