import copy
import torch
from torch import nn
import math

def swish(x):
    return x * torch.sigmoid(x)

def gelu(x):
    return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))

ACT2FN = {"relu": torch.nn.functional.relu, "swish": swish, "gelu": gelu}


class MultiHeadSelfAttention(nn.Module):
    def __init__(self, config):
        super(MultiHeadSelfAttention, self).__init__()
        if config.hidden_size % config.num_attention_heads != 0:
            raise ValueError(
                "The hidden size (%d) is not a multiple of the number of attention "
                "heads (%d)" % (config.hidden_size, config.num_attention_heads))
        self.num_heads = config.num_attention_heads
        self.head_dim = int(config.hidden_size / config.num_attention_heads)
        self.total_dim = self.num_heads * self.head_dim

        self.q_proj = nn.Linear(config.hidden_size, self.total_dim)
        self.k_proj = nn.Linear(config.hidden_size, self.total_dim)
        self.v_proj = nn.Linear(config.hidden_size, self.total_dim)

    def forward(self, hidden_states, attention_mask):
        '''
        hidden_states --> (B, S, H), where B = batch_size, S = sequence length, H = hidden_size
        attention_mask --> (B, S), where B = batch_size, S = sequence length, it is -1.e10 if there is a padding otherwise it is 0
        请补全transformer中的自注意力层，注意self.query、self.key、self.value、 self.dropout都需要在本层中合理使用
        context_layer --> (B, S, H) 
        '''

        # ==========
        # todo '''请补全代码'''
        # ==========
        pass


class FeedForwardIntermediate(nn.Module):
    def __init__(self, config):
        super(FeedForwardIntermediate, self).__init__()
        self.linear = nn.Linear(config.hidden_size, config.intermediate_size)
        self.act_fn = ACT2FN[config.hidden_act] \
            if isinstance(config.hidden_act, str) else config.hidden_act

    def forward(self, hidden_states):
        hidden_states = self.linear(hidden_states)
        hidden_states = self.act_fn(hidden_states)
        return hidden_states


class EncoderLayer(nn.Module):
    def __init__(self, config):
        super(EncoderLayer, self).__init__()
        self.attn = SelfAttention(config)
        self.ffn = FeedForwardIntermediate(config)
        self.out = FeedForwardOutput(config)

    def forward(self, hidden_states, attention_mask):
        attn_out = self.attn(hidden_states, attention_mask)
        ffn_out = self.ffn(attn_out)
        final_out = self.out(ffn_out, attn_out)
        return final_out


class SelfAttention(nn.Module):
    def __init__(self, config):
        super(SelfAttention, self).__init__()
        self.core = MultiHeadSelfAttention(config)
        self.proj = SelfOutputProjection(config)

    def forward(self, input_tensor, attention_mask):
        core_out = self.core(input_tensor, attention_mask)
        proj_out = self.proj(core_out, input_tensor)
        return proj_out


class LayerNormModule(nn.Module):
    def __init__(self, hidden_size, eps=1e-12):
        super(LayerNormModule, self).__init__()
        self.gamma = nn.Parameter(torch.ones(hidden_size))
        self.beta = nn.Parameter(torch.zeros(hidden_size))
        self.epsilon = eps

    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        var = (x - mean).pow(2).mean(-1, keepdim=True)
        x = (x - mean) / torch.sqrt(var + self.epsilon)
        return self.gamma * x + self.beta


class SelfOutputProjection(nn.Module):
    def __init__(self, config):
        super(SelfOutputProjection, self).__init__()
        self.linear = nn.Linear(config.hidden_size, config.hidden_size)
        self.ln = LayerNormModule(config.hidden_size, eps=1e-12)
        self.drop = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states, input_tensor):
        hidden_states = self.linear(hidden_states)
        hidden_states = self.drop(hidden_states)
        hidden_states = self.ln(hidden_states + input_tensor)
        return hidden_states


class FeedForwardOutput(nn.Module):
    def __init__(self, config):
        super(FeedForwardOutput, self).__init__()
        self.linear = nn.Linear(config.intermediate_size, config.hidden_size)
        self.ln = LayerNormModule(config.hidden_size, eps=1e-12)
        self.drop = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states, input_tensor):
        hidden_states = self.linear(hidden_states)
        hidden_states = self.drop(hidden_states)
        hidden_states = self.ln(hidden_states + input_tensor)
        return hidden_states


class EncoderStack(nn.Module):
    def __init__(self, config):
        super(EncoderStack, self).__init__()
        layer = EncoderLayer(config)
        self.layers = nn.ModuleList([copy.deepcopy(layer) for _ in range(config.num_hidden_layers)])

    def forward(self, hidden_states, attention_mask, output_all_layers=False):
        all_layers = []
        for layer_module in self.layers:
            hidden_states = layer_module(hidden_states, attention_mask)
            if output_all_layers:
                all_layers.append(hidden_states)
        if not output_all_layers:
            all_layers.append(hidden_states)
        return all_layers


class LMHeadProjection(nn.Module):
    def __init__(self, config, embedding_weights):
        super(LMHeadProjection, self).__init__()
        self.transform = PredictionHeadTransform(config)

        self.decoder = nn.Linear(embedding_weights.size(1),
                                 embedding_weights.size(0),
                                 bias=False)
        self.decoder.weight = embedding_weights
        self.bias = nn.Parameter(torch.zeros(embedding_weights.size(0)))

    def forward(self, hidden_states):
        hidden_states = self.transform(hidden_states)
        hidden_states = self.decoder(hidden_states) + self.bias
        return hidden_states


class InputEmbeddings(nn.Module):
    def __init__(self, config):
        super(InputEmbeddings, self).__init__()
        self.word_embed = nn.Embedding(config.vocab_size, config.hidden_size)
        self.type_embed = nn.Embedding(config.type_vocab_size, config.hidden_size)
        self.pos_embed = nn.Embedding(config.max_position_embeddings, config.hidden_size)

        self.ln = LayerNormModule(config.hidden_size, eps=1e-12)
        self.drop = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, input_ids, token_ids=None, position_ids=None):
        seq_len = input_ids.size(1)
        if position_ids is None:
            position_ids = torch.arange(seq_len, dtype=torch.long, device=input_ids.device)
            position_ids = position_ids.unsqueeze(0).expand_as(input_ids)
        if token_ids is None:
            token_ids = torch.zeros_like(input_ids)

        word_emb = self.word_embed(input_ids)
        type_emb = self.type_embed(token_ids)
        pos_emb = self.pos_embed(position_ids)

        embeddings = word_emb + pos_emb + type_emb
        embeddings = self.ln(embeddings)
        embeddings = self.drop(embeddings)
        return embeddings


class PoolerHead(nn.Module):
    def __init__(self, config):
        super(PoolerHead, self).__init__()
        self.linear = nn.Linear(config.hidden_size, config.hidden_size)
        self.act = nn.Tanh()

    def forward(self, hidden_states):
        first_token = hidden_states[:, 0]

        pooled = self.linear(first_token)
        pooled = self.act(pooled)
        return pooled


class PredictionHeadTransform(nn.Module):
    def __init__(self, config):
        super(PredictionHeadTransform, self).__init__()
        self.linear = nn.Linear(config.hidden_size, config.hidden_size)
        self.act_fn = ACT2FN[config.hidden_act] \
            if isinstance(config.hidden_act, str) else config.hidden_act
        self.ln = LayerNormModule(config.hidden_size, eps=1e-12)

    def forward(self, hidden_states):
        hidden_states = self.linear(hidden_states)
        hidden_states = self.act_fn(hidden_states)
        hidden_states = self.ln(hidden_states)
        return hidden_states


class MLMHead(nn.Module):
    def __init__(self, config, embedding_weights):
        super(MLMHead, self).__init__()
        self.pred = LMHeadProjection(config, embedding_weights)

    def forward(self, sequence_output):
        scores = self.pred(sequence_output)
        return scores