import matplotlib.pyplot as plt
import networkx as nx

# Создаём граф
G = nx.DiGraph()

# Добавляем узлы
G.add_node('Источник\n110 кВ')
G.add_node('Узел 1\nНагрузка\n17.5 + j6 МВА')
G.add_node('Узел 2\nНагрузка\n11.2 + j3 МВА')

# Добавляем рёбра с параметрами линий
G.add_edge('Источник\n110 кВ', 'Узел 1\nНагрузка\n17.5 + j6 МВА', label='Z01 = 3.06 + j7.54 Ω')
G.add_edge('Узел 1\nНагрузка\n17.5 + j6 МВА', 'Узел 2\nНагрузка\n11.2 + j3 МВА', label='Z12 = 4.875 + j9.875 Ω')

pos = {
    'Источник\n110 кВ': (0, 0),
    'Узел 1\nНагрузка\n17.5 + j6 МВА': (0, -1),
    'Узел 2\nНагрузка\n11.2 + j3 МВА': (0, -2)
}

plt.figure(figsize=(6,4))
nx.draw(G, pos, with_labels=True, node_size=3000, node_color='lightblue', font_size=9, font_weight='bold', arrowsize=20)
edge_labels = nx.get_edge_attributes(G, 'label')
nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_color='red', font_size=8)

plt.axis('off')
plt.title('Схема замещения электрической сети')
plt.tight_layout()
plt.savefig('schem_replacement.png', dpi=300)
plt.show()