#
#  nn/modules/lstm.py
#  bxtorch
#
#  Created by Oliver Borchert on June 13, 2019.
#  Copyright (c) 2019 Oliver Borchert. All rights reserved.
#

import numpy as np
import torch
import torch.nn as nn
import torch.jit as jit
from typing import List, Tuple, Optional

class StackedLSTM(nn.Module):
    """
    The stacked LSTM is an extension to PyTorch's native LSTM allowing stacked
    LSTMs with different hidden dimensions being stacked.
    """

    batch_first: jit.Final[bool]

    # MARK: Initialization
    def __init__(self, input_size, hidden_sizes, bias=True, batch_first=False,
                 cudnn=True):
        """
        Initializes a new stacked LSTM according to the given parameters.

        Parameters:
        -----------
        - input_size: int
            The dimension of the sequence's elements.
        - hidden_sizes: list of int
            The dimensions of the stacked LSTM's layers.
        - bias: bool, default: True
            Whether to use biases in the LSTM.
        - batch_first: bool, default: False
            Whether the batch or the sequence can be found in the first
            dimension.
        - cudnn: bool, default: True
            Whether to use PyTorch's LSTM implementation which uses cuDNN on
            Nvidia GPUs. You usually don't want to change the default value,
            however, PyTorch's default implementation does not allow higher-
            order gradients of the LSTMCell as of version 1.1.0. When this
            value is set to False, we therefore use a (slower) implementation
            of a LSTM cell which allows higher-order gradients.
        """
        super().__init__()
        self.batch_first = batch_first
        self.stacked_cell = StackedLSTMCell(
            input_size, hidden_sizes, bias=bias, cudnn=cudnn
        )

    # MARK: Instance Methods
    def forward(self, 
                inputs: torch.Tensor,
                initial_states: Optional[
                    List[Tuple[torch.Tensor, torch.Tensor]]] = None, 
                return_sequence: bool = True):
        """
        Computes the forward pass through the stacked LSTM.

        Parameters:
        -----------
        - inputs: torch.FloatTensor [S, B, N]
            The inputs fed to the LSTM one after the other. Sequence length S,
            batch size B, and input size N. If ``batch_first`` is set to True,
            the first and second dimension should be swapped.
        - initial_states: list of tuple of 
                (torch.FloatTensor [H_i], torch.FloatTensor [H_i]),
                default: None
            The initial states for all LSTM layers. The length of the list must
            match the number of layers in the LSTM, the sizes of the states must
            match the hidden sizes of the LSTM layers. If None is given, the
            initial states are defaulted to all zeros.
        - return_sequence: bool, default: True
            Whether to return all outputs from the last LSTM layer or only the
            last one.

        Returns:
        --------
        - torch.FloatTensor [S, B, K] or torch.FloatTensor [B, K]
            Depending on whether sequences are returned, either all outputs
            or only the output from the last cell are returned. If the stacked
            LSTM was initialized with ``batch_first``, the first and second
            dimension are swapped when sequences are returned.
        """
        if self.batch_first:
            inputs = inputs.transpose(1, 0)

        sequence_length = inputs.size(0)

        # Initialize the state to empty vectors is needed for jit to properly
        # compile the function
        if initial_states is None:
            states = [(torch.empty(0), torch.empty(0))]
        else:
            states = initial_states

        # Iterate over sequence
        outputs = []
        for n in range(sequence_length):
            output, states = self.stacked_cell(
                inputs[n], None if states[0][0].size(0) == 0 else states
            )
            if return_sequence or n == sequence_length - 1:
                outputs.append(output)

        if return_sequence:
            result = torch.stack(outputs)
            if self.batch_first: 
                # set batch first, sequence length second
                result = result.transpose(1, 0)
            return result
        else:
            return outputs[0]


class StackedLSTMCell(nn.Module):
    """
    Actually, the StackedLSTMCell can easily be constructed from existing 
    modules, however, a bug in PyTorch's JIT compiler prevents implementing
    anything where a stacked LSTM is used within a loop (see the following 
    issue: https://github.com/pytorch/pytorch/issues/18143). Hence, this
    class provides a single time step for a stacked LSTM.
    """

    cells: jit.Final[int]
    num_stacked: jit.Final[int]

    # MARK: Initialization
    def __init__(self, input_size, hidden_sizes, bias=True, cudnn=True):
        """
        Initializes a new stacked LSTM cell.

        Parameters:
        -----------
        - input_size: int
            The dimension of the input variables.
        - hidden_sizes: list of int
            The hidden dimension of the stacked LSTMs.
        - bias: bool, default: True
            Whether to use a bias term for the LSTM implementation.
        - cudnn: bool, default: True
            Whether to not use cuDNN. In almost all cases, you don't want to
            set this value to false, however, you will need to change it if
            you want to compute higher-order derivatives of a network with
            a stacked LSTM cell.
        """
        super().__init__()

        self.num_stacked = len(hidden_sizes)

        cell_class = nn.LSTMCell if cudnn else _LSTMCell

        cells = []
        dims = zip([input_size] + hidden_sizes, hidden_sizes)
        for i, (in_dim, out_dim) in enumerate(dims):
            cells.append(cell_class(in_dim, out_dim, bias=bias))
        self.cells = nn.ModuleList(cells)
        
    # MARK: Instance Methods
    def forward(self,
                x: torch.Tensor,
                initial_states: Optional[
                    List[Tuple[torch.Tensor, torch.Tensor]]] = None):
        """
        Computes the new hidden states and cell states for each stacked cell.

        Parameters:
        -----------
        - x: torch.Tensor [B, N]
            The input with batch size B and dimension N.
        - states: list of tuple of (torch.Tensor [B, D], torch.Tensor [B, D]),
                default: None
            The states for each of the cells where each state is expected to
            have a size with batch size B and (respective) hidden dimension D.

        Returns:
        --------
        - torch.Tensor [B, D]
            The output, i.e. the hidden state of the deepest cell. Only given
            for convenience as it can be extracted from the other return value.
        - list of tuple of (torch.Tensor [B, D], torch.Tensor [B, D])
            The new hidden states and cell states for all cells.
        """
        if initial_states is None:
            # JIT Compatibility
            states = [
                (torch.empty(0), torch.empty(0))
                for _ in range(self.num_stacked)
            ]
        else:
            states = initial_states
        
        i = 0
        for cell in self.cells:
            x, next_cell = cell(
                x, None if states[i][0].size(0) == 0 else states[i]
            )
            states[i] = (x, next_cell)
            i += 1

        return x, states


class _LSTMCell(nn.Module):
    """
    LSTMCell which does not have cuDNN support but allows for higher-order
    gradients.
    Consult PyTorch's LSTMCell for documentation on the class's initialization
    parameters and how to call it.
    """

    hidden_size: jit.Final[int]
    has_bias: jit.Final[bool]

    # MARK: Initialization
    def __init__(self, input_size, hidden_size, bias=True):
        super().__init__()

        self.hidden_size = hidden_size

        self.input_weight = nn.Parameter(
            torch.FloatTensor(input_size, 4 * hidden_size)
        )
        self.hidden_weight = nn.Parameter(
            torch.FloatTensor(hidden_size, 4 * hidden_size)
        )

        if bias:
            self.bias = nn.Parameter(torch.FloatTensor(4 * hidden_size))
            self.has_bias = True
        else:
            self.has_bias = False

        self.reset_parameters()

    # MARK: Instance Methods
    def reset_parameters(self):
        sqrt_hidden = np.sqrt(1 / self.hidden_size)
        init_from = (-sqrt_hidden, sqrt_hidden)
        for p in self.parameters():
            nn.init.uniform_(p, *init_from)

    def forward(self, 
                x_in: torch.Tensor,
                state: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):

        if state is None:
            size = (x_in.size(0), self.hidden_size)
            hidden_state = torch.zeros(
                *size, dtype=torch.float, device=x_in.device
            )
            cell_state = torch.zeros(
                *size, dtype=torch.float, device=x_in.device
            )
        else:
            hidden_state, cell_state = state

        # 1) Perform matrix multiplications for input and last hidden state
        if self.has_bias:
            x = torch.addmm(self.bias, x_in, self.input_weight)
            h = torch.addmm(self.bias, hidden_state, self.hidden_weight)
        else:
            x = x_in.matmul(self.input_weight)
            h = hidden_state.matmul(self.hidden_weight)

        forget_gate_x, input_gate_x_1, input_gate_x_2, output_gate_x = \
            x.split(self.hidden_size, dim=1)
        forget_gate_h, input_gate_h_1, input_gate_h_2, output_gate_h = \
            h.split(self.hidden_size, dim=1)

        # 2) Forget gate
        forget_gate = torch.sigmoid(forget_gate_x + forget_gate_h)

        # 3) Input gate
        input_gate_1 = torch.sigmoid(input_gate_x_1 + input_gate_h_1)
        input_gate_2 = torch.tanh(input_gate_x_2 + input_gate_h_2)
        input_gate = forget_gate * cell_state + input_gate_1 * input_gate_2

        # 4) Output gate
        output_gate_1 = torch.sigmoid(output_gate_x + output_gate_h)
        output_gate = output_gate_1 * torch.tanh(input_gate)

        return output_gate, input_gate
