r"""Miscelaneous plugin helpers."""

__all__ = [
    "load_cards",
]

import os
import sys
import torch
import yaml

from collections.abc import Iterator
from contextlib import contextmanager
from types import ModuleType, SimpleNamespace


def as_torch_dtype(name: str | None = None) -> torch.dtype:
    if name is None:
        return None

    dtype = getattr(torch, name, None)

    if isinstance(dtype, torch.dtype):
        return dtype
    else:
        raise ValueError(f"Unknown data type '{name}'.")


def load_cards(plugin: ModuleType) -> dict[str, SimpleNamespace]:
    r"""Returns the name-card mapping of pre-trained models available in a plugin.

    Arguments:
        plugin: The plugin module.

    Example:
        >>> cards = load_cards(azula.plugins.adm)
        >>> list(cards.keys())
        ['imagenet_64x64_cond',
         'imagenet_128x128_cond',
         'imagenet_256x256',
         'imagenet_256x256_cond',
         'imagenet_512x512_cond',
         'ffhq_256x256']
    """

    if isinstance(plugin, str):
        plugin = sys.modules[plugin]

    file = os.path.join(os.path.dirname(plugin.__file__), "cards.yaml")

    assert os.path.exists(file), f"{plugin} is not a plugin"

    with open(file) as f:
        cards = yaml.safe_load(f)

    for card in cards.values():
        if "dtype_map" in card:
            card["dtype_map"] = {k: as_torch_dtype(v) for k, v in card["dtype_map"].items()}

    return {name: SimpleNamespace(**card) for name, card in cards.items()}


@contextmanager
def patch_diffusers() -> Iterator[None]:
    from unittest.mock import patch

    with (
        patch("transformers.modeling_utils.logger.warning"),
        patch("transformers.modeling_utils.logger.warning_once"),
    ):
        yield


def patch_mmap_safetensors(*modules: torch.nn.Module) -> None:
    for m in modules:
        for p in m.parameters():
            p.data.sum()

        for b in m.buffers():
            b.data.sum()
