# SPDX-FileCopyrightText: Contributors to the Power Grid Model project <powergridmodel@lfenergy.org>
#
# SPDX-License-Identifier: MPL-2.0

"""Data types for power grid model dataset and component types."""

# This file is automatically generated. DO NOT modify it manually!

{%- set dataset_types = all_map.keys() %}
{%- set components = all_map['input'].keys() %}

from collections.abc import Mapping
from enum import Enum, EnumMeta
from typing import Any, TypeAlias, TypeVar

# fmt: off

class _MetaEnum(EnumMeta):
    def __contains__(cls, member):
        """
        Check if member is part of the Enum.

        Args:
            member: Member to check.

        Returns:
            bool: True if the member is part of the Enum, False otherwise.
        """
        return member in cls.__members__


class DatasetType(str, Enum, metaclass=_MetaEnum):
    """
    A DatasetType is the type of a :class:`Dataset` in power grid model.

    - Examples:

        - DatasetType.input = "input"
        - DatasetType.update = "update"
    """
{% for dataset_type in dataset_types %}
    {{ dataset_type }} = "{{ dataset_type }}"
{%- endfor %}


class ComponentType(str, Enum, metaclass=_MetaEnum):
    """
    A ComponentType is the type of a grid component.

    - Examples:

        - ComponentType.node = "node"
        - ComponentType.line = "line"
    """
{% for component in components %}
    {{ component }} = "{{ component }}"
{%- endfor %}


DatasetTypeLike: TypeAlias = DatasetType | str
DatasetTypeVar = TypeVar("DatasetTypeVar", bound=DatasetTypeLike)  # helper used for type deduction

ComponentTypeLike: TypeAlias = ComponentType | str
ComponentTypeVar = TypeVar("ComponentTypeVar", bound=ComponentTypeLike)  # helper used for type deduction


def _str_to_datatype(data_type: DatasetTypeLike) -> DatasetType:
    """Helper function to transform data_type str to DatasetType."""
    if isinstance(data_type, DatasetType):
        return data_type
    return DatasetType[data_type]


def _map_to_datatypes(data: Mapping[DatasetTypeVar, Any]) -> dict[DatasetType, Any]:
    """Helper function to map datatype str keys to DatasetType."""
    return {_str_to_datatype(key): value for key, value in data.items()}


def _str_to_component_type(component: ComponentTypeLike) -> ComponentType:
    """Helper function to transform component str to ComponentType."""
    if isinstance(component, ComponentType):
        return component
    return ComponentType[component]


def _map_to_component_types(data: Mapping[ComponentTypeVar, Any]) -> dict[ComponentType, Any]:
    """Helper function to map componenttype str keys to ComponentType."""
    return {_str_to_component_type(key): value for key, value in data.items()}

# fmt: on

