"""Utility functions for FL and experiment management.

They assure compatibility with the Flower and wandb APIs.
"""

import ast
import concurrent.futures
import copy
import threading
import time
import types
import warnings
from collections import OrderedDict
from collections.abc import Generator, Iterable
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager
from dataclasses import dataclass, field
from logging import DEBUG, ERROR
from queue import Queue
from typing import Any, cast

import numpy as np
import psutil
import torch
from composer import Timestamp
from composer.utils import dist
from flwr.common import (
    NDArray,
    NDArrays,
    log,
)
from llmfoundry.command_utils.train import (
    validate_config,
)
from llmfoundry.utils.builders import (
    build_composer_model,
    build_tokenizer,
)
from llmfoundry.utils.config_utils import (
    TRAIN_CONFIG_KEYS,
    TrainConfig,
    make_dataclass_and_log_config,
    process_init_device,
)
from omegaconf import DictConfig, OmegaConf
from torch import device as device_type
from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel

import ray
import wandb
from mt_dao.conf.base_schema import BaseConfig
from mt_dao.constants import CUSTOM_RAY_GC, CUSTOM_SHM_GC
from mt_dao.masks_utils import ModelStateNames, create_layers_names_and_types_from_names
from mt_dao.moe_utils.utils import (
    split_experts_layers_tensors,
)
from mt_dao.pytorch_utils.model_parameters_getters import get_tensors_dict
from mt_dao.pytorch_utils.utils import get_n_experts_from_model
from mt_dao.shm.utils import clean_shm_by_name
from ray._private.internal_api import (
    free as ray_free,  # noqa: PLC2701  # type: ignore[reportPrivateImportUsage]  # noqa: PLC2701  # type: ignore[reportPrivateImportUsage]
)


@dataclass
class ClientState:
    """Dataclass for client state."""

    local_steps_cumulative: int
    local_timestamp: dict[str, Any] = field(
        default_factory=lambda: {
            k: v
            for k, v in Timestamp().state_dict().items()
            if is_literal_for_ast(repr(v))
        },
    )
    steps_done: int = 0


class ParameterShapeMismatchError(Exception):
    """Exception raised when parameter shapes mismatch."""

    def __init__(
        self,
        remote_name: str,
        remote_shape: torch.Size,
        local_shape: torch.Size,
    ) -> None:
        """Initialize the exception."""
        msg = f"Shapes don't match for {remote_name}: {remote_shape} != {local_shape}"
        super().__init__(msg)


class NoOpContextManager:
    """A context manager that does nothing."""

    def __enter__(self) -> None:
        """Do nothing."""
        return

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        traceback: types.TracebackType | None,
    ) -> None:
        """Do nothing."""
        return


@contextmanager
def custom_ray_garbage_collector(
    garbage_queue: Queue[ray.ObjectRef | str],
    list_of_threads: list[threading.Thread],
    process_name: str,
    timeout: float = 10,
    *,
    join_at_the_end: bool = True,
) -> Generator[None, Any, None]:
    """Context manager for custom Ray garbage collection.

    This context manager creates a thread that continuously collects and frees
    Ray objects from a given queue. The thread runs in the background and stops
    when the context manager exits.

    Parameters
    ----------
    garbage_queue : Queue[ray.ObjectRef | str]
        A queue containing Ray object references or strings to be freed.
    list_of_threads : list[threading.Thread]
        A list of threads to append the garbage collection thread to.
    process_name : str
        The name of the process.
    timeout : float, optional
        The time in seconds to wait between garbage collection cycles, by default 10.
    join_at_the_end : bool, optional
        Whether to join the garbage collection thread at the end, by default
        True.

    Yields
    ------
    None
        The context manager yields control back to the caller.

    Notes
    -----
    The function assumes the existence of `ray_free` and `log` functions/utilities,
    as well as the `ERROR` constant for logging purposes.

    Example
    -------
    >>> garbage_queue = Queue()
    >>> with custom_ray_garbage_collector(garbage_queue):
    >>>     # Your code here
    >>>     pass

    """
    stop_event = threading.Event()

    def collect_ray_garbage() -> None:
        while not stop_event.is_set():
            time.sleep(timeout)
            while not garbage_queue.empty():
                try:
                    obj = garbage_queue.get(block=True, timeout=0.1)
                    asrt_msg = f"Expected object to be a ray.ObjectRef, got {type(obj)}"
                    assert isinstance(obj, ray.ObjectRef), asrt_msg
                    ray_free([obj])
                except Exception as e:  # noqa: BLE001
                    log(ERROR, "Error in custom RayGC-%s.", process_name, exc_info=e)

    if not list_of_threads:
        thread = threading.Thread(
            target=collect_ray_garbage,
            name=CUSTOM_RAY_GC,
        )
        list_of_threads.append(thread)
        thread.start()

    try:
        yield
    except Exception as e:
        log(
            ERROR,
            "Error in the main scoped caught by RayGC-%s context manager.",
            process_name,
            exc_info=e,
        )
        raise
    finally:
        if join_at_the_end:
            log(DEBUG, "RayGC-%s: Closing.", process_name)
            stop_event.set()  # Signal threads to stop
            for thread in list_of_threads:
                thread.join(timeout=timeout * 2)  # Add timeout to prevent hanging
                if thread.is_alive():
                    log(
                        ERROR,
                        "RayGC-%s: Thread did not terminate gracefully.",
                        process_name,
                    )


@contextmanager
def custom_shm_garbage_collector(
    garbage_queue: Queue[ray.ObjectRef | str],
    list_of_threads: list[threading.Thread],
    process_name: str,
    timeout: float = 10,
    *,
    join_at_the_end: bool = True,
) -> Generator[None, Any, None]:
    """Context manager for custom shared memory garbage collection.

    This context manager creates a thread that continuously collects and cleans
    shared memory objects from a given queue. The thread runs in the background
    and stops when the context manager exits.

    Parameters
    ----------
    garbage_queue : Queue[ray.ObjectRef | str]
        A queue containing shared memory object names to be cleaned.
    list_of_threads : list[threading.Thread]
        A list of threads to append the garbage collection thread to.
    process_name : str
        The name of the process.
    timeout : float, optional
        The time in seconds to wait between garbage collection cycles, by default 10.
    join_at_the_end : bool, optional
        Whether to join the garbage collection thread at the end, by default
        True.

    Yields
    ------
    None
        The context manager yields control back to the caller.

    """
    stop_event = threading.Event()

    def collect_shm_garbage() -> None:
        while not stop_event.is_set():
            time.sleep(timeout)
            while not garbage_queue.empty():
                try:
                    obj = garbage_queue.get(block=True, timeout=0.1)
                    asrt_msg = f"Expected object to be a string, got {type(obj)}"
                    assert isinstance(obj, str), asrt_msg
                    clean_shm_by_name(obj)
                except Exception as e:  # noqa: BLE001
                    log(ERROR, "Error in custom ShmGC-%s.", process_name, exc_info=e)

    if not list_of_threads:
        thread = threading.Thread(
            target=collect_shm_garbage,
            name=CUSTOM_SHM_GC,
        )
        list_of_threads.append(thread)
        thread.start()

    try:
        yield
    except Exception as e:
        log(
            ERROR,
            "Error in the main scoped caught by ShmGC-%s context manager.",
            process_name,
            exc_info=e,
        )
        raise
    finally:
        if join_at_the_end:
            log(DEBUG, "ShmGC-%s: Closing.", process_name)
            stop_event.set()  # Signal threads to stop
            for thread in list_of_threads:
                thread.join(timeout=timeout * 2)  # Add timeout to prevent hanging
                if thread.is_alive():
                    log(
                        ERROR,
                        "ShmGC-%s: Thread did not terminate gracefully.",
                        process_name,
                    )


def set_parameters_in_place(
    old_parameters: NDArrays,
    new_parameters: NDArrays,
    parameter_pos: Iterable[int],
) -> None:
    """Update model parameters in shared memory.

    This function updates the model parameters stored in the shared memory segment with
    new parameters.

    Parameters
    ----------
    old_parameters : NDArrays
        The current model parameters stored in shared memory.
    new_parameters : NDArrays
        The new model parameters to be set in shared memory.
    parameter_pos : Iterable[int]
        An iterable of indices specifying which parameters to update. If None, all
        parameters are updated. Default is None.

    """
    assert len(old_parameters) == len(new_parameters), (
        "The number of parameters in old_parameters and new_parameters must be the "
        "same. Was: "
        f"{len(old_parameters)} != {len(new_parameters)} instead."
    )

    def update_parameter(i: int) -> None:
        if len(new_parameters[i].shape) == 0:
            old_parameters[i] = new_parameters[i]
        else:
            old_parameters[i][:] = new_parameters[i][:]

    with concurrent.futures.ThreadPoolExecutor() as executor:
        executor.map(update_parameter, parameter_pos)


def get_random_model_from_config(
    llm_cfg: DictConfig,
) -> torch.nn.Module:
    """Create a randomly initialized model from the provided configuration.

    This function initializes a model based on the provided configuration object. It
    handles the configuration preprocessing by resolving variables, filtering out
    problematic entries, building the tokenizer, and creating the model on CPU.

    Parameters
    ----------
    llm_cfg : DictConfig
        The configuration object used to create the randomly initialized model.

    Returns
    -------
    torch.nn.Module
        A randomly initialized model instance created according to the configuration.

    Example
    -------
    >>> from omegaconf import OmegaConf
    >>> cfg = OmegaConf.create({...})
    >>> model = get_random_model_from_config(cfg)

    Raises
    ------
    TypeError
        If model_config is not a dict or name is not a string.

    """
    # Resolve all interpolation variables as early as possible
    OmegaConf.resolve(llm_cfg)
    OmegaConf.set_struct(llm_cfg, value=False)
    # Deep copy the configuration to prevent any dangerous modification
    internal_cfg = copy.deepcopy(llm_cfg)

    # Filter deprecation warning from torch internal usage
    warnings.filterwarnings(
        action="ignore",
        category=UserWarning,
        message=(
            "torch.distributed.*_base is a private functionand will be deprecated.*"
        ),
    )

    # NOTE: We need to extract a set of global variables from the configuration object
    # to prevent the dataclass creator to crash
    internal_cfg.pop("data_local", None)
    internal_cfg.pop("data_remote", None)
    internal_cfg.pop("global_seed", None)
    internal_cfg.pop("name", None)
    # NOTE: This contains OUR global parameters for the ICL tasks that the
    # `make_dataclass_and_log_config` cannot interpret, so we need to pop it
    _icl_tasks_config_dict: dict[str, Any] | None = internal_cfg.pop(
        "icl_tasks_config",
        None,
    )
    _logged_cfg, train_cfg = make_dataclass_and_log_config(
        internal_cfg,
        TrainConfig,
        TRAIN_CONFIG_KEYS,
        transforms="all",
        icl_tasks_required=internal_cfg.get("icl_tasks", None) is not None,
    )
    # Check for incompatibilities between the model and data loaders
    validate_config(train_cfg)
    # Build tokenizer
    tokenizer_name = train_cfg.tokenizer["name"]
    tokenizer_kwargs = train_cfg.tokenizer.get("kwargs", {})
    tokenizer = build_tokenizer(tokenizer_name, tokenizer_kwargs)
    # Get model while forcing cpu to prevent any GPU allocation
    model_config = train_cfg.model
    if not isinstance(model_config, dict):
        msg = f"Expected model_config to be a dict, got {type(model_config)}"
        raise TypeError(msg)

    model_config["init_device"] = "cpu"
    name = model_config.pop("name")
    if not isinstance(name, str):
        msg = f"Expected name to be a string, got {type(name)}"
        raise TypeError(msg)

    if not isinstance(model_config, dict):
        msg = f"Expected model_config to be a dict, got {type(model_config)}"
        raise TypeError(msg)

    init_context = process_init_device(model_config, None, None)
    model = build_composer_model(
        name=name,
        tokenizer=tokenizer,
        init_context=init_context,
        master_weights_dtype=model_config.pop("master_weights_dtype", None),
        cfg=model_config,
    )
    # Force model to cpu
    model.cpu()
    return model


def get_random_model_states_from_config(
    cfg: BaseConfig,
) -> tuple[NDArrays, tuple[tuple[str, "ModelStateNames"], ...]]:
    """Retrieve random model states from the configuration.

    This function generates random model states based on the provided configuration.
    It uses the LLM configuration to create the model states and their corresponding
    layer names and types.

    Parameters
    ----------
    cfg : BaseConfig
        The configuration object containing the LLM config.

    Returns
    -------
    'tuple[NDArrays, tuple[tuple[str, ModelStateNames], ...]]'
        The random model states and their layer names and types.

    """
    # Resolve the OmegaConf configuration
    llm_config = cfg.llm_config
    OmegaConf.resolve(llm_config)
    OmegaConf.set_struct(llm_config, value=False)
    # Resolve the configuration and return the randomly initialized model instantiated
    # on the CPU
    model = get_random_model_from_config(llm_config)
    # Get the parameters dictionary
    tensors_dict = get_tensors_dict(
        model,
        # NOTE: The random model has all trainable parameters that we should care
        # about with the correct flag set
        only_requires_grad=True,
        # NOTE: We need to sort the parameters to ensure that the order is
        # consistent with the names later
        sort_dict=True,
    )

    if n_experts := get_n_experts_from_model(model):
        # For MPT models where we need to split the experts
        log(DEBUG, "Splitting experts in the randomly init model parameters.")
        tensors_dict = split_experts_layers_tensors(
            tensors_dict=tensors_dict,
            n_experts=n_experts,
        )
    # Get the trainable model parameters as NDArrays
    parameters_ndarrays_dict = {
        name: val.detach().to("cpu").numpy() for name, val in tensors_dict.items()
    }
    parameters_ndarrays = list(parameters_ndarrays_dict.values())
    parameters_names = list(parameters_ndarrays_dict.keys())
    # Create two zeroed-out copies of the parameters to simulate the momenta shape,
    # i.e., the addition of model states
    zeros = [np.zeros_like(param) for param in parameters_ndarrays]
    parameters_ndarrays.extend(zeros)
    parameters_ndarrays.extend(zeros)

    return parameters_ndarrays, create_layers_names_and_types_from_names(
        parameters_names,
    )


def match_shapes(
    remote_param: torch.Tensor,
    local_param: torch.nn.Parameter | torch.Tensor,
) -> torch.Tensor:
    """Match the shape of a remote parameter to a local parameter if possible.

    This function attempts to match the shape of `remote_param` to `local_param`.
    If the shapes are already identical, `remote_param` is returned unchanged.
    If the total number of elements in both tensors differs, `remote_param`
    is returned unchanged as their shapes are fundamentally incompatible for reshaping.
    If the number of elements is the same but shapes differ:
    - If `local_param` is 1-dimensional, `remote_param` is flattened.
    - If `local_param` is N-dimensional (N > 1 or N = 0 for scalar), `remote_param`
      is reshaped to `local_param.shape`.

    This is useful in scenarios like FSDP model loading, where remote parameters
    (e.g., from a sharded state_dict, potentially flattened) need to conform to
    the local model's parameter shapes.

    Parameters
    ----------
    remote_param : torch.Tensor
        The tensor whose shape needs to be matched.
    local_param : torch.nn.Parameter | torch.Tensor
        The tensor that serves as the shape reference.

    Returns
    -------
    torch.Tensor
        The `remote_param` tensor, potentially reshaped or flattened to match
        `local_param.shape`, or the original `remote_param` if shapes
        cannot be matched due to differing numbers of elements.

    """
    # 1. If shapes already match, return remote_param unchanged.
    if remote_param.shape == local_param.shape:
        return remote_param

    # 2. If the total number of elements differs, shapes cannot be matched.
    #    Return remote_param unchanged.
    if remote_param.numel() != local_param.numel():
        return remote_param

    # At this point, numel() is the same, but shapes are different.
    # 3. If local_param is 1D, flatten remote_param.
    #    Since numel() is the same, remote_param.flatten() will have the same
    #    shape as local_param.
    if local_param.ndim == 1:
        return remote_param.flatten()

    # NOTE: We don't expect scalar tensors (0D) in this context.
    return remote_param.reshape(local_param.shape)


def print_all_keys(d: dict, indent: int = 0) -> list[str]:
    """Recursively extract all keys from nested dictionaries.

    This function traverses a dictionary and collects all keys, including those from
    nested dictionaries. It returns a flat list containing string representations of
    all keys found.

    Parameters
    ----------
    d : dict
        The dictionary to extract keys from. Can contain nested dictionaries.
    indent : int, optional
        The current indentation level for formatting, by default 0.
        This parameter is primarily used for recursive calls to track hierarchy.

    Returns
    -------
    list[str]
        A list containing string representations of all keys in the dictionary and
        its nested dictionaries.

    Example
    -------
    >>> nested_dict = {'a': 1, 'b': {'c': 2, 'd': {'e': 3}}}
    >>> print_all_keys(nested_dict)
    ['a', 'b', 'c', 'd', 'e']

    Notes
    -----
    The function ignores the structure of the nested dictionaries and returns
    a flat list of all keys. The original indentation level is not reflected
    in the returned list.

    """
    message: list[str] = []
    for key, value in d.items():
        message.append(str(key))
        if isinstance(value, dict):
            message.extend(print_all_keys(value, indent + 2))
    return message


# NOTE: This is unused but it is kept for reference
def set_trainable_params_dict(
    model: torch.nn.Module,
    parameters_dict: OrderedDict[str, torch.Tensor],
) -> None:
    """Set the trainable parameters of a model.

    Parameters
    ----------
    model : torch.nn.Module
        The model.
    parameters_dict : OrderedDict[str, torch.Tensor]
        The dictionary of parameters.

    Raises
    ------
    ValueError
        If the shapes don't match.

    """
    # NOTE: This function is weird because the encapsulation done to support FSDP and
    # DDP is weird. Since they are both likely to change, we MUST maintain this very
    # well and implement as many checkers as we can.
    if hasattr(model, "model") and type(model.model) is FullyShardedDataParallel:
        if model.model is None:  # type: ignore[reportUnnecessaryComparison]
            error_message = "Model is None"
            raise ValueError(error_message)
        inner_model = model.model
        # NOTE: This doesn't work in the case in use_orig_params is True if the FSDP
        # configuration as the tensors returned are flattened breaking some assumptions
        # of the rest of the codebase
        with FullyShardedDataParallel.summon_full_params(
            inner_model,
            recurse=True,
            # Writing back is not compatible with rank 0 only
            writeback=True,
            rank0_only=False,
            # Prevent moving to CPU device
            offload_to_cpu=False,
            with_grads=False,
        ):
            # NOTE: !!! THIS REQUIRES INVESTIGATION AS IT DOESN'T WORK AS EXPECTED !!!
            # NOTE: This parameter dict using the above parameters, i.e.,
            # (recurse=True, writeback=True, rank0_only=False, offload_to_cpu=False,
            # with_grads=False,), won't be complete in any rank if the model i sharded.
            # Each rank will have zero-size tensors for those layers that are not
            # "living" in there and the flattened/unflattened complete tensors for those
            # blocks living there.
            # NOTE: If the FSDP configuration use the original parameters
            # (use_orig_params=true), then the tensors in rank 0 have the correct
            # original shape. In the other ranks they are flattened anyway.
            for name, param in inner_model.named_parameters():
                # Set the parameters only if they require gradients & have non-zero size
                if param.requires_grad:
                    param.data = parameters_dict[name].to(param.device)
    else:
        for name, param in model.named_parameters():
            # Set the parameters only if they require gradients
            if param.requires_grad:
                # TODO(anon): Handle globally modulse prefixes filtering
                # NOTE: DDP pre-pends "module." to the name of the parameter
                if name.startswith("module."):
                    param.data = parameters_dict[name.replace("module.", "")].to(
                        param.device,
                    )
                # Single GPU
                else:
                    param.data = parameters_dict[name].to(param.device)
    dist.barrier()


def create_tensors_dict(
    parameters_names: list[str],
    parameters: NDArrays,
) -> OrderedDict[str, torch.Tensor]:
    """Create a dictionary of tensors from parameter names and values.

    Parameters
    ----------
    parameters_names : list[str]
        The names of the parameters.
    parameters : NDArrays
        The values of the parameters.

    Returns
    -------
    OrderedDict[str, torch.Tensor]
        The dictionary of tensors.

    """
    ret_dict = OrderedDict(
        {
            name: torch.as_tensor(value)
            for name, value in zip(parameters_names, parameters, strict=True)
        },
    )
    # TODO(anon): Check what differs here compared to before and why these are all
    # False now
    for tensor in ret_dict.values():
        tensor.requires_grad = True
    return ret_dict


def set_frozen_layers_tensors_dict(
    tensors_dict: OrderedDict[str, torch.Tensor],
    frozen_layers: list[str] | None = None,
) -> OrderedDict[str, torch.Tensor]:
    """Set the frozen layers in a dictionary of tensors.

    Parameters
    ----------
    tensors_dict : OrderedDict[str, torch.Tensor]
        The dictionary of tensors.
    frozen_layers : list[str] | None
        The frozen layers, by default None.

    Returns
    -------
    OrderedDict[str, torch.Tensor]
        The dictionary of tensors with frozen layers set.

    """
    if frozen_layers is not None:
        for name in tensors_dict:
            if any(layer in name or name in layer for layer in frozen_layers):
                log(DEBUG, "Freezing %s", name)
                tensors_dict[name].requires_grad = False
    return tensors_dict


def wandb_init(
    wandb_enabled: bool,  # noqa: FBT001
    *args: dict,
    **kwargs: Any,  # noqa: ANN401
) -> NoOpContextManager | Any | None:  # noqa: ANN401
    """Initialize wandb if enabled.

    Parameters
    ----------
    wandb_enabled : bool
        Whether wandb is enabled.
    args : dict
        The arguments.
    kwargs : dict
        The keyword arguments.

    Returns
    -------
    NoOpContextManager | Any | None
        The wandb context

    Raises
    ------
    ValueError
        If the name is not a string.

    """
    if wandb_enabled:
        # Add server suffix to the name of the run
        name = kwargs.pop("name", "")
        if type(name) is not str:
            error = f"Name must be a string, not {type(name)}"
            raise ValueError(error)
        name += "_server"
        return wandb.init(*args, **kwargs, name=name)  # type: ignore[arg-type,misc]

    return NoOpContextManager()


def sum_of_squares(arrays: NDArrays) -> float:
    """Compute the sum of squares of a list of arrays.

    Parameters
    ----------
    arrays : NDArrays
        List of arrays to compute the sum of squares of.

    Returns
    -------
    float
        The sum of squares of the list of arrays.

    """
    return sum(np.sum(np.square(arr)) for arr in arrays)


def sum_of_squares_per_layer(
    model_state: tuple[NDArrays, list[str]],
) -> dict[str, float]:
    """Return the sum of squares for a model, indexed by its layer.

    This function computes the sum of squares for each layer of the model
    in a threaded fashion to make use of the parallism oppertunity given
    that none of the calculations depend on one another.

    Args:
        model_state (tuple[NDArrays, list[str]]): The model state containing
        the arrays and the corresponding layer names as a tuple.

    Returns:
        dict[str, float]: The sum of squares for each layer as a dictionary.

    """

    def compute_single_process(args: tuple[NDArray, str]) -> tuple[str, float]:
        """Compuute the sum of squares for a layer.

        Parameters
        ----------
        args : tuple[NDArray, str]
            A tuple containing the:
                - The array of the layer in question
                - The name of the layer in question

        Returns
        -------
        tuple[str, float]
            A tuple containing the layer name and its corresponding sum of squares.

        """
        array, layer_name = args
        return layer_name, np.sum(np.square(array))

    arrays, layer_names = model_state

    with ThreadPoolExecutor() as executor:
        results = executor.map(
            compute_single_process,
            zip(arrays, layer_names, strict=False),
        )

    return dict(results)


def l2_norm(arrays: NDArrays) -> float:
    """Compute the L2 norm of a list of arrays.

    Parameters
    ----------
    arrays : NDArrays
        List of arrays to compute the L2 norm of.

    Returns
    -------
    float
        The L2 norm of the list of arrays.

    """
    return float(np.sqrt(sum_of_squares(arrays)))


def get_device() -> device_type:
    """Determine which device to use for PyTorch.

    Returns
    -------
        str: device for PyTorch

    """
    device = "cpu"
    if torch.cuda.is_available():
        device = "cuda"
    elif torch.backends.mps.is_available() and torch.backends.mps.is_built():
        device = "mps"
    return cast("device_type", device)


def get_n_cuda_devices() -> int:
    """Get the number of CUDA devices available.

    Returns
    -------
        int: number of CUDA devices available

    """
    if "cuda" in str(get_device()):
        return torch.cuda.device_count()
    return 0


def get_n_cpu_cores() -> int | None:
    """Get the number of CPU cores available.

    Returns
    -------
        int | None: number of CPU cores available

    """
    try:
        cpus = len(psutil.Process().cpu_affinity())  # type: ignore[reportArgumentType]
    except AttributeError:
        cpus = psutil.cpu_count()
    return cpus


def merge_freq_dicts(
    a: dict[int, int],
    b: dict[int, int],
) -> dict[int, int]:
    """Merge two frequency dictionaries.

    Parameters
    ----------
    a : dict[int, int]
        The first frequency dictionary.
    b : dict[int, int]
        The second frequency dictionary.

    Returns
    -------
    dict[int, int]
        The merged frequency dictionary.

    """
    return a | {k: a.get(k, 0) + v for k, v in b.items()}


def get_unigram_probabilities_tensor(
    stream_freq_dict: dict[int, int],
) -> torch.Tensor:
    """Get the unigram probabilities tensor.

    Parameters
    ----------
    stream_freq_dict : dict[int, int]
        The frequency dictionary.

    Returns
    -------
    torch.Tensor
        The unigram probabilities tensor.

    """
    total_tokens = float(sum(v for v in stream_freq_dict.values()))
    probabilities = {k: v / total_tokens for k, v in stream_freq_dict.items()}
    # Get the max token id
    max_token_id = max(stream_freq_dict.keys())
    # Convert to dense tensor
    probabilities_tensor = torch.zeros(max_token_id + 1)
    for k, v in probabilities.items():
        probabilities_tensor[k] = v
    return probabilities_tensor


def is_literal_for_ast(s: str) -> bool:
    """Check if the given str can be evaluate as a literal.

    Parameters
    ----------
    s : str
        The string to check.

    Returns
    -------
    bool
        Whether the string can be evaluated as a literal

    """
    try:
        ast.literal_eval(s)
    except ValueError:
        return False
    return True


def get_cuda_compute_capability() -> tuple[int, int] | None:
    """Get the CUDA compute capability of the current device.

    Returns
    -------
    tuple[int, int] | None
        The CUDA compute capability of the current device.
        If CUDA is not available, it returns None. It returns
        the major and minor compute capability otherwise.

    """
    if "cuda" in str(get_device()):
        device = torch.cuda.current_device()
        return torch.cuda.get_device_capability(device)
    return None
