import torch
from torch import nn as nn
from binlin.model.modules.attention import AttnBahd

from binlin.model.nn_utils import get_lstm_unit, get_gru_unit
from typing import Tuple

class DecoderLSTM(nn.Module):
    def __init__(self,
                 input_size: int,
                 hidden_size: int,
                 num_layers: int,
                 device: torch.device,
                 batch_first=False,
                 bidirectional: bool = False,
                 ):

        super(DecoderLSTM, self).__init__()
        self._input_size = input_size
        self._hidden_size = hidden_size
        self._num_layers = num_layers
        self._batch_first = batch_first
        self._bidirectional = bidirectional
        self._num_directions = 2 if bidirectional else 1
        self._device = device

        self._lstm = nn.LSTM(input_size=self._input_size,
                           hidden_size=self._hidden_size,
                           num_layers=self._num_layers,
                           batch_first=self._batch_first,
                           bidirectional=self._bidirectional)

    def forward(self,
                prev_y_batch: torch.Tensor,
                prev_h_batch: torch.Tensor,
                prev_c_batch: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        prev_y_batch, (prev_h_batch, prev_c_batch) = self._lstm(prev_y_batch, (prev_h_batch, prev_c_batch))


        return prev_y_batch, prev_h_batch, prev_c_batch

class DecoderGRUAttnBahd(nn.Module):
    """
    At each step, a decode uses a softmax function to compute a distribution over the output classes,
    and samples the highest-scoring one. At each step, the input to the decode is:
        - embedded output of the previous step (`prev_y_embedded`)
        - encoder outputs

    The decode computes attention weights over the encoder outpus, uses these weights to compute
    their weighted sum (call it `context`, or `ctx`). Then it combines `ctx` with `prev_y_embedded` and squeezes
    further to a lower dimension. Finally, the result is fed to an RNN-type cell.

    """
    def __init__(self,
                 input_size: int,
                 hidden_size: int,
                 num_layers: int,
                 device: torch.device,
                 encoder_hidden: int,
                 encoder_num_directions: int,
                 attn_input_size: int,
                 attn_output_size: int,
                 bidirectional: bool = False,
                 ):

        super(DecoderGRUAttnBahd, self).__init__()
        self._input_size = input_size
        self._hidden_size = hidden_size
        self._attn_input_size = attn_input_size # usually equals encoder_output_dim * 2
        self._attn_output_size = attn_output_size # usually equals decoder_hidden_dim
        self._enc_hidden_size = encoder_hidden
        self._enc_num_directions = encoder_num_directions
        self._num_layers = num_layers
        self._bidirectional = bidirectional
        self._num_directions = 2 if bidirectional else 1
        self._device = device

        combine_y_ctx_input_dim = self._input_size + self._enc_hidden_size * self._enc_num_directions
        combine_y_ctx_output_dim = self._hidden_size
        self.W_combine_prev_y_ctx = nn.Linear(combine_y_ctx_input_dim, combine_y_ctx_output_dim)

        self._gru = nn.GRU(input_size=combine_y_ctx_output_dim,
                           hidden_size=self._hidden_size,
                           num_layers=self._num_layers,
                           bidirectional=self._bidirectional)

        self.attn = AttnBahd(input_size=self._attn_input_size,
                             output_size=self._attn_output_size)

        self.relu = nn.ReLU()

    def forward(self,
                prev_y_batch: torch.Tensor,
                prev_h_batch: torch.Tensor,
                encoder_outputs_batch: torch.Tensor) -> Tuple[torch.Tensor,torch.Tensor,torch.Tensor]:

        """
        A forward step of the decode.
        NOTE: The step operates on one-step slice of the target sequence,
        and not on the whole sequence (unlike the encoder)!

        Args:
            prev_y_batch: torch.Tensor with embedded previous predictions of size (batch_size, self._input_size)
            prev_h_batch: torch.Tensor with the current decode state, size (1, batch_size, self._hidden_size)
            encoder_outputs_batch: torch.Tensor containing encoder outputs, size (src_seq_len, batch_size, enc_num_directions * enc_hidden).

        Returns:
            rnn_output: torch.Tensor with the output of the decode GRU cell, size (1, batch_size, dec_out_dim)
            rnn_hidden: torch.Tensor with the hidden state of the decode GRU cell, size (1, batch_size, dec_out_dim)
            attn_weights: torch.Tensor with the attention weights tensor, size (batch_size x src_seq_len)

        """

        # Calculate attention from current RNN state and all encoder outputs
        # size (batch_size x src_seq_len)
        attn_weights = self.attn(prev_h_batch, encoder_outputs_batch)

        # Apply attention weights to encoder outputs to get weighted average

        # size (batch_size, 1, encoder._num_directions * encoder._hidden_size)
        context = torch.bmm(attn_weights.unsqueeze(1), encoder_outputs_batch.transpose(0, 1))

        # size (batch_size, encoder._num_directions * encoder._hidden_size)
        context2D = context.squeeze(1)

        # Combine embedded input word and attended context
        # (batch_size, prev_y_dim + encoder._num_directions * encoder._hidden_size)
        y_ctx = torch.cat((prev_y_batch, context2D), 1)

        # run the combined input through RNN
        # size (batch_size x self._gru.hidden)
        gru_input = self.relu(self.W_combine_prev_y_ctx(y_ctx))

        # (1, batch_size, dec_out_dim), (1, batch_size, dec_out_dim)
        gru_output, gru_hidden = self._gru(gru_input.unsqueeze(0), prev_h_batch)

        # Return final output, hidden state, and attention weights (for visualization)
        return gru_output, gru_hidden, attn_weights