# --------------------------------------------------------------------------------
# 		Import
# --------------------------------------------------------------------------------
import torch
import torch.nn as nn
from ocr_models.language_models.sublayers import Embedder, Norm


# --------------------------------------------------------------------------------
# 		Class
# --------------------------------------------------------------------------------
class LSTMEncoder(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_layer, bidirectional):
        super(LSTMEncoder, self).__init__()
        input_dim = 288
        # Define LSTM layer
        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layer, bidirectional=bidirectional)

    def forward(self, src):
        # src(batch, channel, height, width)
        all_outputs = []
        for row in range(src.size(2)):
            # Covert row size into [seq_length, batch_size, embed_length]
            # (b_size, channel, width) -> (b_size, width, channel) -> (width, b_size, channel)
            inp = src[:, :, row, :].transpose(1, 2).transpose(0, 1)
            output, (hidden_state, hidden_cell) = self.lstm(inp)
            all_outputs.append(output)

        all_outputs = torch.cat(all_outputs, 0)

        return all_outputs, hidden_state, hidden_cell


class GlobalAttentionCell(nn.Module):
    def __init__(self, hidden_dim):
        super(GlobalAttentionCell, self).__init__()

        self.hidden_dim = hidden_dim
        self.linear_in = nn.Linear(hidden_dim, hidden_dim, bias=False)
        self.linear_out = nn.Linear(2 * hidden_dim, hidden_dim)
        self.softmax = nn.Softmax(dim=1)
        self.tanh = nn.Tanh()

    def score(self, h_t, h_s):
        """
        :param h_t: (batch, 1, hidden_dim)
        :param h_s: (batch, src_len, hidden_dim)
        :return: Attention score
        """
        batch, _, hidden_dim = h_t.size()
        # h_t^T W_a
        h_t = self.linear_in(h_t)
        # (batch, hidden_dim, src_len)
        h_s = h_s.transpose(1, 2)
        # (batch, 1, hidden_dim) * (batch, hidden_dim, src_len) --> (batch, 1, s_len)
        return torch.bmm(h_t, h_s)

    def forward(self, inputs, context, context_lengths):
        """
        :param inputs: (batch, 1, hidden_dim)
        :param context: (batch, src_len, hidden_dim)
        :param context_lengths: The source context length
        :return:
        """
        # (batch, 1, src_len)
        align = self.score(inputs, context) # (batch, 1, enc_len)
        batch, _, src_len = align.size()
        align_vectors = self.softmax(align.view(batch, src_len))
        align_vectors = align_vectors.view(batch, 1, src_len)
        # (batch, 1, src_len) * (batch, src_len, hidden_dim) -> (batch, 1, hidden_dim)
        c = torch.bmm(align_vectors, context)

        # \hat{h_t} = tanh(W [c_t, h_t])
        concat_c = torch.cat([c, inputs], 2).view(batch, self.hidden_dim * 2)
        attn_h = self.tanh(self.linear_out(concat_c).view(batch, 1, self.hidden_dim))

        # Transpose will make it non-contiguous
        attn_h = attn_h.transpose(0, 1).contiguous()
        align_vectors = align_vectors.transpose(0, 1).contiguous()
        # (1, batch, hidden_dim) and (1, batch, src_len)
        return attn_h, align_vectors


class LSTMDecoder(nn.Module):
    def __init__(self, vocab_size, input_dim, hidden_dim, num_layer, bidirectional, dropout=0.2):
        super(LSTMDecoder, self).__init__()

        self.vocab_size = vocab_size
        self.hidden_dim = hidden_dim
        self.embed = Embedder(vocab_size, input_dim)
        self.norm = Norm(input_dim)

        # Define LSTM cell
        self.lstm_cell = nn.LSTMCell(input_dim, hidden_dim)
        self.linear_out = nn.Linear(hidden_dim, vocab_size)
        self.attn = GlobalAttentionCell(hidden_dim)
        # self.drop = nn.Dropout(dropout)

    def forward(self, inputs, init_hidden_state, init_hidden_cell,
                context, context_lengths, is_train=True):
        """
        :param inputs: (batch_size, tgt_len, d) --> (tgt_len, batch, d)
        :param init_hidden_state:
        :param init_hidden_cell:
        :param context: (src_len, batch_size, hidden_size), outputs of encoder
        :param context_lengths:
        :return:
        """
        emb = self.embed(inputs)
        emb = self.norm(emb)
        emb = emb.transpose(0, 1)
        tgt_len, batch, _ = emb.size()

        decoder_output = torch.zeros(tgt_len, batch, self.hidden_dim).cuda()

        # hx = init_hidden_state.squeeze(0)
        # cx = init_hidden_cell.squeeze(0)
        hx = torch.cuda.FloatTensor(batch, self.hidden_dim).fill_(0)
        cx = torch.cuda.FloatTensor(batch, self.hidden_dim).fill_(0)

        if is_train:
            for i in range(emb.size(0)):
                h_s, cx = self.lstm_cell(emb[i], (hx, cx))
                # (batch, hidden_dim) -> (batch, 1, hidden_dim)
                hx, _ = self.attn(h_s.unsqueeze_(1),
                                  context.transpose(0, 1).contiguous(),
                                  context_lengths=context_lengths)
                hx = hx.squeeze(0)
                decoder_output[i, :, :] = hx

            outputs = self.linear_out(decoder_output)
        # else:
        #     # SOS token
        #     target = torch.cuda.LongTensor(1).fill_(self.vocab_size - 2)
        #     # Probabilities
        #     probs = torch.cuda.FloatTensor(batch, tgt_len, self.vocab_size).fill_(0)
        #     for i in range(emb.size(0)):
        #         h_s, cx = self.lstm_cell(emb[i], (hx, cx))
        #         hx, _ = self.attn(h_s.unsqueeze_(1),
        #                           context.transpose(0, 1).contiguous(),
        #                           context_lengths=context_lengths)
        #         hx = hx.squeeze(0)



        return outputs


class LSTM(nn.Module):
    def __init__(self, input_dim, enc_hidden_dim, enc_bidirectional, dec_hidden_dim, dec_bidirectional,
                 num_layer, vocab_size):
        super(LSTM, self).__init__()

        self.encoder = LSTMEncoder(input_dim, enc_hidden_dim, num_layer, enc_bidirectional)

        self.decoder = LSTMDecoder(vocab_size, input_dim, dec_hidden_dim, num_layer, dec_bidirectional)

    def forward(self, src, trg):
        context, hidden_state, hidden_cell = self.encoder(src)
        hidden_state, hidden_cell = self.init_decoder(hidden_state, hidden_cell)
        output = self.decoder(trg, hidden_state, hidden_cell, context, context.size(1))

        return output

    def init_decoder(self, enc_hidden_state, enc_hidden_cell):
        direc, batch_size, enc_dim = enc_hidden_state.size()
        enc_hidden_state = enc_hidden_state.contiguous().view(1, batch_size, -1)
        enc_hidden_cell = enc_hidden_cell.contiguous().view(1, batch_size, -1)

        return enc_hidden_state, enc_hidden_cell
