# 设置field，用来负责处理指定列的文本数据
SRC_TEXT = Field(sequential=True, init_token=SOS_TOKEN, eos_token=EOS_TOKEN, tokenize=tokenize,include_lengths=True)

# src和target使用同样的field，因为它们处理的都是中文文本，如果是翻译任务则需要针对不同的语言构造不同的field
train_data, val_data = TabularDataset.splits(
        path='./datasets/', train='small_train.tsv',
        validation='demo.tsv', format='tsv',skip_header=True,
        fields=[('src', SRC_TEXT), ('target', SRC_TEXT)])

# SRC_TEXT会用它负责的列的文本构造词典，上面就是用src和target列
SRC_TEXT.build_vocab(train_data)

class EncoderRNN(nn.Module):
    def __init__(self, input_size, hidden_size, n_layers=1, dropout=0.1):
        super(EncoderRNN, self).__init__()

        self.input_size = input_size
        self.hidden_size = hidden_size
        self.n_layers = n_layers
        self.dropout = dropout

        self.embedding = nn.Embedding(input_size, hidden_size)
        self.gru = nn.GRU(hidden_size, hidden_size, n_layers, bidirectional=True, dropout=self.dropout)

    def forward(self, input_seqs, input_lengths, hidden=None):
        total_length = input_seqs.size()[0]
        embedded = self.embedding(input_seqs)
        packed = torch.nn.utils.rnn.pack_padded_sequence(embedded, input_lengths)
        self.gru.flatten_parameters()
        outputs, hidden = self.gru(packed, hidden)
        # unpack (back to padded)
        outputs, output_lengths = torch.nn.utils.rnn.pad_packed_sequence(outputs, total_length=total_length)
        outputs = outputs[:, :, :self.hidden_size] + outputs[:, :, self.hidden_size:]  # Sum bidirectional outputs
        return outputs, hidden

class Attn(nn.Module):
    def __init__(self, method, hidden_size):
        super(Attn, self).__init__()

        self.method = method
        self.hidden_size = hidden_size

        if self.method == 'general':
            self.attn = nn.Linear(self.hidden_size, hidden_size)

        elif self.method == 'concat':
            self.attn = nn.Linear(self.hidden_size * 2, hidden_size)
            self.v = nn.Parameter(torch.FloatTensor(1, hidden_size))
            torch.nn.init.xavier_normal_(self.v)

    def forward(self, hidden, encoder_outputs):
        # Create variable to store attention energies
        attn_energies = self.score(hidden, encoder_outputs)

        # Normalize energies to weights in range 0 to 1, resize to 1 x B x S
        return F.softmax(attn_energies, dim=-1)

    def score(self, hidden, encoder_output):
        '''
        compute an attention score, 注意，多层的gru只有最底层需要计算attention score
        :param hidden:  (n_layers, batch_size,hidden_dim) deocder的hidden
        :param encoder_output: max_len x batch_size x hidden_size
        :param decoder_output: 1 * batch * hidden_size
        :return: a score (batch_size,1, max_len)
        '''
        if self.method == 'dot':
            energy = hidden[0].unsqueeze(1).bmm(encoder_output.permute(1, 2, 0))  # B 1 D * B D S = B 1 S
            return energy

        elif self.method == 'general':
            max_len, batch_size, _ = encoder_output.size()
            encoder_output = encoder_output.view(batch_size * max_len, -1)
            encoder_output = self.attn(encoder_output)
            encoder_output = encoder_output.view(max_len, batch_size, -1).transpose(0, 1)  # S B D->B S D
            energy = hidden[0].unsqueeze(1).bmm(encoder_output.transpose(1, 2))  # B 1 D * B D S   = B 1 S
            return energy

        elif self.method == 'concat':
            hidden = hidden[0].unsqueeze(1).repeat(1, encoder_output.size()[0], 1)  # B max_len hidden_dim
            concated = torch.cat((hidden, encoder_output.transpose(0, 1)), dim=-1).view((-1, hidden.size()[-1] * 2))
            energy = self.attn(concated)
            energy = energy.view(hidden.size())
            energy = energy.transpose(1, 2)
            return self.v.matmul(energy)  # B ! D * B D  S = B 1 S

class LuongAttnDecoderRNN(nn.Module):
    def __init__(self, attn_model, hidden_size, output_size, n_layers=1, dropout=0.1):
        super(LuongAttnDecoderRNN, self).__init__()

        # Keep for reference
        self.attn_model = attn_model
        self.hidden_size = hidden_size
        self.output_size = output_size
        self.n_layers = n_layers
        self.dropout = dropout

        # Define layers
        self.embedding = nn.Embedding(output_size, hidden_size)
        self.gru = nn.GRU(hidden_size, hidden_size, n_layers,dropout=dropout)
        self.embedding_dropout = nn.Dropout(dropout)
        self.concat = nn.Linear(hidden_size * 2, hidden_size)
        self.out = nn.Linear(hidden_size, output_size)

        # Choose attention model
        if attn_model != 'none':
            self.attn = Attn(attn_model, hidden_size)

    def forward(self, input_seq, context, encoder_outputs):
        # Get the embedding of the current input word (last output word)
        input_seq = input_seq.unsqueeze(0)
        embedded = self.embedding(input_seq)
        embedded = self.embedding_dropout(embedded)
        self.gru.flatten_parameters()
        rnn_output, hidden = self.gru(embedded, context)

        # Calculate attention from current RNN state and all encoder outputs;
        # apply to encoder outputs to get weighted average
        attn_weights = self.attn(hidden, encoder_outputs)  # B 1 S
        context = attn_weights.bmm(encoder_outputs.transpose(1, 0))  # B 1 H

        # Attentional vector using the RNN hidden state and context vector
        # concatenated together (Luong eq. 5)
        context = context.transpose(0, 1)  # 1 x B x N
        concat_input = torch.cat((rnn_output, context), -1).squeeze(0)  # B × 2*N
        concat_output = torch.tanh(self.concat(concat_input))

        # Finally predict next token (Luong eq. 6, without softmax)
        output = self.out(concat_output)
        rnn_output = self.out(rnn_output.squeeze(0))
        return output, hidden, attn_weights

class Seq2Seq(nn.Module):
    def __init__(self, vocabs, hidden_size, n_layers, dropout, max_len, attn_model='concat',
                 sos=0, eos=0, **kwargs):
        super(Seq2Seq, self).__init__()
        encoder = EncoderRNN(len(vocabs['src_vocab']), hidden_size, n_layers, dropout)
        decoder = LuongAttnDecoderRNN(attn_model, hidden_size, len(vocabs['tgt_vocab']), n_layers, dropout)

        self.encoder = encoder
        self.decoder = decoder
        self.sos = sos
        self.eos = eos
        self.max_len = max_len

        assert encoder.hidden_size == decoder.hidden_size, \
            "Hidden dimensions of encoder and decoder must be equal!"
        assert encoder.n_layers == decoder.n_layers, \
            "Encoder and decoder must have equal number of layers!"

    def forward(self, src, tgt=None, src_lengths=None, tgt_lengths=None, teacher_forcing_ratio=1):
        if src_lengths is not None and len(src.size())>1:
            src_lengths = src_lengths.squeeze(0)
        encoder_outputs, encoder_hidden = self.encoder(src, src_lengths, None)

        # Prepare input and output variables
        if tgt is not None:
            decoder_input = tgt[0, :]
        else:
            this_batch_size = src.size()[1]
            decoder_input = torch.LongTensor([self.sos] * this_batch_size).cuda()
        # 对于多层的gru，要排除掉后向的hidden，只使用前向的hidden
        decoder_hidden = encoder_hidden[-self.decoder.n_layers * 2::2].contiguous()  # Use last (forward) hidden state

        max_target_length = tgt.size()[0] if tgt is not None else self.max_len
        decoder_outputs = []

        # Run through decoder one time step at a time
        use_teacher_forcing = random.random() < teacher_forcing_ratio
        for t in range(max_target_length):
            decoder_output, decoder_hidden, decoder_attn = self.decoder(
                decoder_input, decoder_hidden, encoder_outputs
            )

            decoder_outputs.append(decoder_output)
            decoder_input = tgt[t] if use_teacher_forcing else decoder_output.argmax(-1)
        return torch.stack(decoder_outputs)