# -*- coding: utf-8 -*-

# Copyright (c) 2016-2026 by University of Kassel and Fraunhofer Institute for Energy Economics
# and Energy System Technology (IEE), Kassel. All rights reserved.

from __future__ import annotations

import logging
from typing import Iterable, Sequence

import numpy as np
from numpy import nan, any as np_any

from pandapower.auxiliary import pandapowerNet
from pandapower.pp_types import Int, SwitchElementType, SwitchType
from pandapower.create._utils import (
    _check_element,
    _check_multiple_elements,
    _get_index_with_check,
    _get_multiple_index_with_check,
    _set_entries,
    _set_multiple_entries,
)

logger = logging.getLogger(__name__)


def create_switch(
    net: pandapowerNet,
    bus: Int,
    element: Int,
    et: SwitchElementType,
    closed: bool = True,
    type: SwitchType | None = None,
    name: str | None = None,
    index: Int | None = None,
    z_ohm: float = 0,
    in_ka: float = nan,
    **kwargs,
) -> Int:
    """
    Adds a switch in the net["switch"] table.

    Switches can be either between two buses (bus-bus switch) or at the end of a line or transformer
    element (bus-element switch).

    Two buses that are connected through a closed bus-bus switches are fused in the power flow if
    the switch is closed or separated if the switch is open.

    An element that is connected to a bus through a bus-element switch is connected to the bus
    if the switch is closed or disconnected if the switch is open.

    Parameters:
        net: The net within which this switch should be created
        bus: The bus that the switch is connected to
        element: index of the element
        et: element type
        
            - "l" = switch between bus and line
            - "t" = switch between bus and transformer
            - "t3" = switch between bus and transformer3w
            - "b" = switch between two buses
        
        closed: switch position:
         
            - False = open
            - True = closed
        
        type: indicates the type of switch
        
            - "LS" = Load Switch
            - "CB" = Circuit Breaker
            - "LBS" = Load Break Switch
            - "DS" = Disconnecting Switch
        
        z_ohm: indicates the resistance of the switch, which has effect only on bus-bus switches, if sets to 0, the
            buses will be fused like before, if larger than 0 a branch will be created for the switch which has also
            effects on the bus mapping
        name: The name for this switch
        in_ka: maximum current that the switch can carry normal operating conditions without tripping

    Returns:
        The unique switch_id of the created switch

    Example:
        >>> create_switch(net, bus=0, element=1, et='b', type="LS", z_ohm=0.1)
        >>> create_switch(net, bus=0, element=1, et='l')

    """
    _check_element(net, bus)
    match et:
        case "l":
            elm_tab = "line"
            if element not in net[elm_tab].index:
                raise UserWarning("Unknown line index")
            if net[elm_tab]["from_bus"].loc[element] != bus and net[elm_tab]["to_bus"].loc[element] != bus:
                raise UserWarning(f"Line {element} not connected to bus {bus}")
        case "t":
            elm_tab = "trafo"
            if element not in net[elm_tab].index:
                raise UserWarning("Unknown bus index")
            if net[elm_tab]["hv_bus"].loc[element] != bus and net[elm_tab]["lv_bus"].loc[element] != bus:
                raise UserWarning(f"Trafo {element} not connected to bus {bus}")
        case "t3":
            elm_tab = "trafo3w"
            if element not in net[elm_tab].index:
                raise UserWarning("Unknown trafo3w index")
            if (
                net[elm_tab]["hv_bus"].loc[element] != bus
                and net[elm_tab]["mv_bus"].loc[element] != bus
                and net[elm_tab]["lv_bus"].loc[element] != bus
            ):
                raise UserWarning(f"Trafo3w {element} not connected to bus {bus}")
        case "b":
            _check_element(net, element)
        case _:
            raise UserWarning("Unknown element type")

    index = _get_index_with_check(net, "switch", index)

    entries = {
        "bus": bus,
        "element": element,
        "et": et,
        "closed": closed,
        "type": type,
        "name": name,
        "z_ohm": z_ohm,
        "in_ka": in_ka,
        **kwargs,
    }
    _set_entries(net, "switch", index, entries=entries)

    return index


def create_switches(
    net: pandapowerNet,
    buses: Sequence,
    elements: Sequence,
    et: SwitchElementType | Sequence[str],
    closed: bool = True,
    type: SwitchType | None = None,
    name: Iterable[str] | None = None,
    index: Int | Iterable[Int] | None = None,
    z_ohm: float = 0,
    in_ka: float = nan,
    **kwargs,
) -> Int:
    """
    Adds a switch in the net["switch"] table.

    Switches can be either between two buses (bus-bus switch) or at the end of a line or transformer
    element (bus-element switch).

    Two buses that are connected through a closed bus-bus switches are fused in the power flow if
    the switch is closed or separated if the switch is open.

    An element that is connected to a bus through a bus-element switch is connected to the bus
    if the switch is closed or disconnected if the switch is open.

    Parameters:
        net: The net within which this switch should be created
        buses: The bus that the switch is connected to
        element: index of the element
        et: element type
         
            - "l" = switch between bus and line
            - "t" = switch between bus and transformer
            - "t3" = switch between bus and transformer3w
            - "b" = switch between two buses
        
        closed: switch position
         
            - False = open
            - True = closed
        
        type: indicates the type of switch
         
            - "LS" = Load Switch
            - "CB" = Circuit Breaker
            - "LBS" = Load Break Switch or
            - "DS" = Disconnecting Switch
        
        z_ohm: indicates the resistance of the switch, which has effect only on bus-bus switches, if sets to 0,
            the buses will be fused like before, if larger than 0 a branch will be created for the switch which has also
            effects on the bus mapping
        name: The name for this switch
        in_ka: maximum current that the switch can carry normal operating conditions without tripping

    Returns:
        List of ids of the created switches

    Example:
        >>> create_switches(net, buses=[0, 1], element=1, et='b', type="LS", z_ohm=0.1)
        >>> create_switches(net, buses=[0, 1], element=1, et='l')
    """
    index = _get_multiple_index_with_check(net, "switch", index, len(buses), name="Switches")
    _check_multiple_elements(net, buses)
    rel_els = ["b", "l", "t", "t3"]
    matcher = {"b": ["bus", "buses"], "l": ["line", "lines"], "t": ["trafo", "trafos"], "t3": ["trafo3w", "trafo3ws"]}
    for typ in rel_els:
        if et == typ:
            _check_multiple_elements(net, elements, *matcher[typ])
    if np.any(np.isin(et, ["b", "l", "t"])):
        mask_all = np.array([False] * len(et))
        for typ in rel_els:
            et_arr = np.array(et)
            el_arr = np.array(elements)
            mask = et_arr == typ
            mask_all |= mask
            _check_multiple_elements(net, el_arr[mask], *matcher[typ])
        not_def = ~mask_all
        if np.any(not_def):
            raise UserWarning(f"et type {et_arr[not_def]} is not implemented")
    else:
        raise UserWarning(f"et type {et} is not implemented")

    b_arr = np.array(buses)[:, None]
    el_arr = np.array(elements)
    et_arr = np.array([et] * len(buses) if isinstance(et, str) else et)
    # Ensure switches are connected correctly.
    for typ, table, joining_busses in [
        ("l", "line", ["from_bus", "to_bus"]),
        ("t", "trafo", ["hv_bus", "lv_bus"]),
        ("t3", "trafo3w", ["hv_bus", "mv_bus", "lv_bus"]),
    ]:
        el = el_arr[et_arr == typ]
        bs = net[table].loc[el, joining_busses].values
        not_connected_mask = ~np.isin(b_arr[et_arr == typ], bs)
        if np_any(not_connected_mask):
            bus_element_pairs = zip(
                el_arr[et_arr == typ][:, None][not_connected_mask].tolist(),
                b_arr[et_arr == typ][not_connected_mask].tolist(),
            )
            raise UserWarning(f"{table.capitalize()} not connected ({table} element, bus): {list(bus_element_pairs)}")

    entries = {
        "bus": buses,
        "element": elements,
        "et": et,
        "closed": closed,
        "type": type,
        "name": name,
        "z_ohm": z_ohm,
        "in_ka": in_ka,
        **kwargs,
    }

    _set_multiple_entries(net, "switch", index, entries=entries)

    return index
