# Import libraries
import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
from qiskit import QuantumCircuit, Aer, execute
import warnings

# Step 1: Define the problem and setup
num_vertices = 20  # Number of vertices in the graph
edge_probability = 0.5  # Probability of creating an edge between any two vertices
shots = 1000  # Number of shots for quantum circuit execution

# Step 2: Generate a random graph using NetworkX
def generate_random_graph(num_vertices, edge_probability):
    """
    Generate a random graph using NetworkX.

    Parameters:
        num_vertices (int): Number of vertices in the graph.
        edge_probability (float): Probability of creating an edge between any two vertices.

    Returns:
        np.array: Adjacency matrix of the generated graph.
    """
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        graph = nx.erdos_renyi_graph(num_vertices, edge_probability)
        return nx.adjacency_matrix(graph).toarray()

# Step 3: Implement quantum walk on graph
def quantum_walk(graph):
    """
    Implement a quantum walk on a graph.

    Parameters:
        graph (np.array): Adjacency matrix representing the graph.

    Returns:
        dict: Dictionary containing counts of measurement outcomes.
    """
    num_vertices = len(graph)
    qc = QuantumCircuit(num_vertices)
    
    for i in range(num_vertices):
        qc.h(i)
    
    for i in range(num_vertices):
        for j in range(i+1, num_vertices):
            if graph[i][j] == 1:
                qc.cp(np.pi, i, j)
    
    qc.measure_all()
    
    simulator = Aer.get_backend('qasm_simulator')
    result = execute(qc, simulator, shots=shots).result()
    counts = result.get_counts(qc)
    
    return counts

# Step 4: Quantum Approximation Algorithm (Max-Cut)
def max_cut_approximation(graph, counts):
    """
    Implement the Quantum Approximation Algorithm (Max-Cut).

    Parameters:
        graph (np.array): Adjacency matrix representing the graph.
        counts (dict): Dictionary containing counts of measurement outcomes.

    Returns:
        tuple: Tuple containing the max cut size and the partition string.
    """
    max_cut = 0
    partition = None
    
    for partition_string in counts.keys():
        cut = 0
        for i in range(len(partition_string)):
            for j in range(i+1, len(partition_string)):
                if partition_string[i] != partition_string[j] and graph[i][j] == 1:
                    cut += 1
        if cut > max_cut:
            max_cut = cut
            partition = partition_string
    
    return max_cut, partition

# Step 5: Decode the partitioning and visualize the resulting graph with cut
def decode_partition(graph, partition):
    """
    Decode the partitioning and visualize the resulting graph with the cut.

    Parameters:
        graph (np.array): Adjacency matrix representing the graph.
        partition (str): String representing the partition of the vertices.

    Returns:
        None
    """
    G = nx.from_numpy_array(graph)
    colors = ['skyblue', 'salmon']
    node_colors = [colors[int(partition[node])] for node in G.nodes()]
    
    # Identify edges crossing the partition
    cut_edges = [(u, v) for u, v in G.edges() if partition[u] != partition[v]]
    
    # Draw graph with nodes colored based on partition and cut edges highlighted
    pos = nx.spring_layout(G)  # Positions for all nodes
    nx.draw(G, pos, with_labels=True, node_color=node_colors, node_size=1000, font_size=12)
    nx.draw_networkx_edges(G, pos, edgelist=cut_edges, width=2, alpha=0.5, edge_color='red')  # Draw cut edges
    plt.title("Resulted Partitioned Graph with Cut")
    plt.show()

# Execute the steps
# Generate random graph
graph = generate_random_graph(num_vertices, edge_probability)

# Implement quantum walk
counts = quantum_walk(graph)

# Find max-cut approximation
max_cut, partition = max_cut_approximation(graph, counts)
print("Max-Cut:", max_cut)
print(f"The value {max_cut} indicates the maximum number of edges that can be cut by partitioning the vertices into two sets.")
print(f"In other words, if we divide the vertices into two groups according to the partition obtained, we can cut a maximum of {max_cut} edges between these two groups.")
print()
print("Partition:", partition)
print(f"The binary string {partition} represents the partitioning of vertices.")
print(f"Each digit in the string corresponds to a vertex in the graph.")
print(f"In this case, there are {num_vertices} vertices, so the binary string has {num_vertices} digits.")
print("The value '1' indicates that the corresponding vertex belongs to one set (Partition 1), and the value '0' indicates that the vertex belongs to the other set (Partition 0).")


# Decode partition and visualize resulting graph
decode_partition(graph, partition)

