"""
Builds the bus admittance matrix (Ybus) for an arbitrary N-bus network
from line/transformer R, X, B data using the standard pi-model stamping
procedure. Supports removing a branch (for N-1 contingency screening)
without rebuilding the whole dataset by hand.
"""

import numpy as np


def build_ybus(n_bus, branches, outage_index=None):
    """
    Parameters
  
    n_bus : int
        Number of buses (buses assumed numbered 1..n_bus).
    branches : list of (from_bus, to_bus, data, is_xfmr)
        data must contain 'R', 'X', and optionally 'B' (line charging, pu).
    outage_index : int or None
        If given, the branch at this index in `branches` is excluded
        from the Ybus build (simulates an open breaker / line outage).

    Returns

    Ybus : (n_bus, n_bus) complex ndarray
    active_branches : list of branches actually stamped (post outage)
    """
    Ybus = np.zeros((n_bus, n_bus), dtype=complex)
    active_branches = []

    for idx, (f, t, data, is_xfmr) in enumerate(branches):
        if outage_index is not None and idx == outage_index:
            continue
        R = data["R"]
        X = data["X"]
        B = data.get("B", 0.0)
        z = complex(R, X)
        y_series = 1.0 / z
        y_shunt = complex(0.0, B / 2.0)

        fi, ti = f - 1, t - 1
        Ybus[fi, fi] += y_series + y_shunt
        Ybus[ti, ti] += y_series + y_shunt
        Ybus[fi, ti] -= y_series
        Ybus[ti, fi] -= y_series

        active_branches.append((f, t, data, is_xfmr))

    return Ybus, active_branches
