divide.py

from qiskit import *
import networkx as nx
import config
import matplotlib.pyplot as plt
import os

def draw(G):
    # define node colors based on 'M' attribute
    node_colors = ['red' if G.nodes[n].get('M') else 'lightblue' for n in G.nodes()]

    # auto layout
    pos = nx.spring_layout(G)

    # plot
    plt.figure(figsize=(8, 6))
    nx.draw(
        G, pos, with_labels=True,
        node_color=node_colors, edge_color='gray',
        node_size=2000, font_size=14, arrows=True
    )

    plt.title("DiGraph with Special Coloring for M=True Nodes")
    plt.axis('off')
    plt.savefig("graph_output.png", format="png", dpi=300, bbox_inches='tight')



def k_densest_subgraph(G, k):
    selected = set()
    degrees = dict(G.degree)

    start = max(degrees, key=degrees.get)
    selected.add(start)

    while len(selected) < k and len(selected) < len(G.nodes()):
        scores = {}
        for node in set(G.nodes()) - selected:
            score = sum(1 for n in selected if G.has_edge(node, n) or G.has_edge(n, node))
            scores[node] = score
        best = max(scores, key=scores.get)
        selected.add(best)

    return selected

def divide_graph(G, select):
    for i in list(config.qubit_for_g.keys()):
        if all(q not in select for q in config.qubit_for_g[i]):
            if i in G:
                preds = list(G.predecessors(i))
                succs = list(G.successors(i))
                for u in preds:
                    for v in succs:
                        if u != v:
                            G.add_edge(u, v)
                G.remove_node(i)

    return G

def conquer_graph(G):
    for i, j in config.a.items():  # reuse i → j
        m_i = [g for g in G.nodes if G.nodes[g].get("M") == True and config.measure_to_index[g] == i]

        for m in m_i:
            for g in config.gate_on_q.get(j, []):
                if g in G.nodes:
                    if not nx.has_path(G, m, g):
                        G.add_edge(m, g)
                        # print("add edge:", m, g)

    return G

def save_graph(G, round_num):
    pos = nx.spring_layout(G, seed=42)  # 固定 layout 結果
    node_colors = ['red' if G.nodes[n].get('M') else 'skyblue' for n in G.nodes()]
    labels = {n: f"{n}:{G.nodes[n].get('name')}" for n in G.nodes()}

    plt.figure(figsize=(14, 10))
    nx.draw_networkx_nodes(G, pos, node_color=node_colors, node_size=800)
    nx.draw_networkx_edges(G, pos, arrows=True, arrowstyle='->')
    nx.draw_networkx_labels(G, pos, labels=labels, font_size=10, verticalalignment='bottom')  # 避免 label 疊在一起

    plt.title(f"Canonical Graph - Round {round_num}")
    plt.axis('off')
    plt.tight_layout()
    os.makedirs("graph_rounds", exist_ok=True)
    plt.savefig(f"graph_rounds/canonical_round_{round_num}.png")
    plt.close()

def gate_time(G):
    # localT = [(node index1, time1), (node index2, time2), ...]
    # e.g., [(0, 0), (1, 5), (2, 1), (3, 2), (4, 4), (5, 3), (6, 16), (8, 6), (10, 7), (14, 8)]
    if nx.is_directed_acyclic_graph(G):
        topo_order = list(nx.topological_sort(G))
        # print("topo order:", topo_order)


    return topo_order

def remove_measurements_and_resets(qc):
    new_qc = QuantumCircuit(*qc.qregs, *qc.cregs)
    for instr, qargs, cargs in qc.data:
        if instr.name not in ["measure", "reset"]:
            new_qc.append(instr, qargs, cargs)
    return new_qc

def rm(qc):
    new_qc = QuantumCircuit(*qc.qregs, *qc.cregs)

    pseudo_clbit_met = [False] * qc.qregs[1].size # THIS
    
    for instr, qargs, cargs in qc.data:
        if instr.name not in ['measure', 'reset']:
            flag = False
            for qubit in qargs:
                if qubit.register == qc.qregs[1] and pseudo_clbit_met[qubit.index] == False:
                    pseudo_clbit_met[qubit.index] = True
                    flag = True
            if not flag:
                new_qc.append(instr, qargs, cargs)
    return new_qc