import torch.nn as nn
import torch
import torch.nn.functional as F
import numpy as np

# transformer modules
class AddNorm(nn.Module):
    """残差连接后进行层归一化"""

    def __init__(self, normalized, dropout):
        super(AddNorm, self).__init__()
        self.dropout = nn.Dropout(dropout)
        self.ln = nn.LayerNorm(normalized)

    def forward(self, x, y):
        return  self.ln(x + self.dropout(y)) 


class PositionWiseFFN(nn.Module):
    """基于位置的前馈⽹络"""

    def __init__(self, ffn_input, ffn_hiddens,mlp_bias=True):
        super(PositionWiseFFN, self).__init__()
        self.ffn = nn.Sequential(
            nn.Linear(ffn_input, ffn_hiddens, bias=mlp_bias),
            nn.ReLU(),
            nn.Linear(ffn_hiddens, ffn_input, bias=mlp_bias),
        )

    def forward(self, x):
        return self.ffn(x)
    
class PositionalEncoding(nn.Module):
    """位置编码"""

    def __init__(self, num_hiddens, dropout, max_len=1000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(dropout)
        # 创建⼀个⾜够⻓的P
        self.P = torch.zeros((1, max_len, num_hiddens))
        X = torch.arange(max_len, dtype=torch.float32).reshape(-1, 1) / torch.pow(10000, torch.arange(0, num_hiddens, 2,
                                                                                                      dtype=torch.float32) / num_hiddens)
        self.P[:, :, 0::2] = torch.sin(X)
        self.P[:, :, 1::2] = torch.cos(X)

    def forward(self, X):
        X = X + self.P[:, :X.shape[1], :].to(X.device)
        return self.dropout(X)
    
class AttentionEncode(nn.Module):

    def __init__(self, dropout, embedding_size, num_heads,ffn=False):
        super(AttentionEncode, self).__init__()
        self.dropout = dropout
        self.embedding_size = embedding_size
        self.num_heads = num_heads
        self.seq_len = 50
        self.is_ffn = ffn
        
        self.att = nn.MultiheadAttention(embed_dim=self.embedding_size,
                                         num_heads=num_heads,
                                         dropout=0.6
                                         )
    
        self.addNorm = AddNorm(normalized=[self.seq_len, self.embedding_size], dropout=self.dropout)

        self.FFN = PositionWiseFFN(ffn_input=self.embedding_size, ffn_hiddens=self.embedding_size*2)

    def forward(self, x):
        bs,_,_ = x.size()
        MHAtt, _ = self.att(x, x, x)
        MHAtt_encode = self.addNorm(x, MHAtt)

        if self.is_ffn:
            ffn_in = MHAtt_encode # bs,seq_len,feat_dims
            ffn_out = self.FFN(ffn_in)
            MHAtt_encode = self.addNorm(ffn_in,ffn_out)

        return MHAtt_encode

class FAN_encode(nn.Module):

    def __init__(self, dropout, shape):
        super(FAN_encode, self).__init__()
        self.dropout = dropout
        self.addNorm = AddNorm(normalized=[1, shape], dropout=self.dropout)
        self.FFN = PositionWiseFFN(ffn_input=shape, ffn_hiddens=(2*shape))
        self.ln = nn.LayerNorm(shape)

    def forward(self, x):
        #x = self.ln(x)
        ffn_out = self.FFN(x)
        encode_output = self.addNorm(x, ffn_out)

        return encode_output