import typing
import numbers
from typing import List, Optional, Tuple, Union, cast

import numpy as np

# import torch
from numpy.typing import ArrayLike, NDArray

from zyplib.utils.tensor_like import TensorLike, get_tensorlike_backend


def _as_int(value: int | float | complex) -> int:
    if isinstance(value, complex):
        if value.imag != 0:
            raise ValueError(f'y must contain integer-valued labels, got {value!r}')
        value = value.real

    real_value = float(cast(numbers.Real, value))
    if not real_value.is_integer():
        raise ValueError(f'y must contain integer-valued labels, got {value!r}')
    return int(real_value)


def to_onehot(
    y: Union[TensorLike, typing.Sequence[int]],
    squeeze_2_class: bool = True,
    num_classes: Optional[int] = None,
) -> TensorLike:
    """将标签编码为独热码 (已使用 tensor_like 改造)

    Parameters
    ----------
    - `y` : `np.ndarray[int] | torch.Tensor[int] | list[int]`
        - 标签, 一个一维的整数数组或张量
    - `squeeze_2_class` : `bool`, optional
        - 如果为 True 且标签为二分类, 则返回一个 [N, 1] 维度的独热码, by default True
    - `num_classes` : `int`, optional
        - 类别数量, 如果为 None, 则自动从数据中推断, by default None

    Returns
    ----------
    - `np.ndarray` or `torch.Tensor`
        - 独热码，返回类型与输入类型一致

    Raises
    ----------
    - `TypeError`
        - 如果输入类型不是 list, NumPy 数组或 PyTorch Tensor
    """
    if isinstance(y, list):
        return to_onehot(np.array(y), squeeze_2_class, num_classes)

    tensor_y = cast(TensorLike, y)
    B = get_tensorlike_backend(tensor_y)

    if B.numel(tensor_y) == 0:
        n_classes = num_classes if num_classes is not None else 0
        # 对于 one-hot 编码，即使是空输入，也应该返回二维数组
        # 如果 squeeze_2_class=True 且 n_classes=2，则返回 (0, 1)
        if squeeze_2_class and n_classes == 2:
            return B.zeros((0, 1), dtype=B.dtype(tensor_y))
        return B.zeros((0, n_classes), dtype=B.dtype(tensor_y))

    indices = B.astype(tensor_y, B.type('int64'))
    if not B.allclose(tensor_y, B.astype(indices, B.dtype(tensor_y)), rtol=0.0, atol=0.0):
        raise ValueError('y must contain integer-valued labels')

    # 1. 确定最终的类别数 n_classes
    if num_classes is None:
        # 从数据推断
        n_classes = _as_int(B.item(B.max(indices))) + 1
        # 确定实际存在的类别数
        actual_classes = B.shape(cast(TensorLike, B.unique(indices)))[0]

        # 2. 基于推断出的 n_classes 和 actual_classes 判断是否为特殊情况
        # 如果推断出的总类别数<=1，或实际存在的类别数是2且需要压缩
        if actual_classes <= 1 or (n_classes == 2 and squeeze_2_class):
            return B.reshape(tensor_y, (-1, 1))
    else:
        # 用户已指定 num_classes
        n_classes = num_classes
        # 3. 如果用户指定了 num_classes，只有当 num_classes=2 且需要压缩时，才返回 (N, 1)
        # 注意：这里不再关心 actual_classes，因为用户的指定优先
        if n_classes == 2 and squeeze_2_class:
            return B.reshape(tensor_y, (-1, 1))

    # 对于所有其他情况（多分类，或不压缩的二分类），执行标准的 one-hot 编码
    one_hot = B.eye(n_classes, dtype=B.dtype(tensor_y), device=B.device(tensor_y))[
        indices
    ]
    return one_hot


def from_onehot(y: TensorLike) -> TensorLike:
    """将独热码转换回标签 (已使用 tensor_like 改造)

    Parameters
    ----------
    - `y` : `np.ndarray` or `torch.Tensor`
        - 独热码, 一个一维或二维的浮点数/整数数组或张量

    Returns
    ----------
    - `np.ndarray` or `torch.Tensor`
        - 标签, 一个一维的整数数组或张量

    Raises
    ----------
    - `TypeError`
        - 输入必须是 NumPy 数组或 PyTorch Tensor
    """
    # if not isinstance(y, (np.ndarray, torch.Tensor)):
    #     raise TypeError('Input must be a NumPy array or PyTorch Tensor')

    B = get_tensorlike_backend(y)

    ndim = B.ndim(y)

    # 如果输入本身就是一维的（已经是标签格式），则确保其为整数类型并返回
    if ndim == 1:
        return B.astype(y, B.type('int64'))

    ## 如果是二维的
    if ndim == 2:
        # 如果是 (N, 1) 的形状，说明是压缩的标签，直接展平即可
        if B.shape(y)[1] == 1:
            return B.astype(B.squeeze(y, axis=1), B.type('int64'))

        # 否则，认为是标准 one-hot，沿类别轴找最大值索引
        return B.argmax(y, axis=1)

    # 对于更高维度的输入，可以抛出错误或保持 argmax 行为
    # 这里我们假设输入最多是二维的，与 to_onehot 的输出对应
    raise ValueError(f'from_onehot expects 1D or 2D input, but got {ndim}D')


def percentile_digitalize(
    x: Union[TensorLike, typing.Sequence],
    n_bins: int = 4,
    axis: int = -1,  # axis 不再是 Optional，默认为 -1
) -> TensorLike:
    """将连续数据按照分位点数字化 (已使用 tensor_like 改造)

    此函数使用分位数将连续数据转换为离散的箱柜（bins），
    默认沿最后一个维度独立进行数字化。

    Parameters
    ----------
    x : Union[np.ndarray, torch.Tensor, list]
        输入数据。
    n_bins : int, optional
        分位点的个数（箱柜的数量），默认为4。
    axis : int, optional
        计算分位点的轴。默认为 -1 (最后一个维度)。不支持 None。

    Returns
    -------
    Union[np.ndarray, torch.Tensor]
        数字化后的数据，类型与输入一致，取值范围为 [0, n_bins-1]。

    Raises
    ------
    ValueError
        如果 axis 被设置为 None。
    """
    if axis is None:
        raise ValueError('`axis=None` is not supported. Please specify an integer axis.')

    if isinstance(x, list):
        if not x:
            return np.array([], dtype=np.int32)
        return percentile_digitalize(np.array(x), n_bins=n_bins, axis=axis)

    x = cast(TensorLike, x)
    B = get_tensorlike_backend(x)
    x_tensor = B.as_tensorlike(x)

    if B.numel(x_tensor) == 0:
        return B.astype(x_tensor, B.type('int32'))

    # 使用 'linear' 插值法
    interpolation_method = 'linear'

    # 1. 计算轴向分位点，并保持维度以便广播
    quantiles = B.linspace(
        1 / n_bins, 1 - 1 / n_bins, n_bins - 1, device=B.device(x_tensor)
    )
    q_values = B.quantile(
        x_tensor, quantiles, axis=axis, keepdims=True, interpolation=interpolation_method
    )

    # 2. 使用广播进行比较和数字化
    # 扩展 x_tensor，在最前面增加一个维度以进行广播
    x_expanded = B.unsqueeze(x_tensor, 0)  # shape: (1, ...)

    # 3. 比较并求和
    # x_expanded:      (1,         batch, features, ...)
    # q_values:        (n_bins-1,  batch, 1,        ...) (如果 axis=1)
    # 广播后的 comparison 张量，其第一个维度代表与哪个分位点的比较结果
    comparison = x_expanded >= q_values

    # 沿着分位数维度(axis=0)求和，得到每个元素大于等于多少个分界点
    summed = B.sum(B.astype(comparison, B.type('int32')), axis=0)

    # 确保最终返回的是 int32
    return B.astype(summed, B.type('int32'))


def continuous_smoothing(
    one_hot: NDArray, alpha: float = 0.2, decay: float = 0.25
) -> NDArray:
    """对独热编码进行连续平滑处理

    Parameters
    ----------
    one_hot : NDArray
        独热编码数据，[0, 0, 1, 0] 这种
    alpha : float, optional
        平滑系数，默认为0.2
    decay : float, optional
        衰减系数，默认为0.25

    Returns
    -------
    NDArray
        平滑后的标签
    """
    classes = one_hot.shape[-1]
    max_idx = np.argmax(one_hot, axis=-1)[..., np.newaxis]
    max_idx = np.repeat(max_idx, classes, axis=-1)
    one_hot_idx = one_hot.copy()
    one_hot_idx[..., :] = np.arange(classes)
    delta = np.abs(one_hot_idx - max_idx)
    scale = np.power(decay, delta).sum(axis=-1, keepdims=True) - 1
    running_alpha = alpha / scale
    label = np.zeros(one_hot.shape, dtype=np.float32)
    label = one_hot * (1 - alpha) + (1 - one_hot) * running_alpha * np.power(decay, delta)
    return label
