Repository: Qiskit/qiskit

Issue:
Control flow mpl visualization not respecting layout
### Environment

- **Qiskit Terra version**: 0.25.0
- **Python version**: 3.11
- **Operating system**:  Linux


### What is happening?

When running the matplotlib based circuit visualization with a transpiled circuit that is using control flow. The new control flow visualization support is drawing the gates on the incorrect qubit. It looks like it's not doing the control flow block qubit mapping correctly to factor in the layout. This results in the gates inside a control flow block ending up outside the box drawn for the control flow block.

For example:

![drawing](https://github.com/Qiskit/qiskit-terra/assets/2447371/e33ee18d-338c-410b-98c7-b08b3deabfbf)

### How can we reproduce the issue?

```python
from qiskit.circuit import QuantumCircuit, QuantumRegister, ClassicalRegister
from qiskit.compiler import transpile
from qiskit.providers.fake_provider import FakeBelemV2
from qiskit.circuit import SwitchCaseOp
​
qreg = QuantumRegister(2)
creg = ClassicalRegister(2)
qc = QuantumCircuit(qreg, creg)
​
qc.h([0, 1])
qc.h([0, 1])
qc.h([0, 1])
qc.measure([0, 1], [0, 1])
with qc.switch(creg) as case:
  with case(0):  # if the register is '00'
    qc.z(0)
  with case(1, 2):  # if the register is '01' or '10'
    qc.cx(0, 1)
  with case(case.DEFAULT):  # the default case
    qc.h(0)
​
backend = FakeBelemV2()
backend.target.add_instruction(SwitchCaseOp, name="switch_case")
​
transpile(qc, backend, optimization_level=2, seed_transpiler=671_42).draw('mpl')
```



### What should happen?

The gates on the right end of the circuit on ancilla_0 should be on q2_1 (and the cx should be shifted down by one row) in the same position.

### Any suggestions?

_No response_
Control flow mpl visualization not respecting layout
### Environment

- **Qiskit Terra version**: 0.25.0
- **Python version**: 3.11
- **Operating system**:  Linux


### What is happening?

When running the matplotlib based circuit visualization with a transpiled circuit that is using control flow. The new control flow visualization support is drawing the gates on the incorrect qubit. It looks like it's not doing the control flow block qubit mapping correctly to factor in the layout. This results in the gates inside a control flow block ending up outside the box drawn for the control flow block.

For example:

![drawing](https://github.com/Qiskit/qiskit-terra/assets/2447371/e33ee18d-338c-410b-98c7-b08b3deabfbf)

### How can we reproduce the issue?

```python
from qiskit.circuit import QuantumCircuit, QuantumRegister, ClassicalRegister
from qiskit.compiler import transpile
from qiskit.providers.fake_provider import FakeBelemV2
from qiskit.circuit import SwitchCaseOp
​
qreg = QuantumRegister(2)
creg = ClassicalRegister(2)
qc = QuantumCircuit(qreg, creg)
​
qc.h([0, 1])
qc.h([0, 1])
qc.h([0, 1])
qc.measure([0, 1], [0, 1])
with qc.switch(creg) as case:
  with case(0):  # if the register is '00'
    qc.z(0)
  with case(1, 2):  # if the register is '01' or '10'
    qc.cx(0, 1)
  with case(case.DEFAULT):  # the default case
    qc.h(0)
​
backend = FakeBelemV2()
backend.target.add_instruction(SwitchCaseOp, name="switch_case")
​
transpile(qc, backend, optimization_level=2, seed_transpiler=671_42).draw('mpl')
```



### What should happen?

The gates on the right end of the circuit on ancilla_0 should be on q2_1 (and the cx should be shifted down by one row) in the same position.

### Any suggestions?

_No response_
Fix inner gate placement for control flow blocks
<!--
⚠️ If you do not respect this template, your pull request will be closed.
⚠️ Your pull request title should be short detailed and understandable for all.
⚠️ Also, please add a release note file using reno if the change needs to be
  documented in the release notes.
⚠️ If your pull request fixes an open issue, please link to the issue.

- [ ] I have added the tests to cover my changes.
- [ ] I have updated the documentation accordingly.
- [ ] I have read the CONTRIBUTING document.
-->

### Summary

This commit fixes a bug in the mpl visualization of control flow blocks. It was assuming the qubit order of the control flow block matched the qubit order in the outer circuit. This could result in the output visualization putting the gates from a control flow block on the incorrect qubits. This commit corrects this oversight by mapping the input qargs to the control flow op to the inner qubits in the block circuit. This results in placing the gates on the correct wires in the output visualization regardless of where the qubits in the circuit are.

### Details and comments

Fixes #10601

Co-authored-by:  Edwin Navarro <enavarro@comcast.net>


Discussion:

Resolving patch:
diff --git a/qiskit/visualization/circuit/matplotlib.py b/qiskit/visualization/circuit/matplotlib.py
--- a/qiskit/visualization/circuit/matplotlib.py
+++ b/qiskit/visualization/circuit/matplotlib.py
@@ -130,19 +130,27 @@ def __init__(
         self._global_phase = self._circuit.global_phase
         self._calibrations = self._circuit.calibrations
 
-        for node in itertools.chain.from_iterable(self._nodes):
-            if node.cargs and node.op.name != "measure":
-                if cregbundle:
-                    warn(
-                        "Cregbundle set to False since an instruction needs to refer"
-                        " to individual classical wire",
-                        RuntimeWarning,
-                        3,
-                    )
-                self._cregbundle = False
-                break
-        else:
-            self._cregbundle = True if cregbundle is None else cregbundle
+        def check_clbit_in_inst(circuit, cregbundle):
+            if cregbundle is False:
+                return False
+            for inst in circuit.data:
+                if isinstance(inst.operation, ControlFlowOp):
+                    for block in inst.operation.blocks:
+                        if check_clbit_in_inst(block, cregbundle) is False:
+                            return False
+                elif inst.clbits and not isinstance(inst.operation, Measure):
+                    if cregbundle is not False:
+                        warn(
+                            "Cregbundle set to False since an instruction needs to refer"
+                            " to individual classical wire",
+                            RuntimeWarning,
+                            3,
+                        )
+                    return False
+
+            return True
+
+        self._cregbundle = check_clbit_in_inst(circuit, cregbundle)
 
         self._lwidth1 = 1.0
         self._lwidth15 = 1.5
@@ -308,12 +316,12 @@ def draw(self, filename=None, verbose=False):
         self._set_bit_reg_info(wire_map, qubits_dict, clbits_dict, glob_data)
 
         # get layer widths - flow gates are initialized here
-        layer_widths = self._get_layer_widths(node_data, wire_map, glob_data)
+        layer_widths = self._get_layer_widths(node_data, wire_map, self._circuit, glob_data)
 
         # load the coordinates for each top level gate and compute number of folds.
         # coordinates for flow gates are loaded before draw_ops
         max_x_index = self._get_coords(
-            node_data, wire_map, layer_widths, qubits_dict, clbits_dict, glob_data
+            node_data, wire_map, self._circuit, layer_widths, qubits_dict, clbits_dict, glob_data
         )
         num_folds = max(0, max_x_index - 1) // self._fold if self._fold > 0 else 0
 
@@ -382,6 +390,7 @@ def draw(self, filename=None, verbose=False):
             self._nodes,
             node_data,
             wire_map,
+            self._circuit,
             layer_widths,
             qubits_dict,
             clbits_dict,
@@ -399,7 +408,7 @@ def draw(self, filename=None, verbose=False):
             matplotlib_close_if_inline(mpl_figure)
             return mpl_figure
 
-    def _get_layer_widths(self, node_data, wire_map, glob_data):
+    def _get_layer_widths(self, node_data, wire_map, outer_circuit, glob_data):
         """Compute the layer_widths for the layers"""
 
         layer_widths = {}
@@ -415,14 +424,14 @@ def _get_layer_widths(self, node_data, wire_map, glob_data):
                 op = node.op
                 node_data[node] = NodeData()
                 node_data[node].width = WID
-                num_ctrl_qubits = 0 if not hasattr(op, "num_ctrl_qubits") else op.num_ctrl_qubits
+                num_ctrl_qubits = getattr(op, "num_ctrl_qubits", 0)
                 if (
                     getattr(op, "_directive", False) and (not op.label or not self._plot_barriers)
                 ) or isinstance(op, Measure):
                     node_data[node].raw_gate_text = op.name
                     continue
 
-                base_type = None if not hasattr(op, "base_gate") else op.base_gate
+                base_type = getattr(op, "base_gate", None)
                 gate_text, ctrl_text, raw_gate_text = get_gate_ctrl_text(
                     op, "mpl", style=self._style, calibrations=self._calibrations
                 )
@@ -435,7 +444,7 @@ def _get_layer_widths(self, node_data, wire_map, glob_data):
                 if (
                     (len(node.qargs) - num_ctrl_qubits) == 1
                     and len(gate_text) < 3
-                    and (not hasattr(op, "params") or len(op.params) == 0)
+                    and len(getattr(op, "params", [])) == 0
                     and ctrl_text is None
                 ):
                     continue
@@ -450,8 +459,7 @@ def _get_layer_widths(self, node_data, wire_map, glob_data):
                 )
                 # get param_width, but 0 for gates with array params or circuits in params
                 if (
-                    hasattr(op, "params")
-                    and len(op.params) > 0
+                    len(getattr(op, "params", [])) > 0
                     and not any(isinstance(param, np.ndarray) for param in op.params)
                     and not any(isinstance(param, QuantumCircuit) for param in op.params)
                 ):
@@ -508,24 +516,33 @@ def _get_layer_widths(self, node_data, wire_map, glob_data):
                         # Depth of nested ControlFlowOp used for color of box
                         if self._flow_parent is not None:
                             node_data[node].nest_depth = node_data[self._flow_parent].nest_depth + 1
-                        # Update the wire_map with the qubits from the inner circuit
-                        flow_wire_map = {
-                            inner: wire_map[outer]
-                            for outer, inner in zip(self._qubits, circuit.qubits)
-                            if inner not in wire_map
-                        }
-                        if not flow_wire_map:
-                            flow_wire_map = wire_map
+
+                        # Build the wire_map to be used by this flow op
+                        flow_wire_map = wire_map.copy()
+                        flow_wire_map.update(
+                            {
+                                inner: wire_map[outer]
+                                for outer, inner in zip(node.qargs, circuit.qubits)
+                            }
+                        )
+                        for outer, inner in zip(node.cargs, circuit.clbits):
+                            if self._cregbundle and (
+                                (in_reg := get_bit_register(outer_circuit, inner)) is not None
+                            ):
+                                out_reg = get_bit_register(outer_circuit, outer)
+                                flow_wire_map.update({in_reg: wire_map[out_reg]})
+                            else:
+                                flow_wire_map.update({inner: wire_map[outer]})
 
                         # Get the layered node lists and instantiate a new drawer class for
                         # the circuit inside the ControlFlowOp.
-                        qubits, clbits, nodes = _get_layered_instructions(
+                        qubits, clbits, flow_nodes = _get_layered_instructions(
                             circuit, wire_map=flow_wire_map
                         )
                         flow_drawer = MatplotlibDrawer(
                             qubits,
                             clbits,
-                            nodes,
+                            flow_nodes,
                             circuit,
                             style=self._style,
                             plot_barriers=self._plot_barriers,
@@ -539,14 +556,14 @@ def _get_layer_widths(self, node_data, wire_map, glob_data):
                         self._flow_drawers[node].append(flow_drawer)
 
                         # Recursively call _get_layer_widths for the circuit inside the ControlFlowOp
-                        flow_widths = flow_drawer._get_layer_widths(node_data, wire_map, glob_data)
+                        flow_widths = flow_drawer._get_layer_widths(
+                            node_data, flow_wire_map, outer_circuit, glob_data
+                        )
                         layer_widths.update(flow_widths)
 
-                        # Gates within a SwitchCaseOp need to know which case they are in
-                        for flow_layer in nodes:
+                        for flow_layer in flow_nodes:
                             for flow_node in flow_layer:
-                                if isinstance(node.op, SwitchCaseOp):
-                                    node_data[flow_node].case_num = circ_num
+                                node_data[flow_node].circ_num = circ_num
 
                         # Add up the width values of the same flow_parent that are not -1
                         # to get the raw_gate_width
@@ -560,6 +577,7 @@ def _get_layer_widths(self, node_data, wire_map, glob_data):
                         # Minor adjustment so else and case section gates align with indexes
                         if circ_num > 0:
                             raw_gate_width += 0.045
+
                         node_data[node].width.append(raw_gate_width)
 
                 # Otherwise, standard gate or multiqubit gate
@@ -663,12 +681,12 @@ def _get_coords(
         self,
         node_data,
         wire_map,
+        outer_circuit,
         layer_widths,
         qubits_dict,
         clbits_dict,
         glob_data,
         flow_parent=None,
-        is_not_first_block=None,
     ):
         """Load all the coordinate info needed to place the gates on the drawing."""
 
@@ -677,19 +695,13 @@ def _get_coords(
             curr_x_index = prev_x_index + 1
             l_width = []
             for node in layer:
-                # For gates inside if, else, while, or case set the x_index and if it's an
-                # else or case increment by if width. For additional cases increment by
-                # width of previous cases.
+                # For gates inside a flow op set the x_index and if it's an else or case,
+                # increment by if/switch width. If more cases increment by width of previous cases.
                 if flow_parent is not None:
                     node_data[node].x_index = node_data[flow_parent].x_index + curr_x_index + 1
-                    if is_not_first_block:
-                        # Add index space for else or first case if switch/case
-                        node_data[node].x_index += int(node_data[flow_parent].width[0]) + 1
-
-                        # Add index space for remaining cases for switch/case
-                        if node_data[node].case_num > 1:
-                            for width in node_data[flow_parent].width[1 : node_data[node].case_num]:
-                                node_data[node].x_index += int(width) + 1
+                    if node_data[node].circ_num > 0:
+                        for width in node_data[flow_parent].width[: node_data[node].circ_num]:
+                            node_data[node].x_index += int(width) + 1
 
                 # get qubit indexes
                 q_indxs = []
@@ -701,9 +713,12 @@ def _get_coords(
                 c_indxs = []
                 for carg in node.cargs:
                     if carg in self._clbits:
-                        register = get_bit_register(self._circuit, carg)
-                        if register is not None and self._cregbundle:
-                            c_indxs.append(wire_map[register])
+                        if self._cregbundle:
+                            register = get_bit_register(outer_circuit, carg)
+                            if register is not None:
+                                c_indxs.append(wire_map[register])
+                            else:
+                                c_indxs.append(wire_map[carg])
                         else:
                             c_indxs.append(wire_map[carg])
 
@@ -935,25 +950,40 @@ def _draw_regs_wires(self, num_folds, xmax, max_x_index, qubits_dict, clbits_dic
                 )
 
     def _add_nodes_and_coords(
-        self, nodes, node_data, wire_map, layer_widths, qubits_dict, clbits_dict, glob_data
+        self,
+        nodes,
+        node_data,
+        wire_map,
+        outer_circuit,
+        layer_widths,
+        qubits_dict,
+        clbits_dict,
+        glob_data,
     ):
         """Add the nodes from ControlFlowOps and their coordinates to the main circuit"""
         for flow_drawers in self._flow_drawers.values():
-            for i, flow_drawer in enumerate(flow_drawers):
+            for flow_drawer in flow_drawers:
                 nodes += flow_drawer._nodes
                 flow_drawer._get_coords(
                     node_data,
                     flow_drawer._flow_wire_map,
+                    outer_circuit,
                     layer_widths,
                     qubits_dict,
                     clbits_dict,
                     glob_data,
                     flow_parent=flow_drawer._flow_parent,
-                    is_not_first_block=(i > 0),
                 )
                 # Recurse for ControlFlowOps inside the flow_drawer
                 flow_drawer._add_nodes_and_coords(
-                    nodes, node_data, wire_map, layer_widths, qubits_dict, clbits_dict, glob_data
+                    nodes,
+                    node_data,
+                    wire_map,
+                    outer_circuit,
+                    layer_widths,
+                    qubits_dict,
+                    clbits_dict,
+                    glob_data,
                 )
 
     def _draw_ops(
@@ -961,6 +991,7 @@ def _draw_ops(
         nodes,
         node_data,
         wire_map,
+        outer_circuit,
         layer_widths,
         qubits_dict,
         clbits_dict,
@@ -971,7 +1002,14 @@ def _draw_ops(
 
         # Add the nodes from all the ControlFlowOps and their coordinates to the main nodes
         self._add_nodes_and_coords(
-            nodes, node_data, wire_map, layer_widths, qubits_dict, clbits_dict, glob_data
+            nodes,
+            node_data,
+            wire_map,
+            outer_circuit,
+            layer_widths,
+            qubits_dict,
+            clbits_dict,
+            glob_data,
         )
         prev_x_index = -1
         for layer in nodes:
@@ -999,11 +1037,11 @@ def _draw_ops(
                         )
                         for ii in clbits_dict
                     ]
-                    self._condition(node, node_data, wire_map, cond_xy, glob_data)
+                    self._condition(node, node_data, wire_map, outer_circuit, cond_xy, glob_data)
 
                 # draw measure
                 if isinstance(op, Measure):
-                    self._measure(node, node_data, glob_data)
+                    self._measure(node, node_data, outer_circuit, glob_data)
 
                 # draw barriers, snapshots, etc.
                 elif getattr(op, "_directive", False):
@@ -1043,7 +1081,7 @@ def _get_colors(self, node, node_data):
         """Get all the colors needed for drawing the circuit"""
 
         op = node.op
-        base_name = None if not hasattr(op, "base_gate") else op.base_gate.name
+        base_name = getattr(getattr(op, "base_gate", None), "name", None)
         color = None
         if node_data[node].raw_gate_text in self._style["dispcol"]:
             color = self._style["dispcol"][node_data[node].raw_gate_text]
@@ -1087,7 +1125,7 @@ def _get_colors(self, node, node_data):
         node_data[node].sc = sc
         node_data[node].lc = lc
 
-    def _condition(self, node, node_data, wire_map, cond_xy, glob_data):
+    def _condition(self, node, node_data, wire_map, outer_circuit, cond_xy, glob_data):
         """Add a conditional to a gate"""
 
         # For SwitchCaseOp convert the target to a fully closed Clbit or register
@@ -1107,25 +1145,24 @@ def _condition(self, node, node_data, wire_map, cond_xy, glob_data):
         cond_pos = []
 
         if isinstance(condition, expr.Expr):
+
             # If fixing this, please update the docstrings of `QuantumCircuit.draw` and
             # `visualization.circuit_drawer` to remove warnings.
+
             condition_bits = condition_resources(condition).clbits
-            label = "[expression]"
+            label = "[expr]"
             override_fc = True
             registers = collections.defaultdict(list)
             for bit in condition_bits:
-                registers[get_bit_register(self._circuit, bit)].append(bit)
+                registers[get_bit_register(outer_circuit, bit)].append(bit)
             # Registerless bits don't care whether cregbundle is set.
             cond_pos.extend(cond_xy[wire_map[bit] - first_clbit] for bit in registers.pop(None, ()))
             if self._cregbundle:
-                cond_pos.extend(
-                    cond_xy[wire_map[register[0]] - first_clbit] for register in registers
-                )
+                cond_pos.extend(cond_xy[wire_map[register] - first_clbit] for register in registers)
             else:
                 cond_pos.extend(
                     cond_xy[wire_map[bit] - first_clbit]
-                    for register, bits in registers.items()
-                    for bit in bits
+                    for bit in itertools.chain.from_iterable(registers.values())
                 )
             val_bits = ["1"] * len(cond_pos)
         else:
@@ -1146,7 +1183,7 @@ def _condition(self, node, node_data, wire_map, cond_xy, glob_data):
 
             # If it's a register bit and cregbundle, need to use the register to find the location
             elif self._cregbundle and isinstance(cond_bit_reg, Clbit):
-                register = get_bit_register(self._circuit, cond_bit_reg)
+                register = get_bit_register(outer_circuit, cond_bit_reg)
                 if register is not None:
                     cond_pos.append(cond_xy[wire_map[register] - first_clbit])
                 else:
@@ -1171,8 +1208,8 @@ def _condition(self, node, node_data, wire_map, cond_xy, glob_data):
         qubit_b = min(node_data[node].q_xy, key=lambda xy: xy[1])
         clbit_b = min(xy_plot, key=lambda xy: xy[1])
 
-        # For IfElseOp, WhileLoopOp or SwitchCaseOp, place the condition
-        # at almost the left edge of the box
+        # For IfElseOp, WhileLoopOp or SwitchCaseOp, place the condition line
+        # near the left edge of the box
         if isinstance(node.op, (IfElseOp, WhileLoopOp, SwitchCaseOp)):
             qubit_b = (qubit_b[0], qubit_b[1] - (0.5 * HIG + 0.14))
 
@@ -1193,11 +1230,11 @@ def _condition(self, node, node_data, wire_map, cond_xy, glob_data):
         )
         self._line(qubit_b, clbit_b, lc=self._style["cc"], ls=self._style["cline"])
 
-    def _measure(self, node, node_data, glob_data):
+    def _measure(self, node, node_data, outer_circuit, glob_data):
         """Draw the measure symbol and the line to the clbit"""
         qx, qy = node_data[node].q_xy[0]
         cx, cy = node_data[node].c_xy[0]
-        register, _, reg_index = get_bit_reg_index(self._circuit, node.cargs[0])
+        register, _, reg_index = get_bit_reg_index(outer_circuit, node.cargs[0])
 
         # draw gate box
         self._gate(node, node_data, glob_data)
@@ -1445,6 +1482,7 @@ def _flow_op_gate(self, node, node_data, glob_data):
 
         if_width = node_data[node].width[0] + WID
         box_width = if_width
+
         # Add the else and case widths to the if_width
         for ewidth in node_data[node].width[1:]:
             if ewidth > 0.0:
@@ -1493,6 +1531,7 @@ def _flow_op_gate(self, node, node_data, glob_data):
             elif isinstance(node.op, SwitchCaseOp):
                 flow_text = "Switch"
 
+            # Indicate type of ControlFlowOp
             self._ax.text(
                 xpos - x_shift - 0.08,
                 ypos_max + 0.2 - y_shift,
@@ -1542,7 +1581,7 @@ def _flow_op_gate(self, node, node_data, glob_data):
             # If there's an else or a case draw the vertical line and the name
             else_case_text = "Else" if isinstance(node.op, IfElseOp) else "Case"
             ewidth_incr = if_width
-            for case_num, ewidth in enumerate(node_data[node].width[1:]):
+            for circ_num, ewidth in enumerate(node_data[node].width[1:]):
                 if ewidth > 0.0:
                     self._ax.plot(
                         [xpos + ewidth_incr + 0.3 - x_shift, xpos + ewidth_incr + 0.3 - x_shift],
@@ -1564,7 +1603,7 @@ def _flow_op_gate(self, node, node_data, glob_data):
                         zorder=PORDER_FLOW,
                     )
                     if isinstance(node.op, SwitchCaseOp):
-                        jump_val = node_data[node].jump_values[case_num]
+                        jump_val = node_data[node].jump_values[circ_num]
                         # If only one value, e.g. (0,)
                         if len(str(jump_val)) == 4:
                             jump_text = str(jump_val)[1]
@@ -1595,7 +1634,7 @@ def _control_gate(self, node, node_data, glob_data):
         """Draw a controlled gate"""
         op = node.op
         xy = node_data[node].q_xy
-        base_type = None if not hasattr(op, "base_gate") else op.base_gate
+        base_type = getattr(op, "base_gate", None)
         qubit_b = min(xy, key=lambda xy: xy[1])
         qubit_t = max(xy, key=lambda xy: xy[1])
         num_ctrl_qubits = op.num_ctrl_qubits
@@ -1730,7 +1769,7 @@ def _symmetric_gate(self, node, node_data, base_type, glob_data):
         xy = node_data[node].q_xy
         qubit_b = min(xy, key=lambda xy: xy[1])
         qubit_t = max(xy, key=lambda xy: xy[1])
-        base_type = None if not hasattr(op, "base_gate") else op.base_gate
+        base_type = getattr(op, "base_gate", None)
         ec = node_data[node].ec
         tc = node_data[node].tc
         lc = node_data[node].lc
@@ -1901,4 +1940,4 @@ def __init__(self):
         self.inside_flow = False
         self.indexset = ()  # List of indices used for ForLoopOp
         self.jump_values = []  # List of jump values used for SwitchCaseOp
-        self.case_num = 0  # Used for SwitchCaseOp
+        self.circ_num = 0  # Which block is it in op.blocks
diff --git a/qiskit/visualization/circuit/text.py b/qiskit/visualization/circuit/text.py
--- a/qiskit/visualization/circuit/text.py
+++ b/qiskit/visualization/circuit/text.py
@@ -17,7 +17,6 @@
 from warnings import warn
 from shutil import get_terminal_size
 import collections
-import itertools
 import sys
 
 from qiskit.circuit import Qubit, Clbit, ClassicalRegister
@@ -732,19 +731,27 @@ def __init__(
         self.vertical_compression = vertical_compression
         self._wire_map = {}
 
-        for node in itertools.chain.from_iterable(self.nodes):
-            if node.cargs and node.op.name != "measure":
-                if cregbundle:
-                    warn(
-                        "Cregbundle set to False since an instruction needs to refer"
-                        " to individual classical wire",
-                        RuntimeWarning,
-                        2,
-                    )
-                self.cregbundle = False
-                break
-        else:
-            self.cregbundle = True if cregbundle is None else cregbundle
+        def check_clbit_in_inst(circuit, cregbundle):
+            if cregbundle is False:
+                return False
+            for inst in circuit.data:
+                if isinstance(inst.operation, ControlFlowOp):
+                    for block in inst.operation.blocks:
+                        if check_clbit_in_inst(block, cregbundle) is False:
+                            return False
+                elif inst.clbits and not isinstance(inst.operation, Measure):
+                    if cregbundle is not False:
+                        warn(
+                            "Cregbundle set to False since an instruction needs to refer"
+                            " to individual classical wire",
+                            RuntimeWarning,
+                            3,
+                        )
+                    return False
+
+            return True
+
+        self.cregbundle = check_clbit_in_inst(circuit, cregbundle)
 
         if encoding:
             self.encoding = encoding
@@ -1257,7 +1264,13 @@ def build_layers(self):
         layers = [InputWire.fillup_layer(wire_names)]
 
         for node_layer in self.nodes:
-            layer = Layer(self.qubits, self.clbits, self.cregbundle, self._circuit, self._wire_map)
+            layer = Layer(
+                self.qubits,
+                self.clbits,
+                self.cregbundle,
+                self._circuit,
+                self._wire_map,
+            )
             for node in node_layer:
                 if isinstance(node.op, ControlFlowOp):
                     self._nest_depth = 0
@@ -1291,25 +1304,27 @@ def add_control_flow(self, node, layers, wire_map):
             # Update the wire_map with the qubits and clbits from the inner circuit
             flow_wire_map = wire_map.copy()
             flow_wire_map.update(
-                (inner, wire_map[outer]) for outer, inner in zip(node.qargs, circuit.qubits)
-            )
-            flow_wire_map.update(
-                (inner, wire_map[outer]) for outer, inner in zip(node.cargs, circuit.clbits)
+                {inner: wire_map[outer] for outer, inner in zip(node.qargs, circuit.qubits)}
             )
+            for outer, inner in zip(node.cargs, circuit.clbits):
+                if self.cregbundle and (
+                    (in_reg := get_bit_register(self._circuit, inner)) is not None
+                ):
+                    out_reg = get_bit_register(self._circuit, outer)
+                    flow_wire_map.update({in_reg: wire_map[out_reg]})
+                else:
+                    flow_wire_map.update({inner: wire_map[outer]})
+
             if circ_num > 0:
                 # Draw a middle box such as Else and Case
                 flow_layer = self.draw_flow_box(node, flow_wire_map, CF_MID, circ_num - 1)
                 layers.append(flow_layer.full_layer)
 
-            qubits, _, nodes = _get_layered_instructions(circuit, wire_map=flow_wire_map)
+            _, _, nodes = _get_layered_instructions(circuit, wire_map=flow_wire_map)
             for layer_nodes in nodes:
                 # Limit qubits sent to only ones from main circuit, so qubit_layer is correct length
                 flow_layer2 = Layer(
-                    qubits[: len(self.qubits)],
-                    self.clbits,
-                    self.cregbundle,
-                    circuit,
-                    flow_wire_map,
+                    self.qubits, self.clbits, self.cregbundle, self._circuit, flow_wire_map
                 )
                 for layer_node in layer_nodes:
                     if isinstance(layer_node.op, ControlFlowOp):
@@ -1372,7 +1387,13 @@ def draw_flow_box(self, node, flow_wire_map, section, circ_num=0):
         else:
             label = "End-" + depth
 
-        flow_layer = Layer(self.qubits, self.clbits, self.cregbundle, self._circuit, flow_wire_map)
+        flow_layer = Layer(
+            self.qubits,
+            self.clbits,
+            self.cregbundle,
+            self._circuit,
+            flow_wire_map,
+        )
         # If only 1 qubit, draw basic 1 qubit box
         if len(node.qargs) == 1:
             flow_layer.set_qubit(
@@ -1412,7 +1433,9 @@ def draw_flow_box(self, node, flow_wire_map, section, circ_num=0):
             )
         if conditional:
             if isinstance(node.op, SwitchCaseOp):
-                if isinstance(node.op.target, Clbit):
+                if isinstance(node.op.target, expr.Expr):
+                    condition = node.op.target
+                elif isinstance(node.op.target, Clbit):
                     condition = (node.op.target, 1)
                 else:
                     condition = (node.op.target, 2 ** (node.op.target.size) - 1)
@@ -1428,9 +1451,7 @@ class Layer:
 
     def __init__(self, qubits, clbits, cregbundle, circuit, wire_map):
         self.qubits = qubits
-        self.clbits_raw = clbits  # list of clbits ignoring cregbundle change below
         self._circuit = circuit
-
         if cregbundle:
             self.clbits = []
             previous_creg = None
@@ -1626,14 +1647,16 @@ def set_cl_multibox(self, condition, wire_map, top_connect="┴"):
         if isinstance(condition, expr.Expr):
             # If fixing this, please update the docstrings of `QuantumCircuit.draw` and
             # `visualization.circuit_drawer` to remove warnings.
-            label = "<expression>"
+            label = "[expr]"
             out = []
             condition_bits = node_resources(condition).clbits
             registers = collections.defaultdict(list)
             for bit in condition_bits:
                 registers[get_bit_register(self._circuit, bit)].append(bit)
             if registerless := registers.pop(None, ()):
-                out.extend(self.set_cond_bullets(label, ["1"] * len(registerless), registerless))
+                out.extend(
+                    self.set_cond_bullets(label, ["1"] * len(registerless), registerless, wire_map)
+                )
             if self.cregbundle:
                 # It's hard to do something properly sensible here without more major rewrites, so
                 # as a minimum to *not crash* we'll just treat a condition that touches part of a
@@ -1642,7 +1665,7 @@ def set_cl_multibox(self, condition, wire_map, top_connect="┴"):
                     self.set_clbit(register[0], BoxOnClWire(label=label, top_connect=top_connect))
             else:
                 for register, bits in registers.items():
-                    out.extend(self.set_cond_bullets(label, ["1"] * len(bits), bits))
+                    out.extend(self.set_cond_bullets(label, ["1"] * len(bits), bits, wire_map))
             return out
 
         label, val_bits = get_condition_label_val(condition, self._circuit, self.cregbundle)


Test patch:
diff --git a/test/python/visualization/test_circuit_text_drawer.py b/test/python/visualization/test_circuit_text_drawer.py
--- a/test/python/visualization/test_circuit_text_drawer.py
+++ b/test/python/visualization/test_circuit_text_drawer.py
@@ -5281,8 +5281,8 @@ def test_text_drawer_cp437(self):
 class TestCircuitControlFlowOps(QiskitVisualizationTestCase):
     """Test ControlFlowOps."""
 
-    def test_if_op(self):
-        """Test an IfElseOp with if only"""
+    def test_if_op_bundle_false(self):
+        """Test an IfElseOp with if only and cregbundle false"""
         expected = "\n".join(
             [
                 "      ┌────── ┌───┐      ───────┐ ",
@@ -5305,6 +5305,35 @@ def test_if_op(self):
         cr = ClassicalRegister(2, "cr")
         circuit = QuantumCircuit(qr, cr)
 
+        with circuit.if_test((cr[1], 1)):
+            circuit.h(0)
+            circuit.cx(0, 1)
+        self.assertEqual(
+            str(_text_circuit_drawer(circuit, initial_state=False, cregbundle=False)), expected
+        )
+
+    def test_if_op_bundle_true(self):
+        """Test an IfElseOp with if only and cregbundle true"""
+        expected = "\n".join(
+            [
+                "        ┌──────   ┌───┐      ───────┐ ",
+                " q_0: ──┤       ──┤ H ├──■──        ├─",
+                "        │ If-0    └───┘┌─┴─┐  End-0 │ ",
+                " q_1: ──┤       ───────┤ X ├        ├─",
+                "        └──╥───        └───┘ ───────┘ ",
+                " q_2: ─────╫──────────────────────────",
+                "           ║                          ",
+                " q_3: ─────╫──────────────────────────",
+                "      ┌────╨─────┐                    ",
+                "cr: 2/╡ cr_1=0x1 ╞════════════════════",
+                "      └──────────┘                    ",
+            ]
+        )
+
+        qr = QuantumRegister(4, "q")
+        cr = ClassicalRegister(2, "cr")
+        circuit = QuantumCircuit(qr, cr)
+
         with circuit.if_test((cr[1], 1)):
             circuit.h(0)
             circuit.cx(0, 1)
@@ -5351,7 +5380,9 @@ def test_if_else_with_body_specified(self):
 
         circuit.if_else((cr[1], 1), circuit2, None, [0, 1, 2], [0, 1, 2])
         circuit.x(0, label="X1i")
-        self.assertEqual(str(_text_circuit_drawer(circuit, initial_state=False)), expected)
+        self.assertEqual(
+            str(_text_circuit_drawer(circuit, initial_state=False, cregbundle=False)), expected
+        )
 
     def test_if_op_nested_wire_order(self):
         """Test IfElseOp with nested if's and wire_order change."""
@@ -5433,7 +5464,10 @@ def test_if_op_nested_wire_order(self):
         self.assertEqual(
             str(
                 _text_circuit_drawer(
-                    circuit, fold=77, initial_state=False, wire_order=[2, 0, 3, 1, 4, 5, 6]
+                    circuit,
+                    fold=77,
+                    initial_state=False,
+                    wire_order=[2, 0, 3, 1, 4, 5, 6],
                 )
             ),
             expected,
@@ -5473,7 +5507,9 @@ def test_while_loop(self):
             circuit.measure(0, 0)
             with circuit.if_test((cr[2], 1)):
                 circuit.x(0)
-        self.assertEqual(str(_text_circuit_drawer(circuit, initial_state=False)), expected)
+        self.assertEqual(
+            str(_text_circuit_drawer(circuit, initial_state=False, cregbundle=False)), expected
+        )
 
     def test_for_loop(self):
         """Test ForLoopOp."""
@@ -5510,7 +5546,10 @@ def test_for_loop(self):
             circuit.measure(0, 0)
             with circuit.if_test((cr[2], 1)):
                 circuit.z(0)
-        self.assertEqual(str(_text_circuit_drawer(circuit, fold=-1, initial_state=False)), expected)
+        self.assertEqual(
+            str(_text_circuit_drawer(circuit, fold=-1, initial_state=False, cregbundle=False)),
+            expected,
+        )
 
     def test_switch_case(self):
         """Test SwitchCaseOp."""
@@ -5563,7 +5602,10 @@ def test_switch_case(self):
             with case(case.DEFAULT):
                 circuit.cx(0, 1)
         circuit.h(0)
-        self.assertEqual(str(_text_circuit_drawer(circuit, fold=78, initial_state=False)), expected)
+        self.assertEqual(
+            str(_text_circuit_drawer(circuit, fold=78, initial_state=False, cregbundle=False)),
+            expected,
+        )
 
     def test_inner_wire_map_control_op(self):
         """Test that the gates inside ControlFlowOps land on correct qubits when transpiled"""
@@ -5600,7 +5642,10 @@ def test_inner_wire_map_control_op(self):
         backend.target.add_instruction(IfElseOp, name="if_else")
 
         circuit = transpile(qc, backend, optimization_level=2, seed_transpiler=671_42)
-        self.assertEqual(str(_text_circuit_drawer(circuit, fold=78, initial_state=False)), expected)
+        self.assertEqual(
+            str(_text_circuit_drawer(circuit, fold=78, initial_state=False, cregbundle=False)),
+            expected,
+        )
 
     def test_if_else_op_from_circuit_with_conditions(self):
         """Test an IfElseOp built from circuit with conditions inside the if using inner creg"""
@@ -5631,13 +5676,14 @@ def test_if_else_op_from_circuit_with_conditions(self):
         circuit.x(2).c_if(cr[1], 2)
 
         qr2 = QuantumRegister(3, "qr2")
-        cr2 = ClassicalRegister(3, "cr2")
-        qc2 = QuantumCircuit(qr2, cr2)
-        qc2.x(0, label="X1").c_if(cr2, 4)
-        qc2.x(1, label="X2").c_if(cr2[1], 1)
+        qc2 = QuantumCircuit(qr2, cr)
+        qc2.x(0, label="X1").c_if(cr, 4)
+        qc2.x(1, label="X2").c_if(cr[1], 1)
 
         circuit.if_else((cr[1], 1), qc2, None, [0, 1, 2], [0, 1, 2])
-        self.assertEqual(str(_text_circuit_drawer(circuit, initial_state=False)), expected)
+        self.assertEqual(
+            str(_text_circuit_drawer(circuit, initial_state=False, cregbundle=False)), expected
+        )
 
 
 if __name__ == "__main__":
diff --git a/test/visual/mpl/circuit/references/fold_with_conditions.png b/test/visual/mpl/circuit/references/fold_with_conditions.png
Binary files a/test/visual/mpl/circuit/references/fold_with_conditions.png and b/test/visual/mpl/circuit/references/fold_with_conditions.png differ
diff --git a/test/visual/mpl/circuit/references/if_else_op.png b/test/visual/mpl/circuit/references/if_else_op_false.png
similarity index 96%
rename from test/visual/mpl/circuit/references/if_else_op.png
rename to test/visual/mpl/circuit/references/if_else_op_false.png
Binary files a/test/visual/mpl/circuit/references/if_else_op.png and b/test/visual/mpl/circuit/references/if_else_op_false.png differ
diff --git a/test/visual/mpl/circuit/references/if_else_op_true.png b/test/visual/mpl/circuit/references/if_else_op_true.png
new file mode 100644
Binary files /dev/null and b/test/visual/mpl/circuit/references/if_else_op_true.png differ
diff --git a/test/visual/mpl/circuit/references/layout_control_flow.png b/test/visual/mpl/circuit/references/layout_control_flow.png
new file mode 100644
Binary files /dev/null and b/test/visual/mpl/circuit/references/layout_control_flow.png differ
diff --git a/test/visual/mpl/circuit/references/nested_layout_control_flow.png b/test/visual/mpl/circuit/references/nested_layout_control_flow.png
new file mode 100644
Binary files /dev/null and b/test/visual/mpl/circuit/references/nested_layout_control_flow.png differ
diff --git a/test/visual/mpl/circuit/test_circuit_matplotlib_drawer.py b/test/visual/mpl/circuit/test_circuit_matplotlib_drawer.py
--- a/test/visual/mpl/circuit/test_circuit_matplotlib_drawer.py
+++ b/test/visual/mpl/circuit/test_circuit_matplotlib_drawer.py
@@ -23,7 +23,7 @@
 
 from qiskit.test import QiskitTestCase
 from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, transpile
-from qiskit.providers.fake_provider import FakeTenerife
+from qiskit.providers.fake_provider import FakeTenerife, FakeBelemV2
 from qiskit.visualization.circuit.circuit_visualization import _matplotlib_circuit_drawer
 from qiskit.circuit.library import (
     XGate,
@@ -40,7 +40,7 @@
     Isometry,
 )
 from qiskit.circuit.library import MCXVChain
-from qiskit.circuit import Parameter, Qubit, Clbit
+from qiskit.circuit import Parameter, Qubit, Clbit, SwitchCaseOp, IfElseOp
 from qiskit.circuit.library import IQP
 from qiskit.quantum_info.random import random_unitary
 from qiskit.utils import optionals
@@ -1485,8 +1485,8 @@ def test_sidetext_with_condition(self):
 
     def test_fold_with_conditions(self):
         """Test that gates with conditions draw correctly when folding"""
-        qr = QuantumRegister(3)
-        cr = ClassicalRegister(5)
+        qr = QuantumRegister(3, "qr")
+        cr = ClassicalRegister(5, "cr")
         circuit = QuantumCircuit(qr, cr)
 
         circuit.append(U1Gate(0).control(1), [1, 0]).c_if(cr, 1)
@@ -1597,7 +1597,7 @@ def test_if_op(self):
             circuit.cx(0, 1)
 
         fname = "if_op.png"
-        self.circuit_drawer(circuit, filename=fname)
+        self.circuit_drawer(circuit, cregbundle=False, filename=fname)
 
         ratio = VisualTestUtilities._save_diff(
             self._image_path(fname),
@@ -1608,8 +1608,8 @@ def test_if_op(self):
         )
         self.assertGreaterEqual(ratio, 0.9999)
 
-    def test_if_else_op(self):
-        """Test the IfElseOp with else"""
+    def test_if_else_op_bundle_false(self):
+        """Test the IfElseOp with else with cregbundle False"""
         qr = QuantumRegister(4, "q")
         cr = ClassicalRegister(2, "cr")
         circuit = QuantumCircuit(qr, cr)
@@ -1620,7 +1620,32 @@ def test_if_else_op(self):
         with _else:
             circuit.cx(0, 1)
 
-        fname = "if_else_op.png"
+        fname = "if_else_op_false.png"
+        self.circuit_drawer(circuit, cregbundle=False, filename=fname)
+
+        ratio = VisualTestUtilities._save_diff(
+            self._image_path(fname),
+            self._reference_path(fname),
+            fname,
+            FAILURE_DIFF_DIR,
+            FAILURE_PREFIX,
+        )
+        self.assertGreaterEqual(ratio, 0.9999)
+
+    def test_if_else_op_bundle_true(self):
+        """Test the IfElseOp with else with cregbundle True"""
+        qr = QuantumRegister(4, "q")
+        cr = ClassicalRegister(2, "cr")
+        circuit = QuantumCircuit(qr, cr)
+
+        circuit.h(0)
+        with circuit.if_test((cr[1], 1)) as _else:
+            circuit.h(0)
+            circuit.cx(0, 1)
+        with _else:
+            circuit.cx(0, 1)
+
+        fname = "if_else_op_true.png"
         self.circuit_drawer(circuit, filename=fname)
 
         ratio = VisualTestUtilities._save_diff(
@@ -1645,7 +1670,7 @@ def test_if_else_op_textbook_style(self):
             circuit.cx(0, 1)
 
         fname = "if_else_op_textbook.png"
-        self.circuit_drawer(circuit, style="textbook", filename=fname)
+        self.circuit_drawer(circuit, style="textbook", cregbundle=False, filename=fname)
 
         ratio = VisualTestUtilities._save_diff(
             self._image_path(fname),
@@ -1679,7 +1704,7 @@ def test_if_else_with_body(self):
         circuit.x(0, label="X1i")
 
         fname = "if_else_body.png"
-        self.circuit_drawer(circuit, filename=fname)
+        self.circuit_drawer(circuit, cregbundle=False, filename=fname)
 
         ratio = VisualTestUtilities._save_diff(
             self._image_path(fname),
@@ -1829,7 +1854,7 @@ def test_while_loop_op(self):
                 circuit.x(0)
 
         fname = "while_loop.png"
-        self.circuit_drawer(circuit, filename=fname)
+        self.circuit_drawer(circuit, cregbundle=False, filename=fname)
 
         ratio = VisualTestUtilities._save_diff(
             self._image_path(fname),
@@ -1858,7 +1883,7 @@ def test_for_loop_op(self):
                 circuit.z(0)
 
         fname = "for_loop.png"
-        self.circuit_drawer(circuit, filename=fname)
+        self.circuit_drawer(circuit, cregbundle=False, filename=fname)
 
         ratio = VisualTestUtilities._save_diff(
             self._image_path(fname),
@@ -1890,8 +1915,71 @@ def test_switch_case_op(self):
         circuit.h(0)
 
         fname = "switch_case.png"
-        self.circuit_drawer(circuit, filename=fname)
+        self.circuit_drawer(circuit, cregbundle=False, filename=fname)
+
+        ratio = VisualTestUtilities._save_diff(
+            self._image_path(fname),
+            self._reference_path(fname),
+            fname,
+            FAILURE_DIFF_DIR,
+            FAILURE_PREFIX,
+        )
+        self.assertGreaterEqual(ratio, 0.9999)
 
+    def test_control_flow_layout(self):
+        """Test control flow with a layout set."""
+        qreg = QuantumRegister(2, "qr")
+        creg = ClassicalRegister(2, "cr")
+        qc = QuantumCircuit(qreg, creg)
+        qc.h([0, 1])
+        qc.h([0, 1])
+        qc.h([0, 1])
+        qc.measure([0, 1], [0, 1])
+        with qc.switch(creg) as case:
+            with case(0):
+                qc.z(0)
+            with case(1, 2):
+                qc.cx(0, 1)
+            with case(case.DEFAULT):
+                qc.h(0)
+        backend = FakeBelemV2()
+        backend.target.add_instruction(SwitchCaseOp, name="switch_case")
+        tqc = transpile(qc, backend, optimization_level=2, seed_transpiler=671_42)
+        fname = "layout_control_flow.png"
+        self.circuit_drawer(tqc, filename=fname)
+        ratio = VisualTestUtilities._save_diff(
+            self._image_path(fname),
+            self._reference_path(fname),
+            fname,
+            FAILURE_DIFF_DIR,
+            FAILURE_PREFIX,
+        )
+        self.assertGreaterEqual(ratio, 0.9999)
+
+    def test_control_flow_nested_layout(self):
+        """Test nested control flow with a layout set."""
+        qreg = QuantumRegister(2, "qr")
+        creg = ClassicalRegister(2, "cr")
+        qc = QuantumCircuit(qreg, creg)
+        qc.h([0, 1])
+        qc.h([0, 1])
+        qc.h([0, 1])
+        qc.measure([0, 1], [0, 1])
+        with qc.switch(creg) as case:
+            with case(0):
+                qc.z(0)
+            with case(1, 2):
+                with qc.if_test((creg[0], 0)):
+                    qc.cx(0, 1)
+            with case(case.DEFAULT):
+                with qc.if_test((creg[1], 0)):
+                    qc.h(0)
+        backend = FakeBelemV2()
+        backend.target.add_instruction(SwitchCaseOp, name="switch_case")
+        backend.target.add_instruction(IfElseOp, name="if_else")
+        tqc = transpile(qc, backend, optimization_level=2, seed_transpiler=671_42)
+        fname = "nested_layout_control_flow.png"
+        self.circuit_drawer(tqc, filename=fname)
         ratio = VisualTestUtilities._save_diff(
             self._image_path(fname),
             self._reference_path(fname),