import torch
from torch import nn, Tensor
from typing import Optional, Tuple
from .utils import get_activation_func, dot_attention


class MultiheadAttention(nn.modules.Module):
    def __init__(self, d_model, num_heads, dropout=0.1, bias=True, add_bias_kv=False, add_zero_attn=False, kdim=None, vdim=None, initialization="normal"):
        super(MultiheadAttention, self).__init__()
        self.d_model = d_model
        self.kdim = kdim if kdim is not None else d_model
        self.vdim = vdim if vdim is not None else d_model
        self._qkv_same_embed_dim = self.kdim == d_model and self.vdim == d_model
        self.initialization = initialization

        self.num_heads = num_heads
        self.dropout = dropout
        self.head_dim = d_model // num_heads
        assert self.head_dim * num_heads == self.d_model, "d_model must be divisible by num_heads"

        if self._qkv_same_embed_dim is False:
            self.q_proj_weight = nn.parameter.Parameter(torch.empty(d_model, d_model))
            self.k_proj_weight = nn.parameter.Parameter(torch.empty(d_model, self.kdim))
            self.v_proj_weight = nn.parameter.Parameter(torch.empty(d_model, self.vdim))
            self.register_parameter('in_proj_weight', None)
        else:
            self.in_proj_weight = nn.parameter.Parameter(torch.empty(3 * d_model, d_model))
            self.register_parameter('q_proj_weight', None)
            self.register_parameter('k_proj_weight', None)
            self.register_parameter('v_proj_weight', None)

        if bias:
            self.in_proj_bias = nn.parameter.Parameter(torch.empty(3 * d_model))
        else:
            self.register_parameter('in_proj_bias', None)
            
        self.out_proj_weight = nn.parameter.Parameter(torch.empty(d_model, d_model))
        if bias:
            self.out_proj_bias = nn.parameter.Parameter(torch.empty(d_model))
        else:
            self.register_parameter('out_proj_bias', None)

        if add_bias_kv:
            self.bias_k = nn.parameter.Parameter(torch.empty(1, 1, d_model))
            self.bias_v = nn.parameter.Parameter(torch.empty(1, 1, d_model))
        else:
            self.bias_k = self.bias_v = None
        
        self.add_zero_attn = add_zero_attn
        
        self._reset_parameters()

    def _reset_parameters(self):
        if self._qkv_same_embed_dim:
            if self.initialization=="normal":
                nn.init.normal_(self.in_proj_weight, mean=0.0, std=0.02)
            elif self.initialization=="xavier_uniform":
                nn.init.xavier_uniform_(self.in_proj_weight)
        else:
            if self.initialization=="normal":
                nn.init.normal_(self.q_proj_weight, mean=0.0, std=0.02)
                nn.init.normal_(self.k_proj_weight, mean=0.0, std=0.02)
                nn.init.normal_(self.v_proj_weight, mean=0.0, std=0.02)
            elif self.initialization=="xavier_uniform":
                nn.init.xavier_uniform_(self.q_proj_weight)
                nn.init.xavier_uniform_(self.k_proj_weight)
                nn.init.xavier_uniform_(self.v_proj_weight)
                
        if self.initialization=="normal":
            nn.init.normal_(self.out_proj_weight, mean=0.0, std=0.02)
        elif self.initialization=="xavier_uniform":
            nn.init.xavier_uniform_(self.out_proj_weight)

        if self.in_proj_bias is not None:
            nn.init.constant_(self.in_proj_bias, 0.)
        if self.out_proj_bias is not None:
            nn.init.constant_(self.out_proj_bias, 0.)
        if self.bias_k is not None:
            nn.init.constant_(self.bias_k, 0.)
        if self.bias_v is not None:
            nn.init.constant_(self.bias_v, 0.)

    def __setstate__(self, state):
        # Support loading old MultiheadAttention checkpoints generated by v1.1.0
        if '_qkv_same_embed_dim' not in state:
            state['_qkv_same_embed_dim'] = True
        super(MultiheadAttention, self).__setstate__(state)

    def forward(self, query: Tensor, key: Tensor, value: Tensor, key_padding_mask: Optional[Tensor] = None,
                need_weights: bool = True, attn_mask: Optional[Tensor] = None) -> Tuple[Tensor, Optional[Tensor]]:
        # batch_first default
        query, key, value = [x.transpose(1, 0) for x in (query, key, value)]

        if not self._qkv_same_embed_dim:
            output, attention_weights = nn.functional.multi_head_attention_forward(
                query, key, value, self.d_model, self.num_heads,
                self.in_proj_weight, self.in_proj_bias,
                self.bias_k, self.bias_v, self.add_zero_attn,
                self.dropout, self.out_proj_weight, self.out_proj_bias,
                training=self.training,
                key_padding_mask=key_padding_mask, need_weights=need_weights,
                attn_mask=attn_mask, use_separate_proj_weight=True,
                q_proj_weight=self.q_proj_weight, k_proj_weight=self.k_proj_weight,
                v_proj_weight=self.v_proj_weight)
        else:
            output, attention_weights = nn.functional.multi_head_attention_forward(
                query, key, value, self.d_model, self.num_heads,
                self.in_proj_weight, self.in_proj_bias,
                self.bias_k, self.bias_v, self.add_zero_attn,
                self.dropout, self.out_proj_weight, self.out_proj_bias,
                training=self.training,
                key_padding_mask=key_padding_mask, need_weights=need_weights,
                attn_mask=attn_mask)
        # batch_first default
        output = output.transpose(1, 0)
        return output, attention_weights


class CodeAttention(nn.modules.Module):
    def __init__(self, m, d_model):
        super(CodeAttention, self).__init__()
        self.m = m
        self.d_model = d_model
        # parameters
        code_ids = torch.arange(self.m, dtype=torch.long)
        self.register_buffer("code_ids", code_ids)
        self.code_embedding_layer = nn.Embedding(self.m, self.d_model)

    def forward(self, key: Tensor, value: Tensor = None) -> Tuple[Tensor, Optional[Tensor]]:
        '''
        :param key: (batch_size, timesteps, d_model)
        :param value: (batch_size, timesteps, d_model)
        :return:
        '''
        if value is None: value = key
        batch_size = key.shape[0]
        # self.code_ids: (batch_size, )
        # code_ids: (batch_size, m)
        code_ids = self.code_ids.unsqueeze(0).expand(batch_size, self.m)
        # codes: (batch_size, m, d_model)
        codes = self.code_embedding_layer(code_ids)
        # context_embed: (batch_size, m, d_model)
        context_embed, context_attention_weight = dot_attention(query=codes, key=key, value=value)
        return context_embed, context_attention_weight


class PositionwiseFeedForward(nn.modules.Module):
    def __init__(self, d_model, d_ff, activation="gelu", dropout=0.1, bias=True, layer_norm_epsilon=1e-5,
                 initialization="normal"):
        super(PositionwiseFeedForward, self).__init__()
        self.d_model = d_model
        self.d_ff = d_ff
        self.activation = activation
        self.activation_func = get_activation_func(activation)
        self.dropout = dropout
        self.bias = bias
        self.layer_norm_epsilon = layer_norm_epsilon
        self.initialization = initialization

        self.f1_weight = nn.parameter.Parameter(torch.empty(d_ff, d_model))
        self.f2_weight = nn.parameter.Parameter(torch.empty(d_model, d_ff))
        if bias:
            self.f1_bias = nn.parameter.Parameter(torch.empty(d_ff))
            self.f2_bias = nn.parameter.Parameter(torch.empty(d_model))
        else:
            self.register_parameter('f1_bias', None)
            self.register_parameter('f2_bias', None)

        self.layer_normalization = nn.modules.normalization.LayerNorm(d_model, eps=layer_norm_epsilon)
        self.f1_dropout = nn.modules.dropout.Dropout(dropout)
        self.f2_dropout = nn.modules.dropout.Dropout(dropout)

        self._reset_parameters()

    def _reset_parameters(self):
        if self.initialization == "normal":
            nn.init.normal_(self.f1_weight, mean=0.0, std=0.02)
            nn.init.normal_(self.f2_weight, mean=0.0, std=0.02)
        elif self.initialization == "xavier_uniform":
            nn.init.xavier_uniform_(self.f1_weight)
            nn.init.xavier_uniform_(self.f2_weight)

        if self.f1_bias is not None:
            nn.init.constant_(self.f1_bias, 0.)
        if self.f2_bias is not None:
            nn.init.constant_(self.f2_bias, 0.)

    def forward(self, inputs: Tensor) -> Tensor:
        f1_output = self.f1_dropout(self.activation_func(nn.functional.linear(inputs, self.f1_weight, self.f1_bias)))
        f2_output = self.f2_dropout(nn.functional.linear(f1_output, self.f2_weight, self.f2_bias))
        residual_output = inputs + f2_output
        output = self.layer_normalization(residual_output)
        return output