# Copyright 2020 Pulser Development Team
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""The Sequence class, where a pulse sequence is defined."""

from __future__ import annotations

from collections import namedtuple
from collections.abc import Callable, Generator, Iterable, Set
import copy
from functools import wraps
from itertools import chain
import json
from sys import version_info
from typing import Any, Dict, cast, NamedTuple, Optional, Tuple, Union
# from typing_extensions import ParamSpecKwargs
import warnings

import numpy as np
from numpy.ma import diff
from numpy.typing import ArrayLike
from qutip.states import projection

from braket.analog.quera.lattice import Lattice
from braket.analog.quera.channels import Channel, Hyperfine, Rydberg
from braket.analog.quera.waveforms import Waveform, ConstantWaveform, CustomWaveform
from braket.analog.quera._seq_drawer import draw_sequence

# import pulser
# from pulser.channels import Channel
# from pulser.devices import MockDevice
# from pulser.devices._device_datacls import Device
# from pulser.json.coders import PulserEncoder, PulserDecoder
# from pulser.json.utils import obj_to_dict
# from pulser.parametrized import Parametrized, Variable
# from pulser.pulse import Pulse
# from pulser.register import Register


# if version_info[:2] >= (3, 8):  # pragma: no cover
#     from typing import Literal, get_args
# else:  # pragma: no cover
#     try:
#         from typing_extensions import Literal, get_args  # type: ignore
#     except ImportError:
#         raise ImportError(
#             "Using pulser with Python version 3.7 requires the"
#             " `typing_extensions` module. Install it by running"
#             " `pip install typing-extensions`."
#         )


QubitId = Union[int, str]

class _TimeSlot(NamedTuple):
    """Auxiliary class to store the information in the schedule."""

    type: Union[Tuple, str] # Union[Pulse, str]
    # Ω: Waveform
    # Δ: Waveform
    # ϕ: Waveform
    ti: int
    tf: int
    targets: Set[QubitId]

# Encodes a sequence building calls
_Call = namedtuple("_Call", ["name", "args", "kwargs"])

class Sequence:
    def __init__(self, lattice: Lattice):
        """Initializes a new pulse sequence."""
        if not isinstance(lattice, Lattice):
            raise TypeError("'lattice' must be of type 'Lattice'.")

        self._lattice: Lattice = lattice
        self._channels: dict[str, Channel] = {}
        self._qids: set[QubitId] = set(self.qubit_info.keys())
        # self._duration: int = 0
        self._schedule: dict[str, list[_TimeSlot]] = {}
        # self._calls: list[_Call] = [_Call("__init__", (lattice), {})]
        # Last time a target happened, by channel
        self._last_target: dict[str, int] = {}

    @property
    def qubit_info(self) -> dict[QubitId, np.ndarray]:
        """Dictionary with the qubit's IDs and positions."""
        return self._lattice.qubits

    @property
    def declared_channels(self) -> dict[str, Channel]:
        """Channels declared in this Sequence."""
        return dict(self._channels)

    def declare_channel(
        self,
        channel_name: str,
        channel_id: str,
        initial_target: Optional[Union[Iterable[QubitId],QubitId]] = None,
    ) -> None:
        """Declares a new channel to the Sequence.

        The first declared channel implicitly defines the sequence's mode of
        operation (i.e. the underlying Hamiltonian). In particular, if the
        first declared channel is of type ``Microwave``, the sequence will work
        in "XY Mode" and will not allow declaration of channels that do not
        address the 'XY' basis. Inversely, declaration of a channel of another
        type will block the declaration of ``Microwave`` channels.

        Note:
            Regular devices only allow a channel to be declared once, but
            ``MockDevice`` channels can be repeatedly declared if needed.

        Args:
            channel_name (str): Unique channel_name for the channel in the sequence.
            channel_id (str): How the channel is identified in the device.
                Consult ``Sequence.available_channels`` to see which channel
                ID's are still available and the associated channel's
                description.
            initial_target (Optional[Union[int, str, Iterable]]): For 'Local'
                addressing channels only. Declares the initial target of the
                channel. If left as None, the initial target will have to be
                set manually as the first addition to this channel.
        """
        if channel_name in self._channels:
            raise ValueError("The given channel name is already in use.")

        # if channel_id not in self._device.channels:
        if channel_id not in ['Hyperfine', 'Rydberg']:
            raise ValueError("No channel %s in the device." % channel_id)
        
        ch = Rydberg() if channel_id == "Rydberg" else Hyperfine()

        self._channels[channel_name] = ch
        self._schedule[channel_name] = []
        self._last_target[channel_name] = 0

        if initial_target == None:
            self._add_to_schedule(channel_name, _TimeSlot("target", -1, 0, self._qids))
        else:
            self._target(
                cast(Union[Iterable, QubitId], initial_target), channel_name
            )

        # # Manually store the channel declaration as a regular call
        # self._calls.append(
        #     _Call(
        #         "declare_channel",
        #         (channel_name, channel_id),
        #         {"initial_target": initial_target},
        #     )
        # )

    def _validate_channel(self, channel_name: str) -> None:
        if channel_name not in self._channels:
            raise ValueError("Use the channel name of a declared channel.")

    def _add_to_schedule(self, channel_name: str, timeslot: _TimeSlot) -> None:
        if hasattr(self, "_measurement"):
            raise RuntimeError(
                "The sequence has already been measured. "
                "Nothing more can be added."
            )
        self._schedule[channel_name].append(timeslot)

    def _delay(self, duration: int, channel: str) -> None:
        self._validate_channel(channel)
        # if self.is_parametrized():
        #     return

        last = self._last(channel)
        ti = last.tf
        tf = ti + self._channels[channel].validate_duration(duration)
        self._add_to_schedule(
            channel, _TimeSlot("delay", ti, tf, last.targets)
        )

    def _last(self, channel: str) -> _TimeSlot:
        """Shortcut to last element in the channel's schedule."""
        try:
            return self._schedule[channel][-1]
        except IndexError:
            raise ValueError("The chosen channel has no target.")

    def _target(
        self, qubits: Union[Iterable[QubitId], QubitId], channel: str
    ) -> None:
        self._validate_channel(channel)

        try:
            qubits_set = (
                set(cast(Iterable, qubits))
                if not isinstance(qubits, str)
                else {qubits}
            )
        except TypeError:
            qubits_set = {qubits}

        if len(qubits_set) > cast(int, self._channels[channel].max_targets):
            raise ValueError(
                "This channel can target at most "
                f"{self._channels[channel].max_targets} qubits at a time"
            )

        if not qubits_set.issubset(self._qids):
            raise ValueError("All given qubits must belong to the register.")

        try:
            last = self._last(channel)
            if last.targets == qubits_set:
                return
            ti = last.tf
            retarget = cast(int, self._channels[channel].retarget_time)
            elapsed = ti - self._last_target[channel]
            delta = cast(int, np.clip(retarget - elapsed, 0, retarget))
            if delta != 0:
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore")
                    delta = self._channels[channel].validate_duration(
                        16 if delta < 16 else delta
                    )
            tf = ti + delta

        except ValueError:
            ti = -1
            tf = 0

        self._last_target[channel] = tf
        self._add_to_schedule(channel, _TimeSlot("target", ti, tf, qubits_set))

    def _get_time_index(self, t: Union[float, int], t_tol: float = 1.0e-3) -> int:
        duration = self.get_duration()
        sim_times = np.array(range(duration+1))

        diff = np.where(abs(t - sim_times) < t_tol)[0]
        if len(diff)==0:
            raise IndexError(
                f"Given time {t} is absent from Simulation times within"
                + f" tolerance {t_tol}."
            )
        else:
            return diff[0]

    def _min_pulse_duration(self) -> float:
        duration_list = []
        for ch_schedule in self._schedule.values():
            for slot in ch_schedule:
                if isinstance(slot.type, tuple):
                    duration_list.append(slot.tf - slot.ti)
        return min(duration_list)

    def add(
        self,
        Ω: Union[Waveform, float, int], 
        Δ: Union[Waveform, float, int], 
        ϕ: Union[Waveform, float, int], 
        channel_name: str,
        duration: float = None, 
        protocol: str = "min-delay",
    ) -> None:

        # Validate the input
        durations = [None, None, None]
        if isinstance(Ω, Waveform): durations[0] = Ω.duration
        if isinstance(Δ, Waveform): durations[1] = Δ.duration
        if isinstance(ϕ, Waveform): durations[2] = ϕ.duration

        durations2 = [x for x in durations if x is not None]

        if len(set(durations2))==0 and duration is None:
            raise ValueError("If the inputs Ω, Δ and ϕ are all numbers, the duration of the waveforms needs to be provided")
        elif len(set(durations2)) >=2:
            raise ValueError("The durations of the input waveforms should be equal to each other")
        else:
            if len(set(durations2))>0: duration = durations2[0]
            if durations[0] == None: Ω = ConstantWaveform(duration, Ω)
            if durations[1] == None: Δ = ConstantWaveform(duration, Δ)
            if durations[2] == None: ϕ = ConstantWaveform(duration, ϕ)
        
        channel_name = cast(str, channel_name)
        self._validate_channel(channel_name)

        valid_protocols = ["min-delay", "no-delay", "wait-for-all"]
        if protocol not in valid_protocols:
            raise ValueError(
                f"Invalid protocol '{protocol}', only accepts protocols: "
                + ", ".join(valid_protocols)
            )


        last_TimeSlot = self._last(channel_name)
        t0 = last_TimeSlot.tf  # Preliminary ti

        # basis = channel_obj.basis
        # phase_barriers = [
        #     self._phase_ref[basis][q].last_time for q in last_TimeSlot.targets
        # ]
        current_max_t = t0 # max(t0, *phase_barriers)
        if protocol != "no-delay":
            for ch, _ in self._schedule.items():
                if ch == channel_name:
                    continue
                for op in self._schedule[ch][::-1]:
                    if op.tf <= current_max_t:
                        break
                    # if not isinstance(op.type, Pulse):
                    if not isinstance(op.type, Tuple):
                        continue
                    if op.targets & last_TimeSlot.targets or protocol == "wait-for-all":
                        current_max_t = op.tf
                        break

        # figure out ti,                     
        ti = current_max_t
        if ti > t0:
            # Insert a delay
            delay_duration = ti - t0

            # Delay must not be shorter than the min duration for this channel_name
            min_duration = self._channels[channel_name].min_duration
            if delay_duration < min_duration:
                ti += min_duration - delay_duration
                delay_duration = min_duration

            self._delay(delay_duration, channel_name)

        # figure out tf, 
        tf = ti + Ω.duration # pulse.duration

        # prs = {self._phase_ref[basis][q].last_phase for q in last.targets}
        # if len(prs) != 1:
        #     raise ValueError(
        #         "Cannot do a multiple-target pulse on qubits with different "
        #         "phase references for the same basis."
        #     )
        # else:
        #     phase_ref = prs.pop()

        # if phase_ref != 0:
        #     # Has to copy to keep the original pulse intact
        #     pulse = copy.deepcopy(pulse)
        #     pulse.phase = (pulse.phase + phase_ref) % (2 * np.pi)

        # figure out targets
        targets = last_TimeSlot.targets

        self._add_to_schedule(channel_name, _TimeSlot((Ω, Δ, ϕ), ti, tf, targets))

    def cut(self, t_initial: Union[float, int], t_final: Union[float, int], t_tol: float = 1.0e-3) -> Sequence:
        """Returns a portion of the given sequence between the given initial and final time.

        Keyword Args:
            channel (Optional[str]): A specific channel to return the duration
                of. If left as None, it will return the duration of the whole
                sequence.

        Returns:
            Sequence: The duration of the channel or sequence, in ns.
        """        
        ind_initial = self._get_time_index(t_initial)
        ind_final = self._get_time_index(t_final)
        duration = self.get_duration()

        if ind_initial >= ind_final:
            raise ValueError(f"The given t_initial {t_initial} is not strictly smaller than the given t_final {t_final} within tolerance {t_tol}")

        if ind_initial==duration:
            pass # This is an error which has been addressed above
        else:
            # seq = copy.deepcopy(self)
            lattice = self._lattice
            seq = Sequence(lattice)
            seq._channels = self.declared_channels

            for ch in seq.declared_channels.keys():
                seq._schedule[ch] = [self._schedule[ch][0]]
            
                intervals = []
                for timeslot in self._schedule[ch]:
                    intervals.append([timeslot.ti, timeslot.tf])
                
                # print(ch, intervals)

                if len(intervals)==1:
                    continue
                else:
                    t0 = seq._last(ch).tf
                    for ind, interval in enumerate(intervals[1:]): 
                        ind += 1
                        if interval[1]<=ind_initial or interval[0]>=ind_final:
                            pass
                        else:
                            # print("t0", t0)
                            timeslot = self._schedule[ch][ind]
                            timeslot_duration = timeslot.tf-timeslot.ti

                            # Determine the set of points in this interval to keep 
                            ti = max(ind_initial, timeslot.ti)
                            tf = min(ind_final, timeslot.tf)
                            
                            if ti==timeslot.ti:
                                ind2 = range(tf-ti)
                            elif tf == timeslot.tf:
                                ind2 = range(ti-(tf-timeslot_duration), timeslot_duration)
                            else:
                                ind2 = range(ti-timeslot.ti, tf-timeslot.ti)
                            
                            # print("ti, tf, timeslot.ti, timeslot.tf, ind2 ", ti, tf, timeslot.ti, timeslot.tf, ind2)

                            # shift the starting point to match the previous interval
                            tf = tf - (ti-t0)
                            ti = t0
                            t0 = tf 

                            if type(timeslot.type) == tuple:
                                Ω = CustomWaveform([timeslot.type[0].samples[i] for i in ind2])
                                Δ = CustomWaveform([timeslot.type[1].samples[i] for i in ind2])
                                ϕ = CustomWaveform([timeslot.type[2].samples[i] for i in ind2])
                                new_timeslot = _TimeSlot((Ω, Δ, ϕ), ti, tf, timeslot.targets)
                            else:
                                new_timeslot = _TimeSlot(timeslot.type, ti, tf, timeslot.targets)
                            seq._add_to_schedule(ch, new_timeslot)   

            return seq




    def get_duration(self, channel: Optional[str] = None) -> int:
        """Returns the current duration of a channel or the whole sequence.

        Keyword Args:
            channel (Optional[str]): A specific channel to return the duration
                of. If left as None, it will return the duration of the whole
                sequence.

        Returns:
            int: The duration of the channel or sequence, in ns.
        """
        if channel is None:
            durations = [
                self._last(ch).tf
                for ch in self._schedule
                if self._schedule[ch]
            ]
            return 0 if not durations else max(durations)

        self._validate_channel(channel)
        return self._last(channel).tf if self._schedule[channel] else 0

    def draw(
        self,
        draw_phase_area: bool = False,
        # draw_interp_pts: bool = True,
        # draw_phase_shifts: bool = False,
    ) -> None:
        """Draws the sequence in its current state.

        Keyword Args:
            draw_phase_area (bool): Whether phase and area values need to be
                shown as text on the plot, defaults to False.
            draw_interp_pts (bool): When the sequence has pulses with waveforms
                of type InterpolatedWaveform, draws the points of interpolation
                on top of the respective waveforms (defaults to True).
            draw_phase_shifts (bool): Whether phase shift and reference
                information should be added to the plot, defaults to False.

        See Also:
            Simulation.draw(): Draws the provided sequence and the one used by
            the solver.
        """
        draw_sequence(
            self,
            draw_phase_area=draw_phase_area,
            # draw_interp_pts=draw_interp_pts,
            # draw_phase_shifts=draw_phase_shifts,
        )

# Copyright 2020 Pulser Development Team
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""The Sequence class, where a pulse sequence is defined."""

from __future__ import annotations

from collections import namedtuple
from collections.abc import Callable, Generator, Iterable, Set
import copy
from functools import wraps
from itertools import chain
import json
from sys import version_info
from typing import Any, Dict, cast, NamedTuple, Optional, Tuple, Union
# from typing_extensions import ParamSpecKwargs
import warnings

import numpy as np
from numpy.ma import diff
from numpy.typing import ArrayLike
from qutip.states import projection

from braket.analog.quera.lattice import Lattice
from braket.analog.quera.channels import Channel, Hyperfine, Rydberg
from braket.analog.quera.waveforms import Waveform, ConstantWaveform, CustomWaveform
from braket.analog.quera._seq_drawer import draw_sequence

# import pulser
# from pulser.channels import Channel
# from pulser.devices import MockDevice
# from pulser.devices._device_datacls import Device
# from pulser.json.coders import PulserEncoder, PulserDecoder
# from pulser.json.utils import obj_to_dict
# from pulser.parametrized import Parametrized, Variable
# from pulser.pulse import Pulse
# from pulser.register import Register


# if version_info[:2] >= (3, 8):  # pragma: no cover
#     from typing import Literal, get_args
# else:  # pragma: no cover
#     try:
#         from typing_extensions import Literal, get_args  # type: ignore
#     except ImportError:
#         raise ImportError(
#             "Using pulser with Python version 3.7 requires the"
#             " `typing_extensions` module. Install it by running"
#             " `pip install typing-extensions`."
#         )


QubitId = Union[int, str]

class _TimeSlot(NamedTuple):
    """Auxiliary class to store the information in the schedule."""

    type: Union[Tuple, str] # Union[Pulse, str]
    # Ω: Waveform
    # Δ: Waveform
    # ϕ: Waveform
    ti: int
    tf: int
    targets: Set[QubitId]

# Encodes a sequence building calls
_Call = namedtuple("_Call", ["name", "args", "kwargs"])

class Sequence:
    def __init__(self, lattice: Lattice):
        """Initializes a new pulse sequence."""
        if not isinstance(lattice, Lattice):
            raise TypeError("'lattice' must be of type 'Lattice'.")

        self._lattice: Lattice = lattice
        self._channels: dict[str, Channel] = {}
        self._qids: set[QubitId] = set(self.qubit_info.keys())
        # self._duration: int = 0
        self._schedule: dict[str, list[_TimeSlot]] = {}
        # self._calls: list[_Call] = [_Call("__init__", (lattice), {})]
        # Last time a target happened, by channel
        self._last_target: dict[str, int] = {}

    @property
    def qubit_info(self) -> dict[QubitId, np.ndarray]:
        """Dictionary with the qubit's IDs and positions."""
        return self._lattice.qubits

    @property
    def declared_channels(self) -> dict[str, Channel]:
        """Channels declared in this Sequence."""
        return dict(self._channels)

    def declare_channel(
        self,
        channel_name: str,
        channel_id: str,
        initial_target: Optional[Union[Iterable[QubitId],QubitId]] = None,
    ) -> None:
        """Declares a new channel to the Sequence.

        The first declared channel implicitly defines the sequence's mode of
        operation (i.e. the underlying Hamiltonian). In particular, if the
        first declared channel is of type ``Microwave``, the sequence will work
        in "XY Mode" and will not allow declaration of channels that do not
        address the 'XY' basis. Inversely, declaration of a channel of another
        type will block the declaration of ``Microwave`` channels.

        Note:
            Regular devices only allow a channel to be declared once, but
            ``MockDevice`` channels can be repeatedly declared if needed.

        Args:
            channel_name (str): Unique channel_name for the channel in the sequence.
            channel_id (str): How the channel is identified in the device.
                Consult ``Sequence.available_channels`` to see which channel
                ID's are still available and the associated channel's
                description.
            initial_target (Optional[Union[int, str, Iterable]]): For 'Local'
                addressing channels only. Declares the initial target of the
                channel. If left as None, the initial target will have to be
                set manually as the first addition to this channel.
        """
        if channel_name in self._channels:
            raise ValueError("The given channel name is already in use.")

        # if channel_id not in self._device.channels:
        if channel_id not in ['Hyperfine', 'Rydberg']:
            raise ValueError("No channel %s in the device." % channel_id)
        
        ch = Rydberg() if channel_id == "Rydberg" else Hyperfine()

        self._channels[channel_name] = ch
        self._schedule[channel_name] = []
        self._last_target[channel_name] = 0

        if initial_target == None:
            self._add_to_schedule(channel_name, _TimeSlot("target", -1, 0, self._qids))
        else:
            self._target(
                cast(Union[Iterable, QubitId], initial_target), channel_name
            )

        # # Manually store the channel declaration as a regular call
        # self._calls.append(
        #     _Call(
        #         "declare_channel",
        #         (channel_name, channel_id),
        #         {"initial_target": initial_target},
        #     )
        # )

    def _validate_channel(self, channel_name: str) -> None:
        if channel_name not in self._channels:
            raise ValueError("Use the channel name of a declared channel.")

    def _add_to_schedule(self, channel_name: str, timeslot: _TimeSlot) -> None:
        if hasattr(self, "_measurement"):
            raise RuntimeError(
                "The sequence has already been measured. "
                "Nothing more can be added."
            )
        self._schedule[channel_name].append(timeslot)

    def _delay(self, duration: int, channel: str) -> None:
        self._validate_channel(channel)
        # if self.is_parametrized():
        #     return

        last = self._last(channel)
        ti = last.tf
        tf = ti + self._channels[channel].validate_duration(duration)
        self._add_to_schedule(
            channel, _TimeSlot("delay", ti, tf, last.targets)
        )

    def _last(self, channel: str) -> _TimeSlot:
        """Shortcut to last element in the channel's schedule."""
        try:
            return self._schedule[channel][-1]
        except IndexError:
            raise ValueError("The chosen channel has no target.")

    def _target(
        self, qubits: Union[Iterable[QubitId], QubitId], channel: str
    ) -> None:
        self._validate_channel(channel)

        try:
            qubits_set = (
                set(cast(Iterable, qubits))
                if not isinstance(qubits, str)
                else {qubits}
            )
        except TypeError:
            qubits_set = {qubits}

        if len(qubits_set) > cast(int, self._channels[channel].max_targets):
            raise ValueError(
                "This channel can target at most "
                f"{self._channels[channel].max_targets} qubits at a time"
            )

        if not qubits_set.issubset(self._qids):
            raise ValueError("All given qubits must belong to the register.")

        try:
            last = self._last(channel)
            if last.targets == qubits_set:
                return
            ti = last.tf
            retarget = cast(int, self._channels[channel].retarget_time)
            elapsed = ti - self._last_target[channel]
            delta = cast(int, np.clip(retarget - elapsed, 0, retarget))
            if delta != 0:
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore")
                    delta = self._channels[channel].validate_duration(
                        16 if delta < 16 else delta
                    )
            tf = ti + delta

        except ValueError:
            ti = -1
            tf = 0

        self._last_target[channel] = tf
        self._add_to_schedule(channel, _TimeSlot("target", ti, tf, qubits_set))

    def _get_time_index(self, t: Union[float, int], t_tol: float = 1.0e-3) -> int:
        duration = self.get_duration()
        sim_times = np.array(range(duration+1))

        diff = np.where(abs(t - sim_times) < t_tol)[0]
        if len(diff)==0:
            raise IndexError(
                f"Given time {t} is absent from Simulation times within"
                + f" tolerance {t_tol}."
            )
        else:
            return diff[0]

    def _min_pulse_duration(self) -> float:
        duration_list = []
        for ch_schedule in self._schedule.values():
            for slot in ch_schedule:
                if isinstance(slot.type, tuple):
                    duration_list.append(slot.tf - slot.ti)
        return min(duration_list)

    def add(
        self,
        Ω: Union[Waveform, float, int], 
        Δ: Union[Waveform, float, int], 
        ϕ: Union[Waveform, float, int], 
        channel_name: str,
        duration: float = None, 
        protocol: str = "min-delay",
    ) -> None:

        # Validate the input
        durations = [None, None, None]
        if isinstance(Ω, Waveform): durations[0] = Ω.duration
        if isinstance(Δ, Waveform): durations[1] = Δ.duration
        if isinstance(ϕ, Waveform): durations[2] = ϕ.duration

        durations2 = [x for x in durations if x is not None]

        if len(set(durations2))==0 and duration is None:
            raise ValueError("If the inputs Ω, Δ and ϕ are all numbers, the duration of the waveforms needs to be provided")
        elif len(set(durations2)) >=2:
            raise ValueError("The durations of the input waveforms should be equal to each other")
        else:
            if len(set(durations2))>0: duration = durations2[0]
            if durations[0] == None: Ω = ConstantWaveform(duration, Ω)
            if durations[1] == None: Δ = ConstantWaveform(duration, Δ)
            if durations[2] == None: ϕ = ConstantWaveform(duration, ϕ)
        
        channel_name = cast(str, channel_name)
        self._validate_channel(channel_name)

        valid_protocols = ["min-delay", "no-delay", "wait-for-all"]
        if protocol not in valid_protocols:
            raise ValueError(
                f"Invalid protocol '{protocol}', only accepts protocols: "
                + ", ".join(valid_protocols)
            )


        last_TimeSlot = self._last(channel_name)
        t0 = last_TimeSlot.tf  # Preliminary ti

        # basis = channel_obj.basis
        # phase_barriers = [
        #     self._phase_ref[basis][q].last_time for q in last_TimeSlot.targets
        # ]
        current_max_t = t0 # max(t0, *phase_barriers)
        if protocol != "no-delay":
            for ch, _ in self._schedule.items():
                if ch == channel_name:
                    continue
                for op in self._schedule[ch][::-1]:
                    if op.tf <= current_max_t:
                        break
                    # if not isinstance(op.type, Pulse):
                    if not isinstance(op.type, Tuple):
                        continue
                    if op.targets & last_TimeSlot.targets or protocol == "wait-for-all":
                        current_max_t = op.tf
                        break

        # figure out ti,                     
        ti = current_max_t
        if ti > t0:
            # Insert a delay
            delay_duration = ti - t0

            # Delay must not be shorter than the min duration for this channel_name
            min_duration = self._channels[channel_name].min_duration
            if delay_duration < min_duration:
                ti += min_duration - delay_duration
                delay_duration = min_duration

            self._delay(delay_duration, channel_name)

        # figure out tf, 
        tf = ti + Ω.duration # pulse.duration

        # prs = {self._phase_ref[basis][q].last_phase for q in last.targets}
        # if len(prs) != 1:
        #     raise ValueError(
        #         "Cannot do a multiple-target pulse on qubits with different "
        #         "phase references for the same basis."
        #     )
        # else:
        #     phase_ref = prs.pop()

        # if phase_ref != 0:
        #     # Has to copy to keep the original pulse intact
        #     pulse = copy.deepcopy(pulse)
        #     pulse.phase = (pulse.phase + phase_ref) % (2 * np.pi)

        # figure out targets
        targets = last_TimeSlot.targets

        self._add_to_schedule(channel_name, _TimeSlot((Ω, Δ, ϕ), ti, tf, targets))

    def cut(self, t_initial: Union[float, int], t_final: Union[float, int], t_tol: float = 1.0e-3) -> Sequence:
        """Returns a portion of the given sequence between the given initial and final time.

        Keyword Args:
            channel (Optional[str]): A specific channel to return the duration
                of. If left as None, it will return the duration of the whole
                sequence.

        Returns:
            Sequence: The duration of the channel or sequence, in ns.
        """        
        ind_initial = self._get_time_index(t_initial)
        ind_final = self._get_time_index(t_final)
        duration = self.get_duration()

        if ind_initial >= ind_final:
            raise ValueError(f"The given t_initial {t_initial} is not strictly smaller than the given t_final {t_final} within tolerance {t_tol}")

        if ind_initial==duration:
            pass # This is an error which has been addressed above
        else:
            # seq = copy.deepcopy(self)
            lattice = self._lattice
            seq = Sequence(lattice)
            seq._channels = self.declared_channels

            for ch in seq.declared_channels.keys():
                seq._schedule[ch] = [self._schedule[ch][0]]
            
                intervals = []
                for timeslot in self._schedule[ch]:
                    intervals.append([timeslot.ti, timeslot.tf])
                
                # print(ch, intervals)

                if len(intervals)==1:
                    continue
                else:
                    t0 = seq._last(ch).tf
                    for ind, interval in enumerate(intervals[1:]): 
                        ind += 1
                        if interval[1]<=ind_initial or interval[0]>=ind_final:
                            pass
                        else:
                            # print("t0", t0)
                            timeslot = self._schedule[ch][ind]
                            timeslot_duration = timeslot.tf-timeslot.ti

                            # Determine the set of points in this interval to keep 
                            ti = max(ind_initial, timeslot.ti)
                            tf = min(ind_final, timeslot.tf)
                            
                            if ti==timeslot.ti:
                                ind2 = range(tf-ti)
                            elif tf == timeslot.tf:
                                ind2 = range(ti-(tf-timeslot_duration), timeslot_duration)
                            else:
                                ind2 = range(ti-timeslot.ti, tf-timeslot.ti)
                            
                            # print("ti, tf, timeslot.ti, timeslot.tf, ind2 ", ti, tf, timeslot.ti, timeslot.tf, ind2)

                            # shift the starting point to match the previous interval
                            tf = tf - (ti-t0)
                            ti = t0
                            t0 = tf 

                            if type(timeslot.type) == tuple:
                                Ω = CustomWaveform([timeslot.type[0].samples[i] for i in ind2])
                                Δ = CustomWaveform([timeslot.type[1].samples[i] for i in ind2])
                                ϕ = CustomWaveform([timeslot.type[2].samples[i] for i in ind2])
                                new_timeslot = _TimeSlot((Ω, Δ, ϕ), ti, tf, timeslot.targets)
                            else:
                                new_timeslot = _TimeSlot(timeslot.type, ti, tf, timeslot.targets)
                            seq._add_to_schedule(ch, new_timeslot)   

            return seq




    def get_duration(self, channel: Optional[str] = None) -> int:
        """Returns the current duration of a channel or the whole sequence.

        Keyword Args:
            channel (Optional[str]): A specific channel to return the duration
                of. If left as None, it will return the duration of the whole
                sequence.

        Returns:
            int: The duration of the channel or sequence, in ns.
        """
        if channel is None:
            durations = [
                self._last(ch).tf
                for ch in self._schedule
                if self._schedule[ch]
            ]
            return 0 if not durations else max(durations)

        self._validate_channel(channel)
        return self._last(channel).tf if self._schedule[channel] else 0

    def draw(
        self,
        draw_phase_area: bool = False,
        # draw_interp_pts: bool = True,
        # draw_phase_shifts: bool = False,
    ) -> None:
        """Draws the sequence in its current state.

        Keyword Args:
            draw_phase_area (bool): Whether phase and area values need to be
                shown as text on the plot, defaults to False.
            draw_interp_pts (bool): When the sequence has pulses with waveforms
                of type InterpolatedWaveform, draws the points of interpolation
                on top of the respective waveforms (defaults to True).
            draw_phase_shifts (bool): Whether phase shift and reference
                information should be added to the plot, defaults to False.

        See Also:
            Simulation.draw(): Draws the provided sequence and the one used by
            the solver.
        """
        draw_sequence(
            self,
            draw_phase_area=draw_phase_area,
            # draw_interp_pts=draw_interp_pts,
            # draw_phase_shifts=draw_phase_shifts,
        )

# Copyright 2020 Pulser Development Team # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # #    http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. The Sequence class, where a pulse sequence is defined. # from typing_extensions import ParamSpecKwargs # import pulser # from pulser.channels import Channel # from pulser.devices import MockDevice # from pulser.devices._device_datacls import Device # from pulser.json.coders import PulserEncoder, PulserDecoder # from pulser.json.utils import obj_to_dict # from pulser.parametrized import Parametrized, Variable # from pulser.pulse import Pulse # from pulser.register import Register # if version_info[:2] >= (3, 8):  # pragma: no cover #     from typing import Literal, get_args # else:  # pragma: no cover #     try: #         from typing_extensions import Literal, get_args  # type: ignore #     except ImportError: #         raise ImportError( #             "Using pulser with Python version 3.7 requires the" #             " `typing_extensions` module. Install it by running" #             " `pip install typing-extensions`." #         ) Auxiliary class to store the information in the schedule. # Union[Pulse, str] # Ω: Waveform # Δ: Waveform # ϕ: Waveform # Encodes a sequence building calls Initializes a new pulse sequence. # self._duration: int = 0 # self._calls: list[_Call] = [_Call("__init__", (lattice), {})] # Last time a target happened, by channel Dictionary with the qubit's IDs and positions. Channels declared in this Sequence. Declares a new channel to the Sequence.          The first declared channel implicitly defines the sequence's mode of         operation (i.e. the underlying Hamiltonian). In particular, if the         first declared channel is of type ``Microwave``, the sequence will work         in "XY Mode" and will not allow declaration of channels that do not         address the 'XY' basis. Inversely, declaration of a channel of another         type will block the declaration of ``Microwave`` channels.          Note:             Regular devices only allow a channel to be declared once, but             ``MockDevice`` channels can be repeatedly declared if needed.          Args:             channel_name (str): Unique channel_name for the channel in the sequence.             channel_id (str): How the channel is identified in the device.                 Consult ``Sequence.available_channels`` to see which channel                 ID's are still available and the associated channel's                 description.             initial_target (Optional[Union[int, str, Iterable]]): For 'Local'                 addressing channels only. Declares the initial target of the                 channel. If left as None, the initial target will have to be                 set manually as the first addition to this channel. # if channel_id not in self._device.channels: # # Manually store the channel declaration as a regular call # self._calls.append( #     _Call( #         "declare_channel", #         (channel_name, channel_id), #         {"initial_target": initial_target}, #     ) # ) # if self.is_parametrized(): #     return Shortcut to last element in the channel's schedule. # Validate the input # Preliminary ti # basis = channel_obj.basis # phase_barriers = [ #     self._phase_ref[basis][q].last_time for q in last_TimeSlot.targets # ] # max(t0, *phase_barriers) # if not isinstance(op.type, Pulse): # figure out ti, # Insert a delay # Delay must not be shorter than the min duration for this channel_name # figure out tf, # pulse.duration # prs = {self._phase_ref[basis][q].last_phase for q in last.targets} # if len(prs) != 1: #     raise ValueError( #         "Cannot do a multiple-target pulse on qubits with different " #         "phase references for the same basis." #     ) # else: #     phase_ref = prs.pop() # if phase_ref != 0: #     # Has to copy to keep the original pulse intact #     pulse = copy.deepcopy(pulse) #     pulse.phase = (pulse.phase + phase_ref) % (2 * np.pi) # figure out targets Returns a portion of the given sequence between the given initial and final time.          Keyword Args:             channel (Optional[str]): A specific channel to return the duration                 of. If left as None, it will return the duration of the whole                 sequence.          Returns:             Sequence: The duration of the channel or sequence, in ns. # This is an error which has been addressed above # seq = copy.deepcopy(self) # print(ch, intervals) # print("t0", t0) # Determine the set of points in this interval to keep # print("ti, tf, timeslot.ti, timeslot.tf, ind2 ", ti, tf, timeslot.ti, timeslot.tf, ind2) # shift the starting point to match the previous interval Returns the current duration of a channel or the whole sequence.          Keyword Args:             channel (Optional[str]): A specific channel to return the duration                 of. If left as None, it will return the duration of the whole                 sequence.          Returns:             int: The duration of the channel or sequence, in ns. # draw_interp_pts: bool = True, # draw_phase_shifts: bool = False, Draws the sequence in its current state.          Keyword Args:             draw_phase_area (bool): Whether phase and area values need to be                 shown as text on the plot, defaults to False.             draw_interp_pts (bool): When the sequence has pulses with waveforms                 of type InterpolatedWaveform, draws the points of interpolation                 on top of the respective waveforms (defaults to True).             draw_phase_shifts (bool): Whether phase shift and reference                 information should be added to the plot, defaults to False.          See Also:             Simulation.draw(): Draws the provided sequence and the one used by             the solver. # draw_interp_pts=draw_interp_pts, # draw_phase_shifts=draw_phase_shifts,