"""
runs the full analysis and writes results/figures 

Sections
1. Base case Newton-Raphson power flow + convergence trace
2. Cross validation against independent Gauss Seidel solver
3. N-1 contingency screening (6 lines + 3 GSU transformers)
4. Meshed ring vs radial reconfiguration loss/voltage comparison
5. Voltage stability / increased demand loading scan (P-V behaviour)
6. Transformer thermal loading scan under increasing demand
7. Three phase symmetrical fault analysis at representative buses
"""

import json
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import pandapower as pp

from network_setup import build_case
from newton_raphson import run_newton_raphson, branch_flows
from gauss_seidel import run_gauss_seidel
from contingency import run_n1_screening
from fault_analysis import three_phase_fault

RESULTS_DIR = "results"
FIG_DIR = "figures"

plt.rcParams.update({
    "figure.facecolor": "white",
    "axes.facecolor": "white",
    "font.size": 11,
    "axes.grid": True,
    "grid.alpha": 0.3,
})
COLOR_PRIMARY = "#1f4e79"
COLOR_ACCENT = "#c55a11"
COLOR_WARN = "#c00000"
COLOR_OK = "#548235"
# The colours are used for the diagrams

def section(title):
    print("\n" + "=" * 70)
    print(title)
    print("=" * 70)


# ---------------------------------------------------------------------
# 1. Base case
section("1. BASE CASE NEWTON-RAPHSON POWER FLOW")

base = build_case(load_scale=1.0)
res = run_newton_raphson(base["Ybus"], base["bus_types"], base["P_spec"], base["Q_spec"],
                          V0=base["V0"], theta0=base["theta0"], tol=1e-6, max_iter=25)
flows = branch_flows(res.V, res.theta, base["branches"], base["n_bus"])
total_loss = sum(f["loss_MW"] for f in flows)
total_gen = sum(base["P_spec"][i] for i in range(9) if base["P_spec"][i] > 0) * 100
# recompute exact slack generation from injection
P_inj, Q_inj = np.zeros(9), np.zeros(9)

print(f"Converged: {res.converged} in {res.iterations} iterations "
      f"(flat start informed initial guess; generator buses at scheduled setpoints)")
print("Bus voltages (pu) / angles (deg):")
for i in range(9):
    print(f"  Bus {i+1}: V={res.V[i]:.4f} pu, angle={np.degrees(res.theta[i]):+.3f} deg  "
          f"[{base['bus_types'][i]}]")
print(f"Total real-power transmission losses: {total_loss:.2f} MW")

base_case_table = [
    {"bus": i + 1, "type": base["bus_types"][i], "V_pu": round(float(res.V[i]), 4),
     "angle_deg": round(float(np.degrees(res.theta[i])), 3)}
    for i in range(9)
]

# Also solve fully flat start (V=1, theta=0 everywhere) for comparison
V0_flat, th0_flat = np.ones(9), np.zeros(9)
res_flat = run_newton_raphson(base["Ybus"], base["bus_types"], base["P_spec"], base["Q_spec"],
                               V0=V0_flat, theta0=th0_flat, tol=1e-8, max_iter=25)
print(f"\n[Fully flat start V=1,theta=0, tol=1e-8] converged in {res_flat.iterations} iterations")

# ---------------------------------------------------------------------
# 2. Validation: pandapower + Gauss-Seidel 
section("2. VALIDATION: PANDAPOWER + GAUSS-SEIDEL")

from pandapower_validation import main as run_pandapower_validation
run_pandapower_validation()

V0_complex = base["V0"] * np.exp(1j * base["theta0"])
gs = run_gauss_seidel(base["Ybus"], base["bus_types"], base["P_spec"], base["Q_spec"],
                       V0=V0_complex, tol=1e-8, max_iter=3000)
deviation_pct = 100 * np.abs(res.V - gs["V"]) / res.V
max_dev = float(np.max(deviation_pct))
print(f"\n[Secondary cross check] Gauss-Seidel converged: {gs['converged']} in {gs['iterations']} iterations")
print(f"Max |V_NR - V_GS| / V_NR across all buses: {max_dev:.5f} %  (validation threshold: <2%)")

validation_table = [
    {"bus": i + 1, "V_NR_pu": round(float(res.V[i]), 5), "V_GS_pu": round(float(gs["V"][i]), 5),
     "deviation_pct": round(float(deviation_pct[i]), 6)}
    for i in range(9)
]

# ---------------------------------------------------------------------
# 3. N-1 contingency screening
section("3. N-1 CONTINGENCY SCREENING")

n1 = run_n1_screening(load_scale=1.0)
n1_table = []
for r in n1["contingency_results"]:
    n1_table.append({
        "outage": r["outage"], "type": "Transformer" if r["is_xfmr"] else "Line",
        "isolated_buses": r["isolated_buses"], "converged": r["converged"],
        "iterations": r["iterations"],
        "max_flow_redistribution_pct": (None if r["max_flow_redistribution_pct"] is None
                                         else round(r["max_flow_redistribution_pct"], 1)),
        "min_bus_v_pu": (None if r["min_bus_v"] is None else round(r["min_bus_v"], 4)),
        "voltage_violations": r["voltage_violations"],
        "thermal_violations": r["thermal_violations"],
    })
    print(json.dumps(n1_table[-1]))

# ---------------------------------------------------------------------
# 4. Meshed vs radial reconfiguration
section("4. MESHED vs RADIAL RECONFIGURATION")

radial_table = [{"config": "Meshed (ring closed, base case)", "losses_MW": round(total_loss, 2),
                  "delta_MW": 0.0, "v_spread_pu": round(float(res.V.max() - res.V.min()), 4),
                  "min_v_pu": round(float(res.V.min()), 4)}]
for oi in range(6):  # the 6 ring transmission lines
    f_out, t_out, d, isx = base["branches"][oi]
    case = build_case(load_scale=1.0, outage_index=oi)
    r = run_newton_raphson(case["Ybus"], case["bus_types"], case["P_spec"], case["Q_spec"],
                            V0=res.V, theta0=res.theta, max_iter=40)
    if not r.converged:
        continue
    fl = branch_flows(r.V, r.theta, case["active_branches"], case["n_bus"])
    loss = sum(x["loss_MW"] for x in fl)
    radial_table.append({
        "config": f"Radial (open {f_out}-{t_out})", "losses_MW": round(loss, 2),
        "delta_MW": round(loss - total_loss, 2),
        "v_spread_pu": round(float(r.V.max() - r.V.min()), 4),
        "min_v_pu": round(float(r.V.min()), 4),
    })
    print(radial_table[-1])

# ---------------------------------------------------------------------
# 5. Voltage stability / increased demand loading scan
section("5. VOLTAGE STABILITY LOADING SCAN (INCREASING DEMAND)")

load_scales = np.round(np.arange(1.0, 2.25, 0.05), 2)
load_buses = [5, 6, 8]
V_base_ls = res.V.copy()
loading_scan = []
xfmr_loading_scan = []
max_converged_scale = 1.0

for scale in load_scales:
    case = build_case(load_scale=float(scale))
    r = run_newton_raphson(case["Ybus"], case["bus_types"], case["P_spec"], case["Q_spec"],
                            V0=base["V0"], theta0=base["theta0"], tol=1e-8, max_iter=60)
    if not r.converged:
        loading_scan.append({"load_scale": float(scale), "converged": False})
        continue
    max_converged_scale = float(scale)
    row = {"load_scale": float(scale), "converged": True, "iterations": r.iterations}
    for b in load_buses:
        drop_pct = 100 * (V_base_ls[b - 1] - r.V[b - 1]) / V_base_ls[b - 1]
        row[f"bus{b}_V_pu"] = round(float(r.V[b - 1]), 4)
        row[f"bus{b}_drop_pct"] = round(float(drop_pct), 2)
    loading_scan.append(row)

    fl = branch_flows(r.V, r.theta, case["branches"], case["n_bus"])
    xrow = {"load_scale": float(scale)}
    for f in fl:
        if f["is_xfmr"]:
            pct_loaded = 100 * f["loading_MVA"] / f["rate_mva"]
            xrow[f"{f['from']}-{f['to']}_pct_rating"] = round(float(pct_loaded), 1)
    xfmr_loading_scan.append(xrow)

print(f"Maximum converged load scale reached in scan: {max_converged_scale:.2f}x nominal "
      f"(power flow solution fails to converge beyond this point => approaching the "
      f"system's maximum loadability / voltage collapse boundary)")
for row in loading_scan:
    print(row)

# ---------------------------------------------------------------------
# 6. Transformer loading table (peak of scan)
section("6. TRANSFORMER THERMAL LOADING (BASE CASE AND UNDER INCREASED DEMAND)")
for xrow in xfmr_loading_scan:
    print(xrow)

# ---------------------------------------------------------------------
# 7. Three-phase fault analysis
section("7. THREE PHASE SYMMETRICAL FAULT ANALYSIS")

kv_map = {1: 16.5, 2: 18.0, 3: 13.8, 4: 230, 5: 230, 6: 230, 7: 230, 8: 230, 9: 230}
S_BASE = 100.0
Vc = res.V * np.exp(1j * res.theta)
fault_table = []
for bus in [1, 4, 5, 6, 7, 8, 9]:
    fr = three_phase_fault(base["Ybus"], Vc, bus)
    Ibase_kA = S_BASE / (np.sqrt(3) * kv_map[bus])
    Ifault_kA = abs(fr["I_fault_pu"]) * Ibase_kA
    fault_table.append({
        "fault_bus": bus, "I_fault_pu": round(float(abs(fr["I_fault_pu"])), 3),
        "I_fault_kA": round(float(Ifault_kA), 2),
        "min_V_during_fault_pu": round(float(np.min(fr["V_during_fault_pu"])), 3),
    })
    print(fault_table[-1])

# =======================================================================
# SAVE RESULTS
section("SAVING RESULTS")

all_results = {
    "base_case": {
        "converged": res.converged, "iterations": res.iterations,
        "mismatch_history": [float(m) for m in res.mismatch_history],
        "flat_start_iterations": res_flat.iterations,
        "total_losses_MW": round(total_loss, 2),
        "bus_results": base_case_table,
        "branch_flows": [
            {"from": f["from"], "to": f["to"], "is_xfmr": f["is_xfmr"],
             "loading_MVA": round(float(f["loading_MVA"]), 2),
             "rating_MVA": f["rate_mva"], "loss_MW": round(float(f["loss_MW"]), 3)}
            for f in flows
        ],
    },
    "validation_pandapower": {
        "method": "primary",
        "script": "src/pandapower_validation.py",
        "implemented": True,
    },
    "validation_vs_gauss_seidel": {
        "method": "secondary",
        "gs_converged": gs["converged"], "gs_iterations": gs["iterations"],
        "max_deviation_pct": round(max_dev, 6), "threshold_pct": 2.0,
        "bus_table": validation_table,
    },
    "n1_contingency": n1_table,
    "radial_reconfiguration": radial_table,
    "voltage_stability_scan": loading_scan,
    "transformer_loading_scan": xfmr_loading_scan,
    "fault_analysis": fault_table,
}

with open(f"{RESULTS_DIR}/results.json", "w") as fh:
    json.dump(all_results, fh, indent=2)
print(f"Saved {RESULTS_DIR}/results.json")

# =======================================================================
# FIGURES
section("GENERATING FIGURES")

# Fig 1: base-case voltage profile
fig, ax = plt.subplots(figsize=(7, 4.2))
buses = list(range(1, 10))
colors = [COLOR_PRIMARY if base["bus_types"][b - 1] != "PQ" else COLOR_ACCENT for b in buses]
ax.bar(buses, res.V, color=colors, width=0.6, zorder=3)
ax.axhline(1.05, color=COLOR_WARN, ls="--", lw=1, label="±5% operating band")
ax.axhline(0.95, color=COLOR_WARN, ls="--", lw=1)
ax.set_ylim(0.9, 1.08)
ax.set_xlabel("Bus")
ax.set_ylabel("Voltage magnitude (pu)")
ax.set_title("Base-Case Bus Voltage Profile (Newton-Raphson)")
ax.set_xticks(buses)
ax.legend(loc="lower right", fontsize=9)
fig.tight_layout()
fig.savefig(f"{FIG_DIR}/fig1_base_voltage_profile.png", dpi=170)
plt.close(fig)

# Fig 2: NR convergence (mismatch vs iteration), base + flat start
fig, ax = plt.subplots(figsize=(7, 4.2))
ax.semilogy(range(1, len(res.mismatch_history) + 1), res.mismatch_history,
            "o-", color=COLOR_PRIMARY, label="Generator-informed start (base case)")
ax.semilogy(range(1, len(res_flat.mismatch_history) + 1), res_flat.mismatch_history,
            "s--", color=COLOR_ACCENT, label="Fully flat start (V=1.0, θ=0)")
ax.set_xlabel("Newton-Raphson iteration")
ax.set_ylabel("Max active/reactive mismatch (pu, log scale)")
ax.set_title("Newton-Raphson Convergence History")
ax.legend(fontsize=9)
fig.tight_layout()
fig.savefig(f"{FIG_DIR}/fig2_nr_convergence.png", dpi=170)
plt.close(fig)

# Fig 3: N-1 contingency flow redistribution
fig, ax = plt.subplots(figsize=(8, 4.5))
labels, values, bar_colors = [], [], []
for r in n1_table:
    labels.append(r["outage"] + ("\n(xfmr)" if r["type"] == "Transformer" else ""))
    if r["isolated_buses"]:
        values.append(0)
        bar_colors.append("#999999")
    else:
        values.append(r["max_flow_redistribution_pct"] or 0)
        bar_colors.append(COLOR_WARN if abs(r["max_flow_redistribution_pct"] or 0) > 100 else COLOR_PRIMARY)
bars = ax.bar(labels, values, color=bar_colors, zorder=3)
for i, r in enumerate(n1_table):
    if r["isolated_buses"]:
        ax.text(i, 5, "islanding", ha="center", rotation=90, fontsize=8, color="#555")
ax.axhline(0, color="black", lw=0.8)
ax.set_ylabel("Max branch flow redistribution (%)")
ax.set_title("N-1 Contingency Screening: Post-Outage Flow Redistribution")
fig.tight_layout()
fig.savefig(f"{FIG_DIR}/fig3_n1_redistribution.png", dpi=170)
plt.close(fig)

# Fig 4: Voltage-stability (P-V style) curves for heavily loaded buses
fig, ax = plt.subplots(figsize=(7.5, 4.5))
converged_scan = [r for r in loading_scan if r.get("converged")]
scales = [r["load_scale"] for r in converged_scan]
for b, c in zip(load_buses, [COLOR_PRIMARY, COLOR_ACCENT, COLOR_OK]):
    ax.plot(scales, [r[f"bus{b}_V_pu"] for r in converged_scan], "o-", color=c,
            label=f"Bus {b}", markersize=4)
ax.axvspan(1.6, 1.8, color=COLOR_WARN, alpha=0.08, label="5-10% drop band")
ax.set_xlabel("System demand (multiple of nominal load)")
ax.set_ylabel("Bus voltage (pu)")
ax.set_title("Voltage Sensitivity at Heavily Loaded Buses vs Increasing Demand")
ax.legend(fontsize=9)
fig.tight_layout()
fig.savefig(f"{FIG_DIR}/fig4_voltage_stability_scan.png", dpi=170)
plt.close(fig)

# Fig 5: Transformer loading vs demand
fig, ax = plt.subplots(figsize=(7.5, 4.5))
xscales = [r["load_scale"] for r in xfmr_loading_scan]
for key, c, lbl in [("1-4_pct_rating", COLOR_PRIMARY, "T1 (Gen1 GSU, 250 MVA)"),
                     ("2-7_pct_rating", COLOR_ACCENT, "T2 (Gen2 GSU, 200 MVA)"),
                     ("3-9_pct_rating", COLOR_OK, "T3 (Gen3 GSU, 100 MVA)")]:
    ax.plot(xscales, [r.get(key) for r in xfmr_loading_scan], "o-", color=c, label=lbl, markersize=4)
ax.axhline(100, color=COLOR_WARN, ls="--", lw=1.2, label="Thermal rating (100%)")
ax.set_xlabel("System demand (multiple of nominal load)")
ax.set_ylabel("Transformer loading (% of MVA rating)")
ax.set_title("Generator Step-Up Transformer Loading vs Increasing Demand")
ax.legend(fontsize=8.5)
fig.tight_layout()
fig.savefig(f"{FIG_DIR}/fig5_transformer_loading.png", dpi=170)
plt.close(fig)

# Fig 6: Meshed vs radial losses
fig, ax = plt.subplots(figsize=(8, 4.5))
configs = [r["config"].replace("Radial (open ", "Open ").replace(")", "") for r in radial_table]
losses = [r["losses_MW"] for r in radial_table]
bar_colors = [COLOR_OK] + [COLOR_ACCENT] * (len(radial_table) - 1)
ax.bar(configs, losses, color=bar_colors, zorder=3)
ax.set_ylabel("Total real-power transmission losses (MW)")
ax.set_title("System Losses: Meshed (Ring) Operation vs Radial Reconfiguration")
plt.setp(ax.get_xticklabels(), rotation=30, ha="right")
fig.tight_layout()
fig.savefig(f"{FIG_DIR}/fig6_meshed_vs_radial_losses.png", dpi=170)
plt.close(fig)

# Fig 7: Fault current by bus
fig, ax = plt.subplots(figsize=(7, 4.2))
fbuses = [f["fault_bus"] for f in fault_table]
fcur = [f["I_fault_kA"] for f in fault_table]
ax.bar([str(b) for b in fbuses], fcur, color=COLOR_PRIMARY, zorder=3)
ax.set_xlabel("Faulted bus")
ax.set_ylabel("Three-phase fault current (kA)")
ax.set_title("Bolted Three-Phase Fault Current by Bus Location")
fig.tight_layout()
fig.savefig(f"{FIG_DIR}/fig7_fault_currents.png", dpi=170)
plt.close(fig)

print("All figures saved to", FIG_DIR)