"""
Grid Loader — IEEE Test Systems via pandapower
Loads IEEE 33-bus (case33bw) and IEEE 118-bus (case118) systems directly
from pandapower's built-in network library.  No external file I/O required.
"""
from __future__ import annotations

import copy
from typing import Literal

import pandas as pd
import pandapower as pp
import pandapower.networks as pn


_LOADERS = {
    "case33bw": pn.case33bw,
    "case118": pn.case118,
}

# IEEE 33-bus tie line indices (out of service in default config)
IEEE33_TIE_LINES = [32, 33, 34, 35, 36]
# Default open lines for IEEE 33-bus in OptiQ (override to 3 OOS)
IEEE33_DEFAULT_OPEN_LINES = [32, 33, 34]


def load_network(system: Literal["case33bw", "case118"] = "case33bw") -> pp.pandapowerNet:
    """Return a fresh pandapower network for the given IEEE test system.

    Parameters
    ----------
    system : str
        One of ``"case33bw"`` (IEEE 33-bus distribution) or
        ``"case118"`` (IEEE 118-bus transmission).

    Returns
    -------
    pp.pandapowerNet
        A ready-to-simulate network object.
    """
    loader = _LOADERS.get(system)
    if loader is None:
        raise ValueError(f"Unknown system '{system}'. Choose from {list(_LOADERS)}")
    net = loader()
    net["optiq_system"] = system
    net["optiq_is_distribution"] = system == "case33bw"

    if system == "case33bw":
        # Ensure only 3 lines are open by default for 33-bus
        net["optiq_tie_lines"] = list(IEEE33_TIE_LINES)
        net["optiq_default_open_lines"] = list(IEEE33_DEFAULT_OPEN_LINES)
        net.line["in_service"] = True
        for idx in IEEE33_DEFAULT_OPEN_LINES:
            if idx in net.line.index:
                net.line.at[idx, "in_service"] = False

    return net


def clone_network(net: pp.pandapowerNet) -> pp.pandapowerNet:
    """Deep-copy a pandapower network (useful for parallel scenarios)."""
    return copy.deepcopy(net)


def get_line_info(net: pp.pandapowerNet) -> dict:
    """Identify switchable lines (all lines), in-service and out-of-service.

    In IEEE 33-bus, reconfiguration is modelled by toggling ``line.in_service``.
    Lines that are out of service in the default config are the tie lines.

    Returns
    -------
    dict with keys:
        ``"all"``          – all line indices
        ``"in_service"``   – indices of lines currently in service (closed)
        ``"out_of_service"`` – indices of lines currently out of service (open/tie)
        ``"tie_lines"``    – the default tie line indices (for IEEE 33-bus)
    """
    all_idx = net.line.index.tolist()
    in_svc = net.line.index[net.line.in_service].tolist()
    out_svc = net.line.index[~net.line.in_service].tolist()
    tie_lines = net.get("optiq_tie_lines", out_svc)
    n_required_open = len(net.line) - (len(net.bus) - 1)
    return {
        "all": all_idx,
        "in_service": in_svc,
        "out_of_service": out_svc,
        "tie_lines": list(tie_lines),
        "n_required_open": n_required_open,
    }


def get_network_summary(net: pp.pandapowerNet) -> dict:
    """Return a JSON-serialisable summary of the network structure."""
    line_info = get_line_info(net)
    return {
        "n_buses": len(net.bus),
        "n_lines": len(net.line),
        "n_lines_in_service": len(line_info["in_service"]),
        "n_tie_lines": len(line_info["out_of_service"]),
        "n_generators": len(net.gen) + len(net.ext_grid),
        "n_loads": len(net.load),
        "total_load_mw": float(net.load.p_mw.sum()),
        "total_load_mvar": float(net.load.q_mvar.sum()),
        "tie_line_indices": line_info["out_of_service"],
    }


def get_topology_data(net: pp.pandapowerNet) -> list[dict]:
    """Return line data for topology visualization.

    Each entry has: index, from_bus, to_bus, r_ohm, x_ohm, in_service, is_tie.
    """
    line_info = get_line_info(net)
    tie_set = set(line_info["tie_lines"])
    rows = []
    for idx, row in net.line.iterrows():
        rows.append({
            "index": int(idx),
            "from_bus": int(row["from_bus"]),
            "to_bus": int(row["to_bus"]),
            "r_ohm_per_km": float(row["r_ohm_per_km"]),
            "x_ohm_per_km": float(row["x_ohm_per_km"]),
            "length_km": float(row["length_km"]),
            "in_service": bool(row["in_service"]),
            "is_tie": idx in tie_set,
        })
    return rows


def get_bus_data(net: pp.pandapowerNet) -> list[dict]:
    """Return bus data for visualization."""
    rows = []
    for idx, row in net.bus.iterrows():
        load_p = 0.0
        load_q = 0.0
        if len(net.load) > 0:
            bus_loads = net.load[net.load.bus == idx]
            load_p = float(bus_loads.p_mw.sum())
            load_q = float(bus_loads.q_mvar.sum())
        rows.append({
            "index": int(idx),
            "name": str(row.get("name", f"Bus {idx}")),
            "vn_kv": float(row["vn_kv"]),
            "load_mw": load_p,
            "load_mvar": load_q,
            "is_slack": int(idx) in net.ext_grid.bus.values,
        })
    return rows