from collections import OrderedDict
from typing import Dict, Union

import numpy as np
import torch
from torch import nn


def module_grad_norm(module: nn.Module, p_norm: int = 2) -> float:
    """计算 module 当前所有参数的梯度的范数

    moduleGradNorm 的计算结果为 module 中所有子模块的 parameter 组成的向量的 norm

    Parameters
    ----------
    - `module` : `torch.nn.Module`
        - 要计算梯度范数的模块
    - `p_norm` : `int`, optional
        - 范数阶数, by default 2

    Returns
    ----------
    - `float`
        - 所有参数梯度的总范数
    """
    sum_grad_norm = 0.0
    for p in module.parameters():
        if p.grad is None:
            continue
        grad = p.grad.detach().norm(p_norm).item()
        sum_grad_norm += grad**p_norm
    total_grad_norm = sum_grad_norm ** (1 / p_norm)
    return total_grad_norm


def param_grad_norm(module: nn.Module, p_norm: int = 2) -> Dict[str, float]:
    """分别计算 module 内各个 parameter 梯度的范数

    以字典的格式返回每个参数的梯度范数。

    Parameters
    ----------
    - `module` : `torch.nn.Module`
        - 要计算梯度范数的模块
    - `p_norm` : `int`, optional
        - 范数阶数, by default 2

    Returns
    ----------
    - `Dict[str, float]`
        - 参数名到其梯度范数的映射字典，格式为 `{parameter_name: norm}`
    """
    norms = {}
    for name, p in module.named_parameters():
        if p.grad is not None:
            grad = p.grad.detach().norm(p_norm).item()
            norms[name] = grad
    return norms


def grad_norm(
    module: nn.Module, p_norm: int = 2, type: str = 'module'
) -> Union[float, dict[str, float]]:
    """计算 module 梯度的范数

    根据 type 的不同，采用不同的方式计算梯度范数。本质上是调用其他两个函数。

    Parameters
    ----------
    - `module` : `torch.nn.Module`
        - 要计算梯度范数的模块
    - `p_norm` : `int`, optional
        - 范数阶数, by default 2
    - `type` : `str`, optional
        - 返回类型, by default 'module'
        - 'module': 计算整个 module 梯度的范数，调用 `moduleGradNorm`，返回 float
        - 'param': 计算 module 各个 parameter 的范数，调用 `paramGradNorm`，返回 dict

    Returns
    ----------
    - `Union[float, dict]`
        - 当 type='module' 时返回 float，表示总范数
        - 当 type='param' 时返回 dict，表示各参数范数

    Raises
    ----------
    - `ValueError`
        - 当 type 参数不是 'module' 或 'param' 时抛出
    """
    if type == 'module':
        return module_grad_norm(module, p_norm)
    elif type == 'param':
        return param_grad_norm(module, p_norm)
    else:
        raise ValueError(f'Invalid argument type: {type}')
