# 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.

import math
from typing import List, Optional, Tuple

from mxnet import init
from mxnet.gluon import HybridBlock, nn, rnn

from gluonts_customized.core.component import validated
from gluonts_customized.mx import Tensor


class GatedLinearUnit(HybridBlock):
    @validated()
    def __init__(self, axis: int = -1, nonlinear: bool = True, **kwargs):
        super().__init__(**kwargs)
        self.axis = axis
        self.nonlinear = nonlinear

    def hybrid_forward(self, F, x: Tensor) -> Tensor:
        val, gate = F.split(x, axis=self.axis, num_outputs=2)
        if self.nonlinear:
            val = F.tanh(val)
        gate = F.sigmoid(gate)
        return F.broadcast_mul(gate, val)


class GatedResidualNetwork(HybridBlock):
    @validated()
    def __init__(
        self,
        d_hidden: int,
        d_input: Optional[int] = None,
        d_output: Optional[int] = None,
        d_static: Optional[int] = None,
        dropout: float = 0.0,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.d_hidden = d_hidden
        self.d_input = d_input or d_hidden
        self.d_static = d_static or 0
        if d_output is None:
            self.d_output = self.d_input
            self.add_skip = False
        else:
            self.d_output = d_output
            if d_output != self.d_input:
                self.add_skip = True
                with self.name_scope():
                    self.skip_proj = nn.Dense(
                        units=self.d_output,
                        in_units=self.d_input,
                        flatten=False,
                        weight_initializer=init.Xavier(),
                    )
            else:
                self.add_skip = False

        with self.name_scope():
            self.mlp = nn.HybridSequential(prefix="mlp_")
            self.mlp.add(
                nn.Dense(
                    units=self.d_hidden,
                    in_units=self.d_input + self.d_static,
                    flatten=False,
                    weight_initializer=init.Xavier(),
                )
            )
            self.mlp.add(nn.ELU())
            self.mlp.add(
                nn.Dense(
                    units=self.d_hidden,
                    in_units=self.d_hidden,
                    flatten=False,
                    weight_initializer=init.Xavier(),
                )
            )
            self.mlp.add(nn.Dropout(dropout)),
            self.mlp.add(
                nn.Dense(
                    units=self.d_output * 2,
                    in_units=self.d_hidden,
                    flatten=False,
                    weight_initializer=init.Xavier(),
                )
            )
            self.mlp.add(
                GatedLinearUnit(
                    axis=-1,
                    nonlinear=False,
                )
            )
            self.lnorm = nn.LayerNorm(axis=-1, in_channels=self.d_output)

    def hybrid_forward(
        self,
        F,
        x: Tensor,
        c: Optional[Tensor] = None,
    ) -> Tensor:
        if self.add_skip:
            skip = self.skip_proj(x)
        else:
            skip = x
        if self.d_static > 0 and c is None:
            raise ValueError("static variable is expected.")
        if self.d_static == 0 and c is not None:
            raise ValueError("static variable is not accpeted.")
        if c is not None:
            x = F.concat(x, c, dim=-1)
        x = self.mlp(x)
        x = self.lnorm(F.broadcast_add(x, skip))
        return x


class VariableSelectionNetwork(HybridBlock):
    @validated()
    def __init__(
        self,
        d_hidden: int,
        n_vars: int,
        dropout: float = 0.0,
        add_static: bool = False,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.d_hidden = d_hidden
        self.n_vars = n_vars
        self.add_static = add_static

        with self.name_scope():
            self.weight_network = GatedResidualNetwork(
                d_hidden=self.d_hidden,
                d_input=self.d_hidden * self.n_vars,
                d_output=self.n_vars,
                d_static=self.d_hidden if add_static else None,
                dropout=dropout,
            )
            self.variable_network = []
            for n in range(self.n_vars):
                var_net = GatedResidualNetwork(
                    d_hidden=self.d_hidden,
                    dropout=dropout,
                )
                self.register_child(var_net, name=f"var_{n+1}")
                self.variable_network.append(var_net)

    def hybrid_forward(
        self,
        F,
        variables: List[Tensor],
        static: Optional[Tensor] = None,
    ) -> Tuple[Tensor, Tensor]:
        if len(variables) != self.n_vars:
            raise ValueError(
                f"expect {self.n_vars} variables, {len(variables)} given."
            )
        if self.add_static and static is None:
            raise ValueError("static variable is expected.")
        if not self.add_static and static is not None:
            raise ValueError("static variable is not accpeted.")
        flatten = F.concat(*variables, dim=-1)
        if static is not None:
            static = F.broadcast_like(static, variables[0])
        weight = self.weight_network(flatten, static)
        weight = F.expand_dims(weight, axis=-2)
        weight = F.softmax(weight, axis=-1)
        var_encodings = []
        for var, net in zip(variables, self.variable_network):
            var_encodings.append(net(var))
        var_encodings = F.stack(*var_encodings, axis=-1)
        var_encodings = F.sum(F.broadcast_mul(var_encodings, weight), axis=-1)
        return var_encodings, weight


class SelfAttention(HybridBlock):
    @validated()
    def __init__(
        self,
        context_length: int,
        prediction_length: int,
        d_hidden: int,
        n_head: int = 1,
        bias: bool = True,
        share_values: bool = False,
        dropout: float = 0.0,
        temperature: float = 1.0,
        **kwargs,
    ):
        super().__init__(**kwargs)
        if d_hidden % n_head != 0:
            raise ValueError(
                f"hidden dim {d_hidden} cannot be split into {n_head} heads."
            )
        self.context_length = context_length
        self.prediction_length = prediction_length
        self.d_hidden = d_hidden
        self.n_head = n_head
        self.d_head = d_hidden // n_head
        self.bias = bias
        self.share_values = share_values
        self.temperature = temperature

        with self.name_scope():
            self.dropout = nn.Dropout(dropout)
            self.q_proj = nn.Dense(
                units=self.d_hidden,
                in_units=self.d_hidden,
                use_bias=self.bias,
                flatten=False,
                weight_initializer=init.Xavier(),
                prefix="q_proj_",
            )
            self.k_proj = nn.Dense(
                units=self.d_hidden,
                in_units=self.d_hidden,
                use_bias=self.bias,
                flatten=False,
                weight_initializer=init.Xavier(),
                prefix="k_proj_",
            )
            self.v_proj = nn.Dense(
                units=self.d_head if self.share_values else self.d_hidden,
                in_units=self.d_hidden,
                use_bias=self.bias,
                flatten=False,
                weight_initializer=init.Xavier(),
                prefix="v_proj_",
            )
            self.out_proj = nn.Dense(
                units=self.d_hidden,
                in_units=self.d_hidden,
                use_bias=self.bias,
                flatten=False,
                weight_initializer=init.Xavier(),
                prefix="out_proj_",
            )

    def _split_head(self, F, x: Tensor) -> Tensor:
        x = F.reshape(data=x, shape=(0, 0, -4, self.n_head, self.d_head))
        x = F.swapaxes(data=x, dim1=1, dim2=2)
        return x

    def _merge_head(self, F, x: Tensor) -> Tensor:
        x = F.swapaxes(data=x, dim1=1, dim2=2)
        x = F.reshape(data=x, shape=(0, 0, self.d_hidden))
        return x

    def _compute_qkv(self, F, x: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
        cx = F.slice_axis(x, axis=1, begin=-self.prediction_length, end=None)
        q = self.q_proj(cx)
        k = self.k_proj(x)
        q = self._split_head(F, q)
        k = self._split_head(F, k)
        v = self.v_proj(x)
        if self.share_values:
            v = F.broadcast_like(v.expand_dims(axis=1), k)
        else:
            v = self._split_head(F, v)
        return q, k, v

    def _apply_mask(
        self, F, score: Tensor, key_mask: Optional[Tensor]
    ) -> Tensor:
        k_idx = F.contrib.arange_like(score, axis=-1)
        k_idx = (
            k_idx.expand_dims(axis=0).expand_dims(axis=0).expand_dims(axis=0)
        )
        q_idx = F.contrib.arange_like(score, axis=-2) + self.context_length
        q_idx = (
            q_idx.expand_dims(axis=-1).expand_dims(axis=0).expand_dims(axis=0)
        )
        unidir_mask = F.broadcast_lesser_equal(k_idx, q_idx)
        unidir_mask = F.broadcast_like(unidir_mask, score)
        score = F.where(unidir_mask, score, F.ones_like(score) * -1e9)
        if key_mask is not None:
            key_mask = key_mask.expand_dims(axis=1)  # head
            key_mask = key_mask.expand_dims(axis=2)  # query
            key_mask = F.broadcast_like(key_mask, score)
            score = F.where(key_mask, score, F.ones_like(score) * -1e9)
        return score

    def _compute_attn_score(
        self,
        F,
        q: Tensor,
        k: Tensor,
        mask: Optional[Tensor],
    ) -> Tensor:
        score = F.batch_dot(lhs=q, rhs=k, transpose_b=True)
        score = self._apply_mask(F, score, mask)
        score = score / (math.sqrt(self.d_head) * self.temperature)
        score = F.softmax(score, axis=-1)
        score = self.dropout(score)
        return score

    def _compute_attn_output(self, F, score: Tensor, v: Tensor) -> Tensor:
        v = F.batch_dot(score, v)
        v = self._merge_head(F, v)
        v = self.out_proj(v)
        return v

    def hybrid_forward(self, F, x: Tensor, mask: Optional[Tensor]) -> Tensor:
        q, k, v = self._compute_qkv(F, x)
        score = self._compute_attn_score(F, q, k, mask)
        v = self._compute_attn_output(F, score, v)
        return v


class TemporalFusionEncoder(HybridBlock):
    @validated()
    def __init__(
        self,
        context_length: int,
        prediction_length: int,
        d_input: int,
        d_hidden: int,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.context_length = context_length
        self.prediction_length = prediction_length
        with self.name_scope():
            self.encoder_lstm = rnn.HybridSequentialRNNCell(prefix="encoder_")
            self.encoder_lstm.add(
                rnn.LSTMCell(
                    hidden_size=d_hidden,
                    input_size=d_input,
                )
            )
            self.decoder_lstm = rnn.HybridSequentialRNNCell(prefix="decoder_")
            self.decoder_lstm.add(
                rnn.LSTMCell(
                    hidden_size=d_hidden,
                    input_size=d_input,
                )
            )
            self.gate = nn.HybridSequential()
            self.gate.add(
                nn.Dense(units=d_hidden * 2, in_units=d_hidden, flatten=False)
            )
            self.gate.add(GatedLinearUnit(axis=-1, nonlinear=False))
            if d_input != d_hidden:
                self.skip_proj = nn.Dense(
                    units=d_hidden, in_units=d_input, flatten=False
                )
                self.add_skip = True
            else:
                self.add_skip = False
            self.lnorm = nn.LayerNorm(axis=-1, in_channels=d_hidden)

    def hybrid_forward(
        self,
        F,
        ctx_input: Tensor,
        tgt_input: Tensor,
        states: List[Tensor],
    ) -> Tensor:
        ctx_encodings, states = self.encoder_lstm.unroll(
            length=self.context_length,
            inputs=ctx_input,
            begin_state=states,
            merge_outputs=True,
        )
        tgt_encodings, _ = self.decoder_lstm.unroll(
            length=self.prediction_length,
            inputs=tgt_input,
            begin_state=states,
            merge_outputs=True,
        )
        encodings = F.concat(ctx_encodings, tgt_encodings, dim=1)
        skip = F.concat(ctx_input, tgt_input, dim=1)
        if self.add_skip:
            skip = self.skip_proj(skip)
        encodings = self.gate(encodings)
        encodings = self.lnorm(F.broadcast_add(skip, encodings))
        return encodings


class TemporalFusionDecoder(HybridBlock):
    @validated()
    def __init__(
        self,
        context_length: int,
        prediction_length: int,
        d_hidden: int,
        d_var: int,
        n_head: int,
        dropout: float = 0.0,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.context_length = context_length
        self.prediction_length = prediction_length

        with self.name_scope():
            self.enrich = GatedResidualNetwork(
                d_hidden=d_hidden,
                d_static=d_var,
                dropout=dropout,
            )
            self.attention = SelfAttention(
                context_length=context_length,
                prediction_length=prediction_length,
                d_hidden=d_hidden,
                n_head=n_head,
                share_values=True,
                dropout=dropout,
            )
            self.att_net = nn.HybridSequential(prefix="attention_")
            self.att_net.add(nn.Dropout(dropout))
            self.att_net.add(
                nn.Dense(
                    units=d_hidden * 2,
                    in_units=d_hidden,
                    flatten=False,
                    weight_initializer=init.Xavier(),
                )
            )
            self.att_net.add(
                GatedLinearUnit(
                    axis=-1,
                    nonlinear=False,
                )
            )
            self.att_lnorm = nn.LayerNorm(
                axis=-1,
                in_channels=d_hidden,
            )
            self.ff_net = nn.HybridSequential()
            self.ff_net.add(
                GatedResidualNetwork(
                    d_hidden,
                    dropout=dropout,
                )
            )
            self.ff_net.add(
                nn.Dense(
                    units=d_hidden * 2,
                    in_units=d_hidden,
                    flatten=False,
                    weight_initializer=init.Xavier(),
                )
            )
            self.ff_net.add(
                GatedLinearUnit(
                    axis=-1,
                    nonlinear=False,
                )
            )
            self.ff_lnorm = nn.LayerNorm(axis=-1, in_channels=d_hidden)

    def hybrid_forward(
        self, F, x: Tensor, static: Tensor, mask: Tensor
    ) -> Tensor:
        static = F.tile(
            static, reps=(1, self.context_length + self.prediction_length, 1)
        )
        skip = F.slice_axis(x, axis=1, begin=self.context_length, end=None)
        x = self.enrich(x, static)
        mask_pad = F.slice_axis(F.ones_like(mask), axis=1, begin=0, end=1)
        mask_pad = F.tile(mask_pad, reps=(1, self.prediction_length))
        mask = F.concat(mask, mask_pad, dim=1)
        att = self.attention(x, mask)
        att = self.att_net(att)
        x = F.slice_axis(x, axis=1, begin=self.context_length, end=None)
        x = self.att_lnorm(F.broadcast_add(x, att))
        x = self.ff_net(x)
        x = self.ff_lnorm(F.broadcast_add(x, skip))
        return x
