max

import numpy as np
from scipy.optimize import linprog
import networkx as nx
import matplotlib.pyplot as plt


def max_flow(graph, source, sink):
    # Создаем матрицу инцидентности
    num_nodes = len(graph)
    num_edges = sum(len(edges) for edges in graph.values())
    A_eq = np.zeros((num_nodes + num_edges, num_edges))
    b_eq = np.zeros(num_nodes + num_edges)
    c = np.zeros(num_edges)
    edge_index = 0

    for i, node in enumerate(graph):
        for neighbor, capacity in graph[node].items():
            if node != source and node != sink:
                A_eq[i, edge_index] = 1
                A_eq[num_nodes + edge_index, edge_index] = -1
                b_eq[i] = 0
                c[edge_index] = -1
                edge_index += 1
            if node == source:
                A_eq[num_nodes + edge_index, edge_index] = -1
                b_eq[num_nodes + edge_index] = -capacity
                c[edge_index] = -1
                edge_index += 1
            elif node == sink:
                A_eq[num_nodes + edge_index, edge_index] = 1
                b_eq[num_nodes + edge_index] = capacity
                c[edge_index] = -1
                edge_index += 1

    # Решаем задачу линейного программирования
    res = linprog(c, A_eq=A_eq, b_eq=b_eq)

    # Создаем словарь потока
    flow = {}
    for node in graph:
        flow[node] = {}

    # Заполняем значениями потока
    edge_index = 0
    for i, node in enumerate(graph):
        for neighbor, _ in graph[node].items():
            flow[node][neighbor] = res.x[edge_index]
            edge_index += 1

    return flow

def visualize_flow(graph, flow):
    G = nx.DiGraph()
    for node, neighbors in graph.items():
        for neighbor, capacity in neighbors.items():
            G.add_edge(node, neighbor, capacity=capacity)

    pos = nx.spring_layout(G)

    plt.figure(figsize=(10, 5))

    # Рисуем исходный граф
    plt.subplot(1, 2, 1)
    nx.draw(G, pos, with_labels=True, node_color='lightblue', node_size=2000, font_size=10)
    plt.title('Original Graph')

    # Рисуем граф с потоком
    plt.subplot(1, 2, 2)
    for node, neighbors in graph.items():
        for neighbor, capacity in neighbors.items():
            flow_value = flow[node][neighbor]
            nx.draw_networkx_edge_labels(G, pos, edge_labels={(node, neighbor): str(flow_value)}, label_pos=0.3)

    nx.draw(G, pos, with_labels=True, node_color='lightblue', node_size=2000, font_size=10)
    plt.title('Graph with Max Flow')

    plt.show()


# Пример графа для тестирования
graph = {
    'A': {'B': 10, 'C': 5},
    'B': {'D': 15},
    'C': {'D': 10},
    'D': {}
}

source = 'A'
sink = 'D'

# Рассчитываем максимальный поток
flow = max_flow(graph, source, sink)

# Создаем словарь потока для визуализации
flow_dict = {}
nodes = list(graph.keys())
edge_index = 0
for i, node in enumerate(nodes):
    flow_dict[node] = {}
    for neighbor, _ in graph[node].items():
        flow_dict[node][neighbor] = flow[node][neighbor]

# Визуализируем поток
visualize_flow(graph, flow_dict)