"""
Utility functions for converting quantum circuits to graph representations and
vice versa, as well as functions for updating node and edge information in the graph.
"""

import rustworkx as rx


def weight_fn(edge_data):
    """
    Custom weight function for edges.
    """
    if len(edge_data) > 2:
        # your existing logic: tuple[3] is raw_weight
        return edge_data[3]  # Use the weight from the tuple
    # default weight
    return edge_data[1]


max_added = None


def update_added(added_set, added, op):
    """
    Append the two new node indices for a 2-qubit op,
    using a global next-available counter instead of
    scanning `added`.
    """
    global max_added

    q0, q1 = op[1][2]  # the two qubit indices

    # lazy initialize on first call
    if max_added is None:
        if added:
            max_added = max(added)
        else:
            # for the very first gate, we'll set it below
            max_added = None

    # both qubits have been seen before → allocate two brand-new nodes
    if q0 in added_set and q1 in added_set:
        a = max_added + 1 # type: ignore[unsupported-operator]
        b = max_added + 2 # type: ignore[unsupported-operator]
        added.extend([a, b])
        added_set.add(a)
        added_set.add(b)
        max_added += 2 # type: ignore[unsupported-operator]

    # only q0 seen → reuse q1 and allocate one fresh node
    elif q0 in added_set:
        a = q1
        max_added = max(max_added, a) # type: ignore[invalid-argument-type]
        b = max_added + 1
        added.extend([a, b])
        added_set.add(b)
        added_set.add(a)
        max_added += 1

    # only q1 seen → reuse q0 and allocate one fresh node
    elif q1 in added_set:
        a = q0
        max_added = max(max_added, a)  # type: ignore[invalid-argument-type]
        b = max_added + 1
        added.extend([a, b])
        added_set.add(b)
        added_set.add(a)
        max_added += 1

    # neither seen → just append those original qubit-IDs
    else:
        added.extend([q0, q1])
        added_set.add(q0)
        added_set.add(q1)
        # now that they’re in, prime max_added for future
        max_added = max(q0, q1)


def update_nodes_on_qubit(nodes_on_qubit, qubit, node):
    if qubit not in nodes_on_qubit:
        nodes_on_qubit[qubit] = [node]
    else:
        nodes_on_qubit[qubit].append(node)


def circ_to_graph(circuit, gateCutWeight=3, wireCutWeight=4):  # noqa: C901
    """
    Convert a quantum circuit to a graph representation.

    Args:
        circuit: A quantum circuit object.

    Returns:
        A graph representation of the circuit.
    """

    nodes_on_qubit = {}

    ops_2qb = [
        (
            idx,
            (
                x.operation.name,
                x.operation.num_qubits,
                [circuit.find_bit(i).index for i in x.qubits],
            ),
        )
        for idx, x in enumerate(circuit.data)
        if x.operation.num_qubits == 2
    ]

    G = rx.PyGraph()
    indxs = list(
        range(
            max(
                max([x[0] for x in ops_2qb]),
                len(ops_2qb) * 2,
            )
            + 1
        )
    )
    for i in indxs:
        G.add_node(i)

    track_nodes = {}
    added = []
    added_set = set()

    for ind, op in enumerate(ops_2qb):
        if (
            op[1][2][0] not in track_nodes.keys()
            and op[1][2][1] not in track_nodes.keys()
        ):
            update_added(added_set, added, op)
            if op[1][2][0] in added_set:
                track_nodes[op[1][2][0]] = [(added[-2])]
            else:
                track_nodes[op[1][2][0]] = []
            if op[1][2][1] in added_set:
                track_nodes[op[1][2][1]] = [added[-1]]
            else:
                track_nodes[op[1][2][1]] = []

            update_nodes_on_qubit(nodes_on_qubit, op[1][2][0], added[-2])
            update_nodes_on_qubit(nodes_on_qubit, op[1][2][1], added[-1])

            G.add_edge(
                added[-2],
                added[-1],
                (op[0], op[1][0], op[1][2], gateCutWeight),
            )
        elif op[1][2][1] not in track_nodes.keys():
            update_added(added_set, added, op)

            if op[1][2][1] in added_set:
                track_nodes[op[1][2][1]] = [(added[-1])]
            else:
                track_nodes[op[1][2][1]] = []

            track_nodes[op[1][2][0]].append(added[-2])
            ind = track_nodes[op[1][2][0]]

            update_nodes_on_qubit(nodes_on_qubit, op[1][2][0], added[-2])
            update_nodes_on_qubit(nodes_on_qubit, op[1][2][1], added[-1])

            G.add_edge(
                added[-2],
                added[-1],
                (op[0], op[1][0], op[1][2], gateCutWeight),
            )
            if ind == []:
                G.add_edge(op[1][2][0], added[-1], (op[1][2][0], wireCutWeight))
            else:
                if len(ind) == 1:
                    G.add_edge(ind[-1], op[1][2][0], (op[1][2][0], wireCutWeight))
                else:
                    G.add_edge(ind[-1], ind[-2], (op[1][2][0], wireCutWeight))
        elif op[1][2][0] not in track_nodes.keys():
            update_added(added_set, added, op)

            if op[1][2][0] in added_set:
                track_nodes[op[1][2][0]] = [(added[-2])]
            else:
                track_nodes[op[1][2][0]] = []

            track_nodes[op[1][2][1]].append(added[-1])
            ind = track_nodes[op[1][2][1]]

            update_nodes_on_qubit(nodes_on_qubit, op[1][2][0], added[-2])
            update_nodes_on_qubit(nodes_on_qubit, op[1][2][1], added[-1])

            G.add_edge(
                added[-2],
                added[-1],
                (op[0], op[1][0], op[1][2], gateCutWeight),
            )
            if ind == []:
                G.add_edge(op[1][2][1], added[-1], (op[1][2][1], wireCutWeight))
            else:
                if len(ind) == 1:
                    G.add_edge(ind[-1], op[1][2][1], (op[1][2][1], wireCutWeight))
                else:
                    G.add_edge(ind[-1], ind[-2], (op[1][2][1], wireCutWeight))
        else:
            update_added(added_set, added, op)
            track_nodes[op[1][2][0]].append(added[-2])
            track_nodes[op[1][2][1]].append(added[-1])
            ind0 = track_nodes[op[1][2][0]]
            ind1 = track_nodes[op[1][2][1]]

            update_nodes_on_qubit(nodes_on_qubit, op[1][2][0], added[-2])
            update_nodes_on_qubit(nodes_on_qubit, op[1][2][1], added[-1])

            G.add_edge(
                added[-2],
                added[-1],
                (op[0], op[1][0], op[1][2], gateCutWeight),
            )
            if ind0 == []:
                G.add_edge(op[1][2][0], added[-1], (op[1][2][0], wireCutWeight))
            else:
                if len(ind0) == 1:
                    G.add_edge(ind0[-1], op[1][2][0], (op[1][2][0], wireCutWeight))
                else:
                    G.add_edge(ind0[-1], ind0[-2], (op[1][2][0], wireCutWeight))
            if ind1 == []:
                G.add_edge(op[1][2][1], added[-2], (op[1][2][1], wireCutWeight))
            else:
                if len(ind1) == 1:
                    G.add_edge(ind1[-1], op[1][2][1], (op[1][2][1], wireCutWeight))
                else:
                    G.add_edge(ind1[-1], ind1[-2], (op[1][2][1], wireCutWeight))

    # Remove nodes with degree 0 (not connected to anything)
    isolated_nodes = [node for node in G.node_indices() if G.degree(node) == 0]

    G.remove_nodes_from(isolated_nodes)
    G, nodes_on_qubit = final_graph(G, nodes_on_qubit)
    return G, nodes_on_qubit


def final_graph(graph, nodes_on_qubit):
    # Relabel nodes to ensure consecutive numbering
    # Since rustworkx doesn't have relabel_nodes, we need to create a new graph manually
    old_to_new = {
        old_idx: new_idx for new_idx, old_idx in enumerate(graph.node_indices())
    }

    # Create a new graph
    new_graph = rx.PyGraph()

    # Add nodes with consecutive indices (0, 1, 2, ...)
    for i in range(len(graph.node_indices())):
        new_graph.add_node(i)

    # Add edges with remapped indices
    for edge in graph.edge_list():
        old_source, old_target = edge
        new_source = old_to_new[old_source]
        new_target = old_to_new[old_target]
        edge_data = graph.get_edge_data(old_source, old_target)
        new_graph.add_edge(new_source, new_target, edge_data)

    # Update nodes_on_qubit dictionary with new node indices
    updated_nodes_on_qubit = {}
    for qubit, old_node_list in nodes_on_qubit.items():
        updated_nodes_on_qubit[qubit] = [
            old_to_new[old_node] for old_node in old_node_list if old_node in old_to_new
        ]

    return new_graph, updated_nodes_on_qubit
