"""
Automated tests for the power-flow analysis toolchain.

Run with:  python3 -m pytest tests/test_power_flow.py -v
       or:  python3 tests/test_power_flow.py   (falls back to plain asserts)
"""
import sys, os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "data"))

import numpy as np
from network_setup import build_case
from newton_raphson import run_newton_raphson, branch_flows
from gauss_seidel import run_gauss_seidel
from fault_analysis import three_phase_fault


def test_base_case_converges():
    case = build_case(load_scale=1.0)
    res = run_newton_raphson(case["Ybus"], case["bus_types"], case["P_spec"], case["Q_spec"],
                              V0=case["V0"], theta0=case["theta0"])
    assert res.converged, "Base case must converge"
    assert res.iterations <= 15, f"Base case took unexpectedly many iterations: {res.iterations}"


def test_base_case_voltages_within_normal_band():
    case = build_case(load_scale=1.0)
    res = run_newton_raphson(case["Ybus"], case["bus_types"], case["P_spec"], case["Q_spec"],
                              V0=case["V0"], theta0=case["theta0"])
    assert np.all(res.V > 0.9) and np.all(res.V < 1.1), "Base-case voltages should be near 1.0 pu"


def test_power_balance():
    """Total generation should equal total load plus losses (within numerical tolerance)."""
    case = build_case(load_scale=1.0)
    res = run_newton_raphson(case["Ybus"], case["bus_types"], case["P_spec"], case["Q_spec"],
                              V0=case["V0"], theta0=case["theta0"])
    flows = branch_flows(res.V, res.theta, case["branches"], case["n_bus"])
    total_loss_mw = sum(f["loss_MW"] for f in flows)
    assert total_loss_mw > 0, "Transmission losses must be positive"
    assert total_loss_mw < 20, "Losses implausibly high for this network/loading"


def test_newton_raphson_matches_gauss_seidel():
    """Two independently formulated solvers must agree to within 2% (validation threshold)."""
    case = build_case(load_scale=1.0)
    res_nr = run_newton_raphson(case["Ybus"], case["bus_types"], case["P_spec"], case["Q_spec"],
                                 V0=case["V0"], theta0=case["theta0"])
    V0c = case["V0"] * np.exp(1j * case["theta0"])
    res_gs = run_gauss_seidel(case["Ybus"], case["bus_types"], case["P_spec"], case["Q_spec"],
                               V0=V0c, tol=1e-8, max_iter=3000)
    assert res_gs["converged"], "Gauss-Seidel cross-check must converge"
    max_dev_pct = float(np.max(100 * np.abs(res_nr.V - res_gs["V"]) / res_nr.V))
    assert max_dev_pct < 2.0, f"NR/GS deviation {max_dev_pct}% exceeds 2% validation threshold"


def test_gsu_transformer_outage_causes_islanding():
    """Each generator's only link to the network is its own GSU transformer; losing it must island the bus."""
    for outage_index in [6, 7, 8]:  # the three transformer entries in branch_list()
        case = build_case(load_scale=1.0, outage_index=outage_index)
        row_sums = np.sum(np.abs(case["Ybus"]), axis=1) - np.abs(np.diag(case["Ybus"]))
        isolated = [i + 1 for i, s in enumerate(row_sums) if s < 1e-9]
        assert len(isolated) == 1, f"Expected exactly one isolated bus for outage index {outage_index}"


def test_line_outages_all_converge():
    for outage_index in range(6):  # the six transmission lines
        case = build_case(load_scale=1.0, outage_index=outage_index)
        res = run_newton_raphson(case["Ybus"], case["bus_types"], case["P_spec"], case["Q_spec"],
                                  V0=case["V0"], theta0=case["theta0"], max_iter=40)
        assert res.converged, f"Line outage index {outage_index} should converge at nominal load"


def test_meshed_topology_has_lower_losses_than_any_radial_option():
    base_case = build_case(load_scale=1.0)
    res_base = run_newton_raphson(base_case["Ybus"], base_case["bus_types"], base_case["P_spec"],
                                   base_case["Q_spec"], V0=base_case["V0"], theta0=base_case["theta0"])
    flows_base = branch_flows(res_base.V, res_base.theta, base_case["branches"], base_case["n_bus"])
    loss_meshed = sum(f["loss_MW"] for f in flows_base)

    for oi in range(6):
        case = build_case(load_scale=1.0, outage_index=oi)
        res = run_newton_raphson(case["Ybus"], case["bus_types"], case["P_spec"], case["Q_spec"],
                                  V0=res_base.V, theta0=res_base.theta, max_iter=40)
        if not res.converged:
            continue
        flows = branch_flows(res.V, res.theta, case["active_branches"], case["n_bus"])
        loss_radial = sum(f["loss_MW"] for f in flows)
        assert loss_radial > loss_meshed, "Radial reconfiguration should never reduce losses vs. meshed operation"


def test_fault_current_positive_and_voltage_collapses_at_fault_bus():
    case = build_case(load_scale=1.0)
    res = run_newton_raphson(case["Ybus"], case["bus_types"], case["P_spec"], case["Q_spec"],
                              V0=case["V0"], theta0=case["theta0"])
    Vc = res.V * np.exp(1j * res.theta)
    fr = three_phase_fault(case["Ybus"], Vc, fault_bus_1idx=7)
    assert abs(fr["I_fault_pu"]) > 0
    assert fr["V_during_fault_pu"][7 - 1] < 1e-6, "Faulted bus voltage should collapse to ~0"


def test_increasing_load_monotonically_reduces_voltage_at_load_buses():
    base_case = build_case(load_scale=1.0)
    prev_v = None
    for scale in [1.0, 1.3, 1.6, 1.9]:
        case = build_case(load_scale=scale)
        res = run_newton_raphson(case["Ybus"], case["bus_types"], case["P_spec"], case["Q_spec"],
                                  V0=base_case["V0"], theta0=base_case["theta0"], max_iter=60)
        assert res.converged
        v5 = res.V[4]  # bus 5, 0-indexed
        if prev_v is not None:
            assert v5 < prev_v, "Bus 5 voltage should decline monotonically as demand increases"
        prev_v = v5


if __name__ == "__main__":
    tests = [obj for name, obj in list(globals().items()) if name.startswith("test_")]
    passed, failed = 0, 0
    for t in tests:
        try:
            t()
            print(f"PASS: {t.__name__}")
            passed += 1
        except AssertionError as e:
            print(f"FAIL: {t.__name__}: {e}")
            failed += 1
    print(f"\n{passed} passed, {failed} failed out of {len(tests)} tests")
    sys.exit(1 if failed else 0)
