"""
Generates a schematic one-line diagram of the WSCC 9-bus test network
topology (buses, generators, transformers, transmission lines).
"""
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches

NAVY = "#1f4e79"
ACCENT = "#c55a11"
STEEL = "#7FA8C9"
GRAY = "#595959"

# Bus positions (schematic, not geographic)
pos = {
    1: (0.0, 1.0), 4: (1.6, 1.0),
    5: (3.2, 1.9), 6: (3.2, 0.1),
    7: (4.8, 1.9), 9: (4.8, 0.1),
    8: (6.4, 1.0),
    2: (4.8, 3.1), 3: (4.8, -1.1),
}
gens = [1, 2, 3]
loads = [5, 6, 8]
junctions = [4, 7, 9]

lines = [(4, 5), (4, 6), (5, 7), (6, 9), (7, 8), (8, 9)]
xfmrs = [(1, 4), (2, 7), (3, 9)]

fig, ax = plt.subplots(figsize=(7.0, 4.84))

for f, t in lines:
    x1, y1 = pos[f]; x2, y2 = pos[t]
    ax.plot([x1, x2], [y1, y2], color=NAVY, lw=2.0, zorder=1)
for f, t in xfmrs:
    x1, y1 = pos[f]; x2, y2 = pos[t]
    ax.plot([x1, x2], [y1, y2], color=ACCENT, lw=2.0, ls=(0, (4, 2)), zorder=1)
    mx, my = (x1 + x2) / 2, (y1 + y2) / 2
    circ = plt.Circle((mx, my), 0.11, facecolor="white", edgecolor=ACCENT, lw=1.6, zorder=2)
    ax.add_patch(circ)

for b, (x, y) in pos.items():
    if b in gens:
        c = plt.Circle((x, y), 0.28, facecolor=NAVY, edgecolor=NAVY, zorder=3)
        ax.add_patch(c)
        ax.text(x, y, f"G{b}", ha="center", va="center", color="white",
                 fontsize=10, fontweight="bold", zorder=4)
    else:
        face = ACCENT if b in loads else "white"
        edge = ACCENT if b in loads else GRAY
        c = plt.Circle((x, y), 0.20, facecolor=face, edgecolor=edge, lw=1.8, zorder=3)
        ax.add_patch(c)
        txt_color = "white" if b in loads else GRAY
        ax.text(x, y, str(b), ha="center", va="center", color=txt_color,
                 fontsize=9.5, fontweight="bold", zorder=4)
    label_y = y + 0.42 if b not in (3,) else y - 0.42
    if b in loads:
        ax.text(x, y - 0.42, "load", ha="center", va="top", fontsize=7.5, color=GRAY, style="italic")

ax.set_xlim(-0.6, 7.0)
ax.set_ylim(-1.7, 3.7)
ax.set_aspect("equal")
ax.axis("off")

leg = [
    mpatches.Patch(color=NAVY, label="Generator bus (slack / PV)"),
    mpatches.Patch(color=ACCENT, label="Load bus (PQ)"),
    mpatches.Patch(facecolor="white", edgecolor=GRAY, label="Junction bus (PQ)"),
]
ax.legend(handles=leg, loc="lower center", ncol=3, frameon=False, fontsize=8.5,
          bbox_to_anchor=(0.5, -0.06))
ax.text(-0.5, 3.5, "— Line     ---- Transformer", fontsize=8, color=GRAY)

fig.tight_layout()
fig.savefig("/home/claude/pf_project/figures/network_diagram.png", dpi=270,
            facecolor="white")
print("Saved network_diagram.png")

from PIL import Image
im = Image.open("/home/claude/pf_project/figures/network_diagram.png")
print("Final size:", im.size, "aspect:", im.size[0] / im.size[1])
