from __future__ import annotations

from typing import Any

import torch


def count_parameters(model: torch.nn.Module) -> int:
    return sum(parameter.numel() for parameter in model.parameters() if parameter.requires_grad)


def per_channel_mse(
    squared_error_sum: torch.Tensor,
    count: int,
    target_names: tuple[str, ...],
) -> dict[str, float]:
    if count <= 0:
        raise ValueError("Metric count must be positive")
    values = (squared_error_sum / count).detach().cpu().tolist()
    return {name: float(value) for name, value in zip(target_names, values, strict=True)}


def overall_mse(squared_error_sum: torch.Tensor, count: int, target_dim: int) -> float:
    if count <= 0 or target_dim <= 0:
        raise ValueError("Metric count and target dimension must be positive")
    return float((squared_error_sum.sum() / (count * target_dim)).detach().cpu().item())


def device_metrics(device: torch.device) -> dict[str, Any]:
    if device.type != "cuda":
        return {
            "device": str(device),
            "gpu_name": None,
            "gpu_memory_total_mb": None,
            "gpu_memory_peak_allocated_mb": None,
        }

    index = device.index if device.index is not None else torch.cuda.current_device()
    properties = torch.cuda.get_device_properties(index)
    return {
        "device": f"cuda:{index}",
        "gpu_name": torch.cuda.get_device_name(index),
        "gpu_memory_total_mb": int(properties.total_memory // (1024 * 1024)),
        "gpu_memory_peak_allocated_mb": int(torch.cuda.max_memory_allocated(index) // (1024 * 1024)),
    }