"""
Assembles Ybus, bus-type list, and specified P/Q injection vectors from the
raw network dataset (data/network_9bus.py), applying an optional load-
scaling factor (used for the voltage stability loading scan) and an
optional branch outage (used for N-1 contingency screening).
"""

import numpy as np
import sys, os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "data"))
sys.path.insert(0, os.path.dirname(__file__))

import network_9bus as net
from ybus import build_ybus


def build_case(load_scale=1.0, outage_index=None):
    n = net.N_BUS
    branches = net.branch_list()
    Ybus, active_branches = build_ybus(n, branches, outage_index=outage_index)

    bus_types = []
    for b in range(1, n + 1):
        bus_types.append(net.BUS_DATA[b]["type"])

    P_spec = np.zeros(n)
    Q_spec = np.zeros(n)

    for gbus, gd in net.GEN_DATA.items():
        if gd["P_mw"] is not None:
            P_spec[gbus - 1] += gd["P_mw"] / net.S_BASE_MVA

    for lbus, ld in net.LOAD_DATA.items():
        P_spec[lbus - 1] -= (ld["P_mw"] * load_scale) / net.S_BASE_MVA
        Q_spec[lbus - 1] -= (ld["Q_mvar"] * load_scale) / net.S_BASE_MVA

    V0 = np.ones(n)
    theta0 = np.zeros(n)
    for b in range(1, n + 1):
        bt = net.BUS_DATA[b]["type"]
        if bt in ("slack", "PV"):
            V0[b - 1] = net.BUS_DATA[b]["v_set"]

    return {
        "Ybus": Ybus,
        "branches": branches,
        "active_branches": active_branches,
        "bus_types": bus_types,
        "P_spec": P_spec,
        "Q_spec": Q_spec,
        "V0": V0,
        "theta0": theta0,
        "n_bus": n,
    }
