"""
N-1 Contingency Screening

For every branch (line or transformer) in the network, this module:
  1. Removes the branch (simulating breaker trip / outage)
  2. Re-solves the AC power flow with Newton-Raphson from the base case
     solution as the starting point (a "hot start", standard EMS practice)
  3. Detects islanding / singular Ybus conditions (a branch whose removal
     isolates part of the network)
  4. Quantifies the % redistribution of active-power flow on every
     surviving branch relative to the pre-contingency base case
  5. Flags thermal violations against branch MVA ratings and voltage
     violations against +/-5% (0.95-1.05 pu) operating limits
"""

import numpy as np
from network_setup import build_case
from newton_raphson import run_newton_raphson, branch_flows


def run_n1_screening(load_scale=1.0, v_low=0.95, v_high=1.05):
    base_case = build_case(load_scale=load_scale)
    base_res = 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"])
    base_flows = branch_flows(base_res.V, base_res.theta, base_case["branches"],
                               base_case["n_bus"])
    base_flow_map = {(f["from"], f["to"]): abs(f["S_ft_MVA"]) for f in base_flows}


    branches = base_case["branches"]
    n_bus = base_case["n_bus"]
    results = []

    for oi, (f_out, t_out, data_out, is_x) in enumerate(branches):
        case = build_case(load_scale=load_scale, outage_index=oi)

        # Detect islanding
        row_sums = np.sum(np.abs(case["Ybus"]), axis=1) - np.abs(np.diag(case["Ybus"]))
        isolated_buses = [i + 1 for i, s in enumerate(row_sums) if s < 1e-9]

        entry = {
            "outage": f"{f_out}-{t_out}", "is_xfmr": is_x,
            "isolated_buses": isolated_buses,
            "converged": False, "iterations": None,
            "max_flow_redistribution_pct": None,
            "min_bus_v": None, "max_bus_v": None,
            "voltage_violations": [], "thermal_violations": [],
        }

        if isolated_buses:
            results.append(entry)
            continue

        try:
            res = run_newton_raphson(case["Ybus"], case["bus_types"],
                                      case["P_spec"], case["Q_spec"],
                                      V0=base_res.V, theta0=base_res.theta,
                                      max_iter=40)
        except np.linalg.LinAlgError:
            results.append(entry)
            continue

        entry["converged"] = res.converged
        entry["iterations"] = res.iterations

        if res.converged:
            flows = branch_flows(res.V, res.theta, case["active_branches"], n_bus)
            redist = []
            for fl in flows:
                key = (fl["from"], fl["to"])
                base_val = base_flow_map.get(key, 0.0)
                new_val = abs(fl["S_ft_MVA"])
                if base_val > 1.0:  # ignore near-zero base flows to avoid divide blowups :)
                    pct = 100.0 * (new_val - base_val) / base_val
                    redist.append(pct)
                if fl["rate_mva"] and new_val > fl["rate_mva"]:
                    entry["thermal_violations"].append(
                        {"branch": key, "loading_mva": round(new_val, 1),
                         "rating_mva": fl["rate_mva"]})
            entry["max_flow_redistribution_pct"] = max(redist, key=abs) if redist else 0.0
            entry["min_bus_v"] = float(np.min(res.V))
            entry["max_bus_v"] = float(np.max(res.V))
            entry["voltage_violations"] = [
                {"bus": i + 1, "v_pu": round(float(res.V[i]), 4)}
                for i in range(n_bus) if res.V[i] < v_low or res.V[i] > v_high
            ]

        results.append(entry)

    return {
        "base_case_flows": base_flows,
        "base_case_result": base_res,
        "contingency_results": results,
    }