"""
环境层通用工具
=============
提供数组拼接、索引采样、tensor→numpy 转换等辅助功能。
"""

from typing import Dict, Iterable, List, Optional, Union

import numpy as np
from numpy import ndarray
from torch import Tensor
from torch_geometric.data.data import BaseData


def save_concatenate(
    arrays: Iterable[np.ndarray], *args, **kwargs
) -> Optional[np.ndarray]:
    """
    安全拼接多个数组。自动过滤 None 值，空列表返回 None。

    Args:
        arrays: 要拼接的数组列表（可能包含 None）

    Returns:
        拼接后的数组；若全为 None 则返回 None

    Example:
        >>> result = save_concatenate([arr1, None, arr2], axis=1)
    """
    arrays = [array for array in arrays if array is not None]
    if len(arrays) == 0:
        return None
    return np.concatenate(arrays, *args, **kwargs)


class IndexSampler:
    """
    随机索引采样器 — 用于循环缓冲区中随机抽取 PDE 实例。

    内部维护一个随机排列的索引数组，每次调用 next() 返回一个索引。
    遍历完所有索引后自动重新洗牌。

    Example:
        >>> sampler = IndexSampler(100, np.random.RandomState(42))
        >>> idx = sampler.next()  # 随机抽取一个索引
    """

    def __init__(self, size: int, random_state: np.random.RandomState):
        self._size = size
        self._indices = np.arange(size)
        self._random_state = random_state
        self._reset()

    def next(self) -> int:
        """返回下一个随机索引，到底后自动洗牌重排。"""
        if self._position == self._size:
            self._reset()
        index = self._indices[self._position]
        self._position += 1
        return index

    def _reset(self):
        self._position = 0
        self._random_state.shuffle(self._indices)

    def __len__(self):
        return self._size


def detach(
    tensor: Union[Tensor, Dict[str, Tensor], List[Tensor]],
) -> Union[ndarray, Dict[str, ndarray], List[ndarray], BaseData]:
    """
    将 PyTorch tensor 安全转换为 numpy 数组（自动处理 GPU→CPU）。

    Args:
        tensor: PyTorch tensor、tensor 字典或 tensor 列表

    Returns:
        对应的 numpy 数组

    Example:
        >>> action_np = detach(actions_tensor)  # → np.ndarray
    """
    if isinstance(tensor, dict):
        return {key: detach(value) for key, value in tensor.items()}
    elif isinstance(tensor, list):
        return [detach(value) for value in tensor]
    if tensor.is_cuda:
        return tensor.cpu().detach().numpy()
    else:
        return tensor.detach().numpy()