"""注意力机制工具函数

提供注意力机制中常用的工具函数，包括头操作、掩码创建、池化操作等。
这些函数被多个注意力模块共享使用，避免代码重复。
"""

import math
from typing import Optional, Tuple, Union

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


def split_heads(x: torch.Tensor, num_heads: int, head_dim: int) -> torch.Tensor:
    """将输入张量拆分为多头

    Parameters
    ----------
    - `x` : `torch.Tensor`
        - 输入张量，形状: (batch_size, seq_len, embed_dim)
    - `num_heads` : `int`
        - 注意力头的数量
    - `head_dim` : `int`
        - 每个头的维度

    Returns
    ----------
    - `torch.Tensor`
        - 拆分后的张量，形状: (batch_size, num_heads, seq_len, head_dim)

    Examples
    --------
    ```python
    >>> import torch
    >>> from zyplib.nn.modules.attention.attn_utils import split_heads
    >>>
    >>> x = torch.randn(32, 100, 128)  # [batch, seq_len, embed_dim]
    >>> heads = split_heads(x, num_heads=8, head_dim=16)
    >>> print(heads.shape)  # torch.Size([32, 8, 100, 16])
    ```
    """
    batch_size, seq_len, _ = x.shape
    return x.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)


def combine_heads(x: torch.Tensor, embed_dim: int) -> torch.Tensor:
    """合并多头注意力的输出

    Parameters
    ----------
    - `x` : `torch.Tensor`
        - 多头输出，形状: (batch_size, num_heads, seq_len, head_dim)
    - `embed_dim` : `int`
        - 嵌入维度 (num_heads * head_dim)

    Returns
    ----------
    - `torch.Tensor`
        - 合并后的张量，形状: (batch_size, seq_len, embed_dim)

    Examples
    --------
    ```python
    >>> import torch
    >>> from zyplib.nn.modules.attention.attn_utils import combine_heads
    >>>
    >>> x = torch.randn(32, 8, 100, 16)  # [batch, heads, seq_len, head_dim]
    >>> combined = combine_heads(x, embed_dim=128)
    >>> print(combined.shape)  # torch.Size([32, 100, 128])
    ```
    """
    batch_size, _, seq_len, _ = x.shape
    return x.transpose(1, 2).contiguous().view(batch_size, seq_len, embed_dim)


def create_causal_mask(seq_len: int, device: torch.device) -> torch.Tensor:
    """创建因果掩码（下三角掩码）

    Parameters
    ----------
    - `seq_len` : `int`
        - 序列长度
    - `device` : `torch.device`
        - 设备类型

    Returns
    ----------
    - `torch.Tensor`
        - 形状为 (seq_len, seq_len) 的布尔掩码

    Examples
    --------
    ```python
    >>> import torch
    >>> from zyplib.nn.modules.attention.attn_utils import create_causal_mask
    >>>
    >>> mask = create_causal_mask(5, torch.device('cpu'))
    >>> print(mask)
    tensor([[ True, False, False, False, False],
            [ True,  True, False, False, False],
            [ True,  True,  True, False, False],
            [ True,  True,  True,  True, False],
            [ True,  True,  True,  True,  True]])
    ```
    """
    return torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device))
