import torch
import torch.nn as nn
import math


def get_sinusoid_encoding_table(n_position, d_hid, padding_idx=None):
    '''Returns: [seq_len, d_hid]
    '''
    table = torch.zeros(n_position, d_hid)
    position = torch.arange(0, n_position, dtype=torch.float).unsqueeze(-1)
    sin_term = torch.pow(10000, -1 * torch.arange(0, d_hid, 2) / d_hid)
    cons_term = torch.pow(10000, -1 * torch.arange(1, d_hid, 2) / d_hid)
    table[..., ::2] = torch.sin(position * sin_term)
    table[..., 1::2] = torch.cos(position * cons_term)
    return table

    #其他实现方式1
    # embeddings_table = torch.zeros(n_position, d_hid)
    # position = torch.arange(0, n_position, dtype=torch.float).unsqueeze(1)
    # sin_div_term = torch.exp(torch.arange(0, d_hid, 2).float() * (-math.log(10000.0) / d_hid))
    # cos_div_term = torch.exp(torch.arange(1, d_hid, 2).float() * (-math.log(10000.0) / d_hid))
    # embeddings_table[:, 0::2] = torch.sin(position * sin_div_term)
    # embeddings_table[:, 1::2] = torch.cos(position * cos_div_term)
    # return embeddings_table

    # 其他实现方式2
    # position_ids = torch.arange(0, n_position).unsqueeze(1)
    # position_ids = position_ids.expand(-1, d_hid)
    # indices = torch.arange(0, d_hid)
    # position_ids = position_ids * torch.pow(10000, -2 * torch.true_divide(torch.floor_divide(indices, 2), d_hid))
    # position_ids[:, ::2] = torch.sin(position_ids[:, ::2])
    # position_ids[:, 1::2] = torch.cos(position_ids[:, 1::2])
    # return position_ids

class RelativePositionEmbedding(nn.Module):
    def __init__(self,
                 qlen,
                 klen,
                 embedding_size,
                 max_relative_position=127) -> None:
        super(RelativePositionEmbedding, self).__init__()
        vocab_size = max_relative_position * 2 + 1
        distance_mat = torch.arange(klen)[None, :] - torch.arange(qlen)[:, None]  # 列数-行数, [query_len, key_len]
        distance_mat_clipped = torch.clamp(distance_mat, -max_relative_position, max_relative_position)
        final_mat = distance_mat_clipped + max_relative_position

        # sinusoid_encoding编码的位置矩阵
        embeddings_table = get_sinusoid_encoding_table(vocab_size, embedding_size)

        # 实现方式1
        # flat_relative_positions_matrix = final_mat.view(-1)
        # one_hot_relative_positions_matrix = torch.nn.functional.one_hot(flat_relative_positions_matrix, num_classes=vocab_size).float()
        # position_embeddings = torch.matmul(one_hot_relative_positions_matrix, embeddings_table)
        # my_shape = list(final_mat.size())
        # my_shape.append(embedding_size)
        # position_embeddings = position_embeddings.view(my_shape)

        # 实现方式2
        # position_embeddings = torch.take_along_dim(embeddings_table, final_mat.flatten().unsqueeze(1), dim=0)
        # position_embeddings = position_embeddings.reshape(*final_mat.shape, embeddings_table.shape[-1])  # [seq_len, seq_len, hdsz]
        # self.register_buffer('position_embeddings', position_embeddings)
        
        # 实现方式3
        position_embeddings = nn.Embedding.from_pretrained(embeddings_table, freeze=True)(final_mat)
        self.register_buffer('position_embeddings', position_embeddings)

    def forward(self, qlen, klen):
        return self.position_embeddings[:qlen, :klen, :]
    
    
# class SinusoidalPositionEmbedding(nn.Module):
#     """定义Sin-Cos位置Embedding
#     """
#     def __init__(self, max_position, embedding_size):
#         super(SinusoidalPositionEmbedding, self).__init__()
#         position_embeddings = nn.Embedding.from_pretrained(get_sinusoid_encoding_table(max_position, embedding_size), freeze=True)
#         self.register_buffer('position_embeddings', position_embeddings)

#     def forward(self, position_ids):
#         return self.position_embeddings(position_ids)


class SinusoidalPositionEmbedding(nn.Module):
    """定义Sin-Cos位置Embedding
    """
    def __init__(
        self,
        output_size,
        merge_mode='add',
        custom_position_ids=False
    ):
        super(SinusoidalPositionEmbedding, self).__init__()
        self.output_size = output_size
        self.merge_mode = merge_mode
        self.custom_position_ids = custom_position_ids

    def forward(self, inputs):
        input_shape = inputs.shape
        _, seq_len = input_shape[0], input_shape[1]
        position_ids = torch.arange(seq_len).type(torch.float)[None]
        indices = torch.arange(self.output_size // 2).type(torch.float)
        indices = torch.pow(10000.0, -2 * indices / self.output_size)
        embeddings = torch.einsum('bn,d->bnd', position_ids, indices)
        embeddings = torch.stack([torch.sin(embeddings), torch.cos(embeddings)], dim=-1)
        embeddings = torch.reshape(embeddings, (-1, seq_len, self.output_size))

        if self.merge_mode == 'add':
            return inputs + embeddings.to(inputs.device)
        elif self.merge_mode == 'mul':
            return inputs * (embeddings + 1.0).to(inputs.device)
        elif self.merge_mode == 'zero':
            return embeddings.to(inputs.device)
        
        
class RoPEPositionEncoding(nn.Module):
    """旋转式位置编码: https://kexue.fm/archives/8265
    """
    def __init__(self, max_position, embedding_size):
        super(RoPEPositionEncoding, self).__init__()
        position_embeddings = get_sinusoid_encoding_table(max_position, embedding_size)  # [seq_len, hdsz]
        cos_position = position_embeddings[:, 1::2].repeat_interleave(2, dim=-1)
        sin_position = position_embeddings[:, ::2].repeat_interleave(2, dim=-1)
        # register_buffer是为了最外层model.to(device)，不用内部指定device
        self.register_buffer('cos_position', cos_position)
        self.register_buffer('sin_position', sin_position)
    
    def forward(self, qw, seq_dim=-2):
        # 默认最后两个维度为[seq_len, hdsz]
        seq_len = qw.shape[seq_dim]
        qw2 = torch.stack([-qw[..., 1::2], qw[..., ::2]], dim=-1).reshape_as(qw)
        return qw * self.cos_position[:seq_len] + qw2 * self.sin_position[:seq_len]