From 1e95cd6ac9cc8cfe8fe6e2240c1a924266d252ca Mon Sep 17 00:00:00 2001
From: HGSilveri <henrique.silverio@tecnico.ulisboa.pt>
Date: Fri, 21 Jul 2023 15:41:02 +0200
Subject: [PATCH] Improvements to Sequence.switch_device()

---
 pulser-core/pulser/sequence/sequence.py | 21 +++++++++++++++------
 tests/test_sequence.py                  |  6 +++---
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/pulser-core/pulser/sequence/sequence.py b/pulser-core/pulser/sequence/sequence.py
index f29c4b1b..35e47e58 100644
--- a/pulser-core/pulser/sequence/sequence.py
+++ b/pulser-core/pulser/sequence/sequence.py
@@ -496,6 +496,11 @@ def check_retarget(ch_obj: Channel) -> bool:
         channel_match: dict[str, Any] = {}
         strict_error_message = ""
         ch_match_err = ""
+        active_eom_channels = [
+            {**dict(zip(("channel",), call.args)), **call.kwargs}["channel"]
+            for call in self._calls + self._to_build_calls
+            if call.name == "enable_eom_mode"
+        ]
         for old_ch_name, old_ch_obj in self.declared_channels.items():
             channel_match[old_ch_name] = None
             # Find the corresponding channel on the new device
@@ -509,27 +514,31 @@ def check_retarget(ch_obj: Channel) -> bool:
 
                 # We verify the channel class then
                 # check whether the addressing is Global or Local
+                type_match = type(old_ch_obj) == type(new_ch_obj)
                 basis_match = old_ch_obj.basis == new_ch_obj.basis
                 addressing_match = (
                     old_ch_obj.addressing == new_ch_obj.addressing
                 )
                 base_msg = f"No match for channel {old_ch_name}"
-                if not (basis_match and addressing_match):
+                if not (type_match and basis_match and addressing_match):
                     # If there already is a message, keeps it
                     ch_match_err = ch_match_err or (
-                        base_msg + " with the right basis and addressing."
+                        base_msg
+                        + " with the right type, basis and addressing."
                     )
                     continue
-                if any(
-                    call.name == "enable_eom_mode"
-                    for call in self._calls + self._to_build_calls
-                ):
+                if old_ch_name in active_eom_channels:
                     # Uses EOM mode, so the new device needs a matching
                     # EOM configuration
                     if new_ch_obj.eom_config is None:
                         ch_match_err = base_msg + " with an EOM configuration."
                         continue
                     if (
+                        # TODO: Improvements to this check:
+                        # 1. multiple_beam_control doesn't matter when there
+                        # is only one beam
+                        # 2. custom_buffer_time doesn't have to match as long
+                        # as `Channel_eom_buffer_time`` does
                         new_ch_obj.eom_config != old_ch_obj.eom_config
                         and strict
                     ):
diff --git a/tests/test_sequence.py b/tests/test_sequence.py
index c5324b67..79063b19 100644
--- a/tests/test_sequence.py
+++ b/tests/test_sequence.py
@@ -347,7 +347,7 @@ def test_switch_device_down(reg, devices, pulses, mappable_reg, parametrized):
     with pytest.raises(
         TypeError,
         match="No match for channel global2 with the"
-        " right basis and addressing.",
+        " right type, basis and addressing.",
     ):
         # Can't find a match for the 2nd rydberg_global
         seq.switch_device(Chadoq2)
@@ -402,13 +402,13 @@ def test_switch_device_down(reg, devices, pulses, mappable_reg, parametrized):
         mappable_reg=mappable_reg,
     )
     for dev_ in (
-        Chadoq2,  # Different Channels basis
+        Chadoq2,  # Different Channels type / basis
         devices[1],  # Different addressing channels
     ):
         with pytest.raises(
             TypeError,
             match="No match for channel ising with the"
-            + " right basis and addressing.",
+            + " right type, basis and addressing.",
         ):
             seq.switch_device(dev_)