# -*- coding:utf-8 -*-

import math

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


class Bert(nn.Module):
    def __init__(self, d_hidden, n_heads, n_layers, n_groups, vocab_size, max_num_tokens):
        super().__init__()
        self.d_hidden = d_hidden
        self.n_layers = n_layers
        self.n_heads = n_heads
        self.n_groups = n_groups

        self.embedding = Embedding(vocab_size=vocab_size, dims=d_hidden, max_num_tokens=max_num_tokens, dropout=0.1)

        assert n_layers % n_groups == 0 and n_layers >= n_groups
        self.transformers = nn.ModuleList(
            [Transformer(d_hidden, n_heads, d_hidden * 4, dropout=0.1) for _ in range(n_groups)])

    def forward(self, x):
        mask = (x > 0).unsqueeze(1).repeat(1, x.size(1), 1).unsqueeze(1)
        x = self.embedding(x)
        #x = self.expand_fc(x)  # 차원 확장

        for i in range(self.n_layers):  # 각 그룹에 해당하는 Bert 적용
            x = self.transformers[i // (self.n_layers // self.n_groups)].forward(x, mask)
        return x


class Embedding(nn.Module):
    def __init__(self, vocab_size, dims, max_num_tokens, dropout):
        super().__init__()
        self.dims = dims
        self.pe = self.position_embedding(max_num_tokens, dims)  # position embedding
        self.te = nn.Embedding(vocab_size, dims, padding_idx=0)  # token embedding
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        # pe = self.position_embedding(x.shape[-1], self.dims).to(x.device)
        x = self.te(x) + self.pe.to(x.device)
        return self.dropout(x)

    def position_embedding(self, max_num_tokens, dims):
        pe = torch.zeros(max_num_tokens, dims)
        pe.require_grad = False

        position = torch.arange(0, max_num_tokens).float().unsqueeze(1)
        div_term = (torch.arange(0, dims, 2).float() * -(math.log(10000.0) / dims)).exp()

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)

        pe = pe.unsqueeze(0)
        return pe


class Transformer(nn.Module):
    def __init__(self, d_hidden, n_heads, d_fc, dropout=0.1):
        super().__init__()
        self.attention = MultiHeadedAttention(n_heads, d_hidden)
        self.feed_forward = nn.Sequential(
                                nn.Linear(d_hidden, d_fc),
                                nn.GELU(),
                                nn.Dropout(dropout),
                                nn.Linear(d_fc, d_hidden)
                            )
        self.layer_norms = nn.ModuleList([nn.LayerNorm(d_hidden) for _ in range(2)])
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask):
        x = self.layer_norms[0](x)
        att = self.attention(x, x, x, mask=mask)
        x = x + self.dropout(att)

        x = self.layer_norms[1](x)
        att = self.feed_forward(x)
        x = x + self.dropout(att)

        return self.dropout(x)


class MultiHeadedAttention(nn.Module):
    def __init__(self, n_heads, d_hidden, dropout=0.1):
        super().__init__()
        assert d_hidden % n_heads == 0

        self.d_k = d_hidden // n_heads
        self.h = n_heads

        self.projection = nn.ModuleList([nn.Linear(d_hidden, d_hidden) for _ in range(3)])
        self.output_linear = nn.Linear(d_hidden, d_hidden)

        self.dropout = nn.Dropout(p=dropout)

    def forward(self, query, key, value, mask=None):
        batch_size = query.size(0)

        query, key, value = [l(x).view(batch_size, -1, self.h, self.d_k).transpose(1, 2)
                             for l, x in zip(self.projection, (query, key, value))]

        x, attn = self.attention(query, key, value, mask=mask, dropout=self.dropout)

        x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.d_k)

        return self.output_linear(x)

    def attention(self, query, key, value, mask=None, dropout=None):
        scores = torch.matmul(query, key.transpose(-2, -1)) \
                 / math.sqrt(query.size(-1))
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        p_attn = F.softmax(scores, dim=-1)

        if dropout is not None:
            p_attn = dropout(p_attn)

        return torch.matmul(p_attn, value), p_attn
