"""Utility functions for DEPToE models."""
from logging import DEBUG
from typing import Any

from composer import ComposerModel
from flwr.common import log


def log_num_params(model: ComposerModel, logged_cfg: dict[str, Any]) -> None:
    """Log the number of parameters in the model.

    Parameters
    ----------
    model : ComposerModel
        The model for which to log the number of parameters.
    logged_cfg : dict[str, Any]
        The configuration dictionary to update with the number of parameters.

    """
    # Log number of parameters
    n_params = sum(p.numel() for p in model.parameters())
    n_trainable_params = sum(
        p.numel() for p in model.parameters() if p.requires_grad
    )
    if hasattr(model, "n_active_params"):
        n_active_params = model.n_active_params
    else:
        n_active_params = n_params
    logged_cfg.update(
        {
            "n_params": n_params,
            "n_active_params": n_active_params,
            "n_trainable_params": n_trainable_params,
            "n_non_embedding_params": (
                model.n_non_embedding_params
                if hasattr(model, "n_non_embedding_params")
                else "N/A"
            ),
        },
    )
    log(
        DEBUG,
        "Number of model parameters: %s",
        {
            "n_params": logged_cfg["n_params"],
            "n_active_params": logged_cfg["n_active_params"],
            "n_trainable_params": logged_cfg["n_trainable_params"],
            "n_non_embedding_params": logged_cfg["n_non_embedding_params"],
        },
    )
