# 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!

from enum import Enum, EnumMeta
from typing import Any, Mapping

# pylint: disable=invalid-name

# 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__.keys()


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"
    """

    input = "input"
    sym_output = "sym_output"
    asym_output = "asym_output"
    update = "update"
    sc_output = "sc_output"


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

    - Examples:

        - ComponentType.node = "node"
        - ComponentType.line = "line"
    """

    node = "node"
    line = "line"
    link = "link"
    transformer = "transformer"
    transformer_tap_regulator = "transformer_tap_regulator"
    three_winding_transformer = "three_winding_transformer"
    sym_load = "sym_load"
    sym_gen = "sym_gen"
    asym_load = "asym_load"
    asym_gen = "asym_gen"
    shunt = "shunt"
    source = "source"
    sym_voltage_sensor = "sym_voltage_sensor"
    asym_voltage_sensor = "asym_voltage_sensor"
    sym_power_sensor = "sym_power_sensor"
    asym_power_sensor = "asym_power_sensor"
    fault = "fault"


# pylint: enable=invalid-name


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


def _map_to_datatypes(data: Mapping[Any, 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: Any) -> ComponentType:
    """Helper function to transform component str to ComponentType."""
    if isinstance(component, str):
        return ComponentType[component]
    return component


def _map_to_component_types(data: Mapping[Any, 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
