import typing
from collections import namedtuple
from dataclasses import dataclass
import math
from typing import Optional, TypeAlias

import torch
from torch import Tensor

__all__ = [
    'mixup',
    'remix',
    'mamix',
    'mixup_criterion',
    'remix_criterion',
    'mamix_criterion',
]

# MixupResult = namedtuple('MixupResult', ['mixed_x', 'mixed_y', 'y_a', 'y_b', 'lam'])
# RemixResult = namedtuple('RemixResult', ['mixed_x', 'y_a', 'y_b', 'lam_y'])
OnehotLabel: TypeAlias = Tensor  # shape: (N, C), dtype: torch.float32
IndexLabel: TypeAlias = Tensor  # shape: (N,), dtype: torch.long


@dataclass
class MixupResult:
    mixed_x: Tensor
    mixed_y: OnehotLabel
    y_a: IndexLabel
    y_b: IndexLabel
    lam: float


@dataclass
class RemixResult:
    mixed_x: Tensor
    y_a: IndexLabel
    y_b: IndexLabel
    lam_y: Tensor


def mixup(
    x: Tensor, y: Tensor, alpha: float = 1.0, random_state: Optional[int] = None
) -> MixupResult:
    """标准 Mixup 数据增强

    对两个随机样本及其标签进行线性插值来生成新的训练样本。
    要求 y 为 one-hot 编码。

    Parameters
    ----------
    - `x` : `Tensor`
        - 输入数据，形状 (N, ...)
    - `y` : `Tensor`
        - 标签，形状 (N, C)，one-hot 编码
    - `alpha` : `float`, optional
        - Beta 分布参数，默认 1.0
    - `random_state` : `Optional[int]`, optional
        - 随机种子

    Returns
    -------
    - `MixupResult`
        - mixed_x: 混合后的数据
        - mixed_y: 混合后的标签（float, one-hot）
        - y_a, y_b: 原始标签对（整数索引）
        - lam: 混合系数

    References
    ----------
    - Paper: mixup: Beyond Empirical Risk Minimization
    - Link: https://arxiv.org/pdf/1710.09412.pdf

    Examples
    --------
    ```python
    # y_one_hot = F.one_hot(y, num_classes=10)
    # 使用混合标签（适用于 MSE/BCE 等损失）
    result = mixup(x, y_one_hot, alpha=1.0)
    output = model(result.mixed_x)
    loss = criterion(output, result.mixed_y)

    # 使用原始标签对（适用于 CrossEntropy）；更正统
    result = mixup(x, y_one_hot, alpha=1.0)
    output = model(result.mixed_x)
    loss = mixup_criterion(criterion, output, result.y_a, result.y_b, result.lam)
    ```
    """
    if alpha <= 0:
        raise ValueError('alpha 参数必须大于 0')
    if y.ndim != 2:
        raise ValueError('y 必须是 one-hot 编码, 形状为 (N, C)')

    device = x.device
    batch_size = x.size(0)

    generator = torch.Generator(device=device)
    if random_state is not None:
        generator.manual_seed(random_state)

    # 采样 lambda; 这个无法控制随机，就不管了
    lam = torch.distributions.Beta(alpha, alpha).sample().item()  # 需要调整 Beta 采样
    index = torch.randperm(batch_size, generator=generator, device=device)  # 随机配对

    # 混合数据
    mixed_x = lam * x + (1 - lam) * x[index]
    y_a_onehot, y_b_onehot = y, y[index]

    # 混合标签（one-hot）
    mixed_y = lam * y_a_onehot.float() + (1 - lam) * y_b_onehot.float()

    # 转为 index label 以便 mixup_criterion 使用
    y_a = torch.argmax(y_a_onehot, dim=1)
    y_b = torch.argmax(y_b_onehot, dim=1)

    return MixupResult(mixed_x=mixed_x, mixed_y=mixed_y, y_a=y_a, y_b=y_b, lam=lam)


def mixup_criterion(
    criterion: torch.nn.Module, pred: Tensor, y_a: Tensor, y_b: Tensor, lam: float
) -> Tensor:
    """Mixup 损失函数

    Parameters
    ----------
    - `criterion` : `nn.Module`
        - 基础损失函数（如 CrossEntropyLoss）
    - `pred` : `Tensor`
        - 模型预测，形状 (N, C)
    - `y_a`, `y_b` : `Tensor`
        - 标签对，形状 (N,)
    - `lam` : `float`
        - 混合系数

    Returns
    -------
    - `Tensor`
        - 混合损失值
    """
    return lam * criterion(pred, y_a) + (1 - lam) * criterion(pred, y_b)


def remix(
    x: Tensor,
    y: Tensor,
    cls_num_list: typing.Sequence[int],
    k: int = 3,
    tau: float = 0.5,
    alpha: float = 1.0,
    random_state: Optional[int] = None,
) -> RemixResult:
    """Remix: 针对不均衡数据的 Mixup 改进算法

    根据类别频率动态调整标签混合比例，保护少数类样本。
    y 可以是整数索引或 one-hot 编码。

    Parameters
    ----------
    - `x` : `Tensor`
        - 输入数据，形状 (N, ...)
    - `y` : `Tensor`
        - 标签，形状 (N,) 整数类型, 或 (N, C) one-hot 编码
    - `cls_num_list` : `list[int]`
        - 每个类别的样本数，长度为 num_classes
    - `k` : `int`, optional
        - 多数类阈值，默认 3
    - `tau` : `float`, optional
        - lambda 阈值，默认 0.5
    - `alpha` : `float`, optional
        - Beta 分布参数，默认 1.0
    - `random_state` : `Optional[int]`, optional
        - 随机种子

    Returns
    -------
    - `RemixResult`
        - mixed_x: 混合后的数据
        - y_a, y_b: 原始标签对
        - lam_y: 标签混合比例（向量）

    References
    ----------
    - Paper: Remix: Rebalanced Mixup (ECCV 2020)
    - Link: https://arxiv.org/pdf/2007.03943.pdf

    Examples
    --------
    ```python
    # 获取类别频率
    cls_num_list = [5000, 2997, 1796, 1077, 645, 387, 232, 139, 83, 50]

    # 应用 Remix
    result = remix(x, y, cls_num_list, k=3, tau=0.5)

    # 计算损失
    output = model(result.mixed_x)
    loss = remix_criterion(criterion, output, result.y_a, result.y_b, result.lam_y)
    ```
    """
    if alpha <= 0:
        raise ValueError('alpha 参数必须大于 0')

    # 如果 y 是 one-hot，转为 index
    if y.ndim == 2:
        y = torch.argmax(y, dim=1)

    device = x.device
    batch_size = x.size(0)

    generator = torch.Generator(device=device)
    if random_state is not None:
        generator.manual_seed(random_state)

    # 采样 lambda_x
    lam_x = torch.distributions.Beta(alpha, alpha).sample().item()

    # 随机配对
    index = torch.randperm(batch_size, generator=generator, device=device)  # 随机配对

    # **向量化计算 lambda_y**（关键优化）
    cls_num_tensor = torch.tensor(cls_num_list, device=device, dtype=torch.float32)

    # 获取每对样本的类别频率
    n_i = cls_num_tensor[y]  # shape: (N,)
    n_j = cls_num_tensor[y[index]]  # shape: (N,)

    # 计算频率比
    ratio = n_i / (n_j + 1e-8)  # 避免除零

    # 向量化 Remix 规则
    lam_y = torch.full((batch_size,), lam_x, device=device, dtype=x.dtype)

    # 条件 1: ratio >= k 且 lam_x < tau -> lam_y = 0（保留少数类 j）
    mask1 = (ratio >= k) & (lam_x < tau)
    lam_y[mask1] = 0.0

    # 条件 2: ratio <= 1/k 且 (1-lam_x) < tau -> lam_y = 1（保留少数类 i）
    mask2 = (ratio <= 1.0 / k) & ((1.0 - lam_x) < tau)
    lam_y[mask2] = 1.0

    # 数据混合
    mixed_x = lam_x * x + (1.0 - lam_x) * x[index]
    y_a, y_b = y, y[index]

    return RemixResult(mixed_x=mixed_x, y_a=y_a, y_b=y_b, lam_y=lam_y)


def remix_criterion(
    criterion: torch.nn.Module, pred: Tensor, y_a: Tensor, y_b: Tensor, lam_y: Tensor
) -> Tensor:
    """Remix 损失函数

    Parameters
    ----------
    - `criterion` : `nn.Module`
        - 基础损失函数（如 CrossEntropyLoss，需要 reduction='none'）
    - `pred` : `Tensor`
        - 模型预测，形状 (N, C)
    - `y_a`, `y_b` : `Tensor`
        - 标签对，形状 (N,)
    - `lam_y` : `Tensor`
        - 标签混合比例，形状 (N,)

    Returns
    -------
    - `Tensor`
        - 平均损失值
    """
    loss_a = criterion(pred, y_a)
    loss_b = criterion(pred, y_b)

    if loss_a.dim() == 0:  # 如果 criterion 没有 reduction='none'
        raise RuntimeError('criterion 需要设置 reduction="none"')

    return (lam_y * loss_a + (1 - lam_y) * loss_b).mean()

    # 逐样本加权
    # if loss_a.dim() == 0:  # 如果 criterion 没有 reduction='none'
    #     loss = lam_y.mean() * loss_a + (1 - lam_y).mean() * loss_b
    # else:  # 正常情况
    #     loss = lam_y * loss_a + (1 - lam_y) * loss_b
    #     loss = loss.mean()

    # return loss


# def _get_k(n1: float, n2: float, f: float) -> float:
#     """计算 MAMix 的平滑系数"""
#     if n2 <= 0 or n1 <= 0:
#         return 0.0
#     import math

#     log_ratio = math.log(n1 / n2)
#     return f * log_ratio


# def _get_lambda(lam_x: float, k1: float, k2: float) -> float:
#     """使用 sigmoid 计算平滑的 lambda"""
#     import math
#     # Question: 这里应该是 k1+k2 OR k1-k2?
#     return 1.0 / (1.0 + math.exp(-(k1 - k2) * (lam_x - 0.5)))


def mamix(
    x: Tensor,
    y: Tensor,
    cls_num_list: typing.Sequence[int],
    f: float = 1.0,
    alpha: float = 1.0,
    random_state: Optional[int] = None,
) -> RemixResult:
    """MAMix: Majority-Aware Mixup

    使用平滑函数调整标签混合比例，比 Remix 更优雅。
    y 可以是整数索引或 one-hot 编码。

    Parameters
    ----------
    - `x` : `Tensor`
        - 输入数据，形状 (N, ...)
    - `y` : `Tensor`
        - 标签，形状 (N,) 整数类型, 或 (N, C) one-hot 编码
    - `cls_num_list` : `list[int]`
        - 每个类别的样本数
    - `f` : `float`, optional
        - 比例因子，默认 1.0
    - `alpha` : `float`, optional
        - Beta 分布参数，默认 1.0
    - `random_state` : `Optional[int]`, optional
        - 随机种子

    Returns
    -------
    - `RemixResult`
        - mixed_x, y_a, y_b, lam_y

    References
    ----------
    - Paper: From SMOTE to Mixup for Deep Imbalanced Classification (TAAI 2023)
    - Link: https://www.csie.ntu.edu.tw/~htlin/paper/doc/taai23mamix.pdf

    Examples
    --------
    ```python
    # MAMix 通常有最佳性能
    result = mamix(x, y, cls_num_list, f=1.0)

    output = model(result.mixed_x)
    loss = mamix_criterion(criterion, output, result.y_a, result.y_b, result.lam_y)
    ```
    """
    if alpha <= 0:
        raise ValueError('alpha 参数必须大于 0')

    # 如果 y 是 one-hot，转为 index
    if y.ndim == 2:
        y = torch.argmax(y, dim=1)

    generator = torch.Generator(device=x.device)
    if random_state is not None:
        generator.manual_seed(random_state)

    device = x.device
    batch_size = x.size(0)

    # 采样 lambda_x
    lam_x = torch.distributions.Beta(alpha, alpha).sample().item()

    # 随机配对
    index = torch.randperm(batch_size, generator=generator, device=device)  # 随机配对

    # 计算平滑 lambda_y
    cls_num_tensor = torch.tensor(cls_num_list, device=device, dtype=torch.float32)
    # 获取每对样本的类别频率
    n_i = cls_num_tensor[y]  # shape: (N,)
    n_j = cls_num_tensor[y[index]]  # shape: (N,)

    # 计算 k 和 lambda_y（向量化）
    # k = f * log(n_i / n_j)
    # lam_y = 1 / (1 + exp(-k * (lam_x - 0.5)))
    k = f * torch.log(n_i / (n_j + 1e-8))  # 避免除零和 log(0)
    lam_y = 1.0 / (1.0 + torch.exp(-k * (lam_x - 0.5)))

    # 数据混合
    mixed_x = lam_x * x + (1.0 - lam_x) * x[index]
    y_a, y_b = y, y[index]

    return RemixResult(mixed_x=mixed_x, y_a=y_a, y_b=y_b, lam_y=lam_y)


def mamix_criterion(
    criterion: torch.nn.Module, pred: Tensor, y_a: Tensor, y_b: Tensor, lam_y: Tensor
) -> Tensor:
    """MAMix 损失函数

    与 remix_criterion 相同的实现

    Parameters
    ----------
    - `criterion` : `nn.Module`
        - 基础损失函数（需要 reduction='none'）
    - `pred` : `Tensor`
        - 模型预测，形状 (N, C)
    - `y_a`, `y_b` : `Tensor`
        - 标签对，形状 (N,)
    - `lam_y` : `Tensor`
        - 标签混合比例，形状 (N,)

    Returns
    -------
    - `Tensor`
        - 平均损失值
    """
    return remix_criterion(criterion, pred, y_a, y_b, lam_y)
