import inspect
import os
import random
from functools import lru_cache
from typing import Any, Optional, Union, cast

import numpy as np
import torch
from simple_inject import inject
from torch import nn


from zyplib.utils.print import print_debug, print_error


def seed_everything(seed: int):
    """设置所有随机数生成器的种子

    设置 Python、NumPy、PyTorch 的随机数种子，确保实验可重复性。

    Parameters
    ----------
    - `seed` : `int`
        - 随机数种子
    """
    os.environ['PL_GLOBAL_SEED'] = str(seed)
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def step_lr_sched(
    scheduler: torch.optim.lr_scheduler.LRScheduler,
    metrics: Optional[Union[float, torch.Tensor]] = None,
    epoch: Optional[int] = None,
):
    """更新学习率调度器

    根据调度器类型选择合适的更新方式。

    Parameters
    ----------
    - `scheduler` : `torch.optim.lr_scheduler.LRScheduler`
        - 学习率调度器
    - `metrics` : `Union[float, torch.Tensor]`, optional
        - 用于 ReduceLROnPlateau 的指标值, by default None
    - `epoch` : `int`, optional
    """
    if isinstance(scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
        scalar_metric: float = (
            float(metrics.item())
            if isinstance(metrics, torch.Tensor)
            else cast(float, metrics)
        )
        scheduler.step(metrics=cast(float, scalar_metric), epoch=epoch)
    elif epoch is not None:
        scheduler.step(epoch=epoch)
    else:
        scheduler.step()


@lru_cache
def cuda_available():
    """检查 CUDA 是否可用

    Returns
    ----------
    - `bool`
        - CUDA 是否可用
    """
    return torch.cuda.is_available()


@lru_cache
def default_device():
    """获取默认设备

    Returns
    ----------
    - `torch.device`
        - 如果 CUDA 可用返回 'cuda'，否则返回 'cpu'
    """
    return torch.device('cuda' if cuda_available() else 'cpu')


def use_device(
    x: Union[torch.Tensor, nn.Module],
    device: Optional[Union[str, torch.device]] = None,
    hint: bool = False,
):
    """将模型或张量移动到指定设备

    依赖注入：如果 device 为 None，则首先会 inject('device') 来尝试获取设备名称，
    如果仍然为 None，则使用默认设备。

    Parameters
    ----------
    - `x` : `Union[torch.Tensor, nn.Module]`
        - 需要移动的模型或张量
    - `device` : `Union[str, torch.device]`, optional
        - 目标设备名称或设备对象, by default None
    - `hint` : `bool`, optional
        - 是否打印移动信息, by default False

    Returns
    ----------
    - `Union[torch.Tensor, nn.Module]`
        - 移动后的模型或张量
    """
    if not isinstance(x, (torch.Tensor, nn.Module)):
        return x
    elif not hasattr(x, 'to'):
        return x

    if device is None:
        device = inject('device', if_not_found='none')
    if device is None:
        device = default_device()

    if hint:
        print_debug(f'use_device: {x.__class__.__name__} --> {device}')
    return x.to(device)


def load_metrics_from_tensorboard_file(fpath: str) -> dict[str, Any]:
    """从 Tensorboard 的 event 文件中加载所有的指标

    仅限常见的标量指标

    Parameters
    ----------
    - `fpath` : `str`
        - Tensorboard event 文件路径

    Returns
    ----------
    - `dict`
        - 包含指标名称和对应值列表的字典，格式为 {'metric_name': {'step': [list], 'value': [list]}}
    """
    from zyplib.utils.kv import DottableDict

    # 创建事件累积器
    try:
        from tensorboard.backend.event_processing import event_accumulator  # pyright: ignore[reportMissingImports]
    except Exception as e:
        print_error(f'Failed to import event_accumulator from tensorboard: {e}')
        raise e
    ea = event_accumulator.EventAccumulator(fpath)
    ea.Reload()

    # 获取所有标量指标
    metrics: dict[str, Any] = {}
    scalar_keys = ea.Tags()['scalars']

    for key in scalar_keys:
        # 获取标量事件
        scalar_events = ea.Scalars(key)
        # 提取 (step, value) 对
        # metrics[key] = [(event.step, event.value) for event in scalar_events]
        metrics[key] = DottableDict({'step': [], 'value': []})
        for event in scalar_events:
            metrics[key]['step'].append(event.step)
            metrics[key]['value'].append(event.value)

    return metrics


def tensorboard_eventfile_to_json(fpath: str) -> dict[str, Any]:
    """将 Tensorboard 的 event 文件转换为 JSON 格式

    Parameters
    ----------
    - `fpath` : `str`
        - Tensorboard event 文件路径

    Returns
    ----------
    - `dict`
        - 包含指标名称和对应值列表的字典，格式为 {'metric_name': {'step': [list], 'value': [list]}}
    """
    metric = load_metrics_from_tensorboard_file(fpath)
    import json
    from pathlib import Path

    json_path = Path(fpath).with_suffix('.json')
    with open(json_path, 'w', encoding='utf-8') as f:
        json.dump(metric, f, ensure_ascii=False)
    return cast(dict[str, Any], metric)
