"""
Three-Phase Symmetrical Fault Analysis 

Computes bolted three phase fault current and post fault bus voltages
for a fault at any bus, using the standard Zbus method

This uses the same network admittance model as the load flow study.
"""

import numpy as np


def three_phase_fault(Ybus, V_prefault, fault_bus_1idx, Zf=0.0):
    """
    fault_bus_1idx : 1 indexed bus number of the fault location
    Zf : fault impedance in pu (0 = bolted fault)
    """
    n = Ybus.shape[0]
    Zbus = np.linalg.inv(Ybus)
    k = fault_bus_1idx - 1

    Vpre_k = V_prefault[k]
    I_fault = Vpre_k / (Zbus[k, k] + Zf)

    V_fault = np.zeros(n, dtype=complex)
    for i in range(n):
        V_fault[i] = V_prefault[i] - Zbus[i, k] * I_fault

    return {
        "fault_bus": fault_bus_1idx,
        "I_fault_pu": I_fault,
        "I_fault_ka_base": None,  # u can change this if u want
        "V_during_fault_pu": np.abs(V_fault),
        "Zbus_kk": Zbus[k, k],
    }
