#!/usr/bin/env python3
# coding=utf-8

import math

import torch
import torch.nn as nn
import torch.nn.functional as F

from transformers import AutoModel
from model.module.char_embedding import CharEmbedding


class WordDropout(nn.Dropout):
    def forward(self, input_tensor):
        if self.p == 0:
            return input_tensor

        ones = input_tensor.new_ones(input_tensor.shape[:-1])
        dropout_mask = torch.nn.functional.dropout(ones, self.p, self.training, inplace=False)

        return dropout_mask.unsqueeze(-1) * input_tensor


class Encoder(nn.Module):
    def __init__(self, args, dataset):
        super(Encoder, self).__init__()

        self.dim = args.hidden_size
        self.n_layers = args.n_encoder_layers
        self.width_factor = args.query_length

        self.bert = AutoModel.from_pretrained(args.encoder, add_pooling_layer=False)
        # self.bert._set_gradient_checkpointing(self.bert.encoder, value=True)
        if args.encoder_freeze_embedding:
            self.bert.embeddings.requires_grad_(False)
            self.bert.embeddings.LayerNorm.requires_grad_(True)

        if args.freeze_bert:
            self.bert.requires_grad_(False)

        self.use_char_embedding = args.char_embedding
        if self.use_char_embedding:
            self.form_char_embedding = CharEmbedding(dataset.char_form_vocab_size, args.char_embedding_size, self.dim)
            self.word_dropout = WordDropout(args.dropout_word)

        self.post_layer_norm = nn.LayerNorm(self.dim)
        self.subword_attention = nn.Linear(self.dim, 1)

        if self.width_factor > 1:
            self.query_generator = nn.Linear(self.dim, self.dim * self.width_factor)
        else:
            self.query_generator = nn.Identity()

        self.encoded_layer_norm = nn.LayerNorm(self.dim)
        self.scores = nn.Parameter(torch.zeros(self.n_layers, 1, 1, 1), requires_grad=True)

    def forward(self, bert_input, form_chars, to_scatter, n_words):
        tokens, mask = bert_input
        batch_size = tokens.size(0)

        encoded = self.bert(tokens, attention_mask=mask, output_hidden_states=True).hidden_states[1:]
        encoded = torch.stack(encoded, dim=0)  # shape: (12, B, T, H)
        encoded = self.encoded_layer_norm(encoded)

        if self.training:
            time_len = encoded.size(2)
            scores = self.scores.expand(-1, batch_size, time_len, -1)
            dropout = torch.empty(self.n_layers, batch_size, 1, 1, dtype=torch.bool, device=self.scores.device)
            dropout.bernoulli_(0.1)
            scores = scores.masked_fill(dropout, float("-inf"))
        else:
            scores = self.scores

        scores = F.softmax(scores, dim=0)
        encoded = (scores * encoded).sum(0)  # shape: (B, T, H)
        encoded = encoded.masked_fill(mask.unsqueeze(-1) == 0, 0.0)  # shape: (B, T, H)

        subword_attention = self.subword_attention(encoded) / math.sqrt(self.dim)  # shape: (B, T, 1)
        subword_attention = subword_attention.expand_as(to_scatter)  # shape: (B, T_subword, T_word)
        subword_attention = subword_attention.masked_fill(to_scatter == 0, float("-inf"))  # shape: (B, T_subword, T_word)
        subword_attention = torch.softmax(subword_attention, dim=1)  # shape: (B, T_subword, T_word)
        subword_attention = subword_attention.masked_fill(to_scatter.sum(1, keepdim=True) == 0, value=0.0)  # shape: (B, T_subword, T_word)

        encoder_output = torch.einsum("bsd,bsw->bwd", encoded, subword_attention)
        encoder_output = self.post_layer_norm(encoder_output)

        if self.use_char_embedding:
            form_char_embedding = self.form_char_embedding(form_chars[0], form_chars[1], form_chars[2])
            encoder_output = self.word_dropout(encoder_output) + form_char_embedding

        decoder_input = self.query_generator(encoder_output)
        decoder_input = decoder_input.view(batch_size, -1, self.width_factor, self.dim).flatten(1, 2)  # shape: (B, T*Q, D)

        return encoder_output, decoder_input