import torch.nn as nn
from typing import Dict


def grad_norm(model: nn.Module, norm_type: int = 2) -> Dict[str, float]:
    """Compute gradient norm for each parameter in the model."""
    norms = {}
    for name, param in model.named_parameters():
        norms[name] = 0.0
        if param.grad:
            norm = param.grad.data.norm(norm_type).item()
            norms[name] = norm
    return norms


def grad_std(model: nn.Module) -> Dict[str, float]:
    """Compute standard deviation of gradients for each parameter."""
    stds = {}
    for name, param in model.named_parameters():
        stds[name] = 0.0
        if param.grad:
            std = param.grad.data.std().item()
            stds[name] = std
    return stds


def grad_max(model: nn.Module) -> Dict[str, float]:
    """Find the maximum absolute gradient value for each parameter."""
    max_vals = {}
    for name, param in model.named_parameters():
        max_vals[name] = -float("inf")
        if param.grad:
            max_val = param.grad.data.max().item()
            max_vals[name] = max_val

    return max_vals


def grad_min(model: nn.Module) -> Dict[str, float]:
    """Find the minimum absolute gradient value for each parameter."""
    min_vals = {}
    for name, param in model.named_parameters():
        min_vals[name] = float("inf")
        if param.grad:
            min_val = param.grad.data.min().item()
            min_vals[name] = min_val

    return min_vals


def grad_mean(model: nn.Module) -> Dict[str, float]:
    """Compute mean of gradients for each parameter."""
    means = {}
    for name, param in model.named_parameters():
        means[name] = 0.0
        if param.grad:
            mean = param.grad.data.mean().item()
            means[name] = mean

    return means


def grad_nan_num(model: nn.Module) -> Dict[str, int]:
    """Count the number of NaNs in gradients for each parameter."""
    nan_nums = {}
    for name, param in model.named_parameters():
        nan_nums[name] = 0
        if param.grad:
            nan_num = param.grad.isnan().sum().item()
            nan_nums[name] = nan_num
    return nan_nums


def ctx_get_loss(ctx):
    return ctx.loss


def ctx_get_lr(ctx):
    return ctx.optimizer.param_groups[-1]["lr"]


def ctx_get_grad_norm(ctx):
    return grad_norm(ctx.model)


def ctx_get_grad_std(ctx):
    return grad_std(ctx.model)


def ctx_get_grad_max(ctx):
    return grad_max(ctx.model)


def ctx_get_grad_min(ctx):
    return grad_min(ctx.model)


def ctx_get_grad_mean(ctx):
    return grad_mean(ctx.model)


def ctx_get_grad_nan_num(ctx):
    return grad_nan_num(ctx.model)