#
# Copyright (c) 2017 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License 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.
#

"""
Module that defines the LSTM middleware class
"""

from typing import Union
from types import ModuleType

import mxnet as mx
from mxnet.gluon import rnn
from rl_coach.architectures.mxnet_components.layers import Dense
from rl_coach.architectures.mxnet_components.middlewares.middleware import Middleware
from rl_coach.architectures.middleware_parameters import LSTMMiddlewareParameters
from rl_coach.base_parameters import MiddlewareScheme

nd_sym_type = Union[mx.nd.NDArray, mx.sym.Symbol]


class LSTMMiddleware(Middleware):
    def __init__(self, params: LSTMMiddlewareParameters):
        """
        LSTMMiddleware or Long Short Term Memory Middleware can be used in the middle part of the network. It takes the
        embeddings from the input embedders, after they were aggregated in some method (for example, concatenation)
        and passes it through a neural network  which can be customizable but shared between the heads of the network.

        :param params: parameters object containing batchnorm, activation_function, dropout and
            number_of_lstm_cells properties.
        """
        super(LSTMMiddleware, self).__init__(params)
        self.number_of_lstm_cells = params.number_of_lstm_cells
        with self.name_scope():
            self.lstm = rnn.LSTM(hidden_size=self.number_of_lstm_cells)

    @property
    def schemes(self) -> dict:
        """
        Schemes are the pre-defined network architectures of various depths and complexities that can be used for the
        Middleware. Are used to create Block when LSTMMiddleware is initialised, and are applied before the LSTM.

        :return: dictionary of schemes, with key of type MiddlewareScheme enum and value being list of mxnet.gluon.Block.
        """
        return {
            MiddlewareScheme.Empty:
                [],

            # Use for PPO
            MiddlewareScheme.Shallow:
                [
                    Dense(units=64)
                ],

            # Use for DQN
            MiddlewareScheme.Medium:
                [
                    Dense(units=512)
                ],

            MiddlewareScheme.Deep:
                [
                    Dense(units=128),
                    Dense(units=128),
                    Dense(units=128)
                ]
        }

    def hybrid_forward(self,
                       F: ModuleType,
                       x: nd_sym_type,
                       *args, **kwargs) -> nd_sym_type:
        """
        Used for forward pass through LSTM middleware network.
        Applies dense layers from selected scheme before passing result to LSTM layer.

        :param F: backend api, either `mxnet.nd` or `mxnet.sym` (if block has been hybridized).
        :param x: state embedding, of shape (batch_size, in_channels).
        :return: state middleware embedding, where shape is (batch_size, channels).
        """
        x_ntc = x.reshape(shape=(0, 0, -1))
        emb_ntc = super(LSTMMiddleware, self).hybrid_forward(F, x_ntc, *args, **kwargs)
        emb_tnc = emb_ntc.transpose(axes=(1, 0, 2))
        return self.lstm(emb_tnc)