# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file 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.

from typing import List, Union
from unittest.mock import Mock

import numpy as np
import pytest
from openpulse import ast
from oqpy import IntVar

from braket.pulse import ArbitraryWaveform, ConstantWaveform, DragGaussianWaveform, GaussianWaveform
from braket.pulse.ast.approximation_parser import _ApproximationParser
from braket.pulse.frame import Frame
from braket.pulse.port import Port
from braket.pulse.pulse_sequence import PulseSequence
from braket.timings.time_series import TimeSeries, _all_close


@pytest.fixture
def port():
    return Port(port_id="device_port_x0", dt=1e-9, properties={})


def test_bare_pulsequence():
    pulse_seq = PulseSequence()
    results = pulse_seq.to_time_trace()
    verify_results(results, {}, {}, {})


def test_delay(port):
    frame = Frame(frame_id="frame1", port=port, frequency=1e8, phase=0, is_predefined=False)
    pulse_seq = PulseSequence().delay(frame, 3e-9)

    expected_amplitudes = {"frame1": TimeSeries()}
    expected_frequencies = {"frame1": TimeSeries()}
    expected_phases = {"frame1": TimeSeries()}

    expected_amplitudes["frame1"].put(0, 0).put(2e-9, 0)
    expected_frequencies["frame1"].put(0, 1e8).put(2e-9, 1e8)
    expected_phases["frame1"].put(0, 0).put(2e-9, 0)

    parser = _ApproximationParser(program=pulse_seq._program, frames=to_dict(frame))

    verify_results(parser, expected_amplitudes, expected_frequencies, expected_phases)


def test_predefined_frame(port):
    frame = Frame(frame_id="frame1", port=port, frequency=1e8, phase=0, is_predefined=True)
    pulse_seq = PulseSequence().delay(frame, 3e-9)

    expected_amplitudes = {"frame1": TimeSeries()}
    expected_frequencies = {"frame1": TimeSeries()}
    expected_phases = {"frame1": TimeSeries()}

    expected_amplitudes["frame1"].put(0, 0).put(2e-9, 0)
    expected_frequencies["frame1"].put(0, 1e8).put(2e-9, 1e8)
    expected_phases["frame1"].put(0, 0).put(2e-9, 0)

    for statement in pulse_seq._program._state.body:
        assert type(statement) != ast.FrameType

    parser = _ApproximationParser(program=pulse_seq._program, frames=to_dict(frame))

    verify_results(parser, expected_amplitudes, expected_frequencies, expected_phases)


def test_set_shift_phase(port):
    frame = Frame(frame_id="frame1", port=port, frequency=1e8, phase=0, is_predefined=False)
    pulse_seq = (
        PulseSequence()
        .set_phase(frame, 1)
        .delay(frame, 2e-9)
        .shift_phase(frame, 2)
        .delay(frame, 5e-9)
        .set_phase(frame, 0)
    )
    expected_amplitudes = {"frame1": TimeSeries()}
    expected_frequencies = {"frame1": TimeSeries()}
    expected_phases = {"frame1": TimeSeries()}

    expected_amplitudes["frame1"].put(0, 0).put(1e-9, 0)
    expected_frequencies["frame1"].put(0, 1e8).put(1e-9, 1e8)
    expected_phases["frame1"].put(0, 1).put(1e-9, 1)

    expected_amplitudes["frame1"].put(2e-9, 0).put(6e-9, 0)
    expected_frequencies["frame1"].put(2e-9, 1e8).put(6e-9, 1e8)
    expected_phases["frame1"].put(2e-9, 3).put(6e-9, 3)

    parser = _ApproximationParser(program=pulse_seq._program, frames=to_dict(frame))

    verify_results(parser, expected_amplitudes, expected_frequencies, expected_phases)


def test_set_shift_phase_beyond_2_pi(port):
    frame = Frame(frame_id="frame1", port=port, frequency=1e8, phase=0, is_predefined=False)
    pulse_seq = (
        PulseSequence()
        .set_phase(frame, 5 * np.pi / 2)
        .delay(frame, 2e-9)
        .shift_phase(frame, -np.pi)
        .delay(frame, 5e-9)
        .set_phase(frame, 0)
    )
    expected_amplitudes = {"frame1": TimeSeries()}
    expected_frequencies = {"frame1": TimeSeries()}
    expected_phases = {"frame1": TimeSeries()}

    # 5pi/2 is reduced to pi/2
    expected_amplitudes["frame1"].put(0, 0).put(1e-9, 0)
    expected_frequencies["frame1"].put(0, 1e8).put(1e-9, 1e8)
    expected_phases["frame1"].put(0, np.pi / 2).put(1e-9, np.pi / 2)

    # shift_phase adds -pi to the phase of the last point -> 3pi/2
    expected_amplitudes["frame1"].put(2e-9, 0).put(6e-9, 0)
    expected_frequencies["frame1"].put(2e-9, 1e8).put(6e-9, 1e8)
    expected_phases["frame1"].put(2e-9, 3 * np.pi / 2).put(6e-9, 3 * np.pi / 2)

    parser = _ApproximationParser(program=pulse_seq._program, frames=to_dict(frame))

    verify_results(parser, expected_amplitudes, expected_frequencies, expected_phases)


def test_set_shift_frequency(port):
    frame = Frame(frame_id="frame1", port=port, frequency=1e8, phase=0, is_predefined=False)
    pulse_seq = (
        PulseSequence()
        .set_frequency(frame, 2e8)
        .delay(frame, 20e-9)
        .shift_frequency(frame, -0.1e8)
        .delay(frame, 10e-9)
    )

    expected_amplitudes = {"frame1": TimeSeries()}
    expected_frequencies = {"frame1": TimeSeries()}
    expected_phases = {"frame1": TimeSeries()}

    expected_amplitudes["frame1"].put(0, 0).put(19e-9, 0)
    expected_frequencies["frame1"].put(0, 2e8).put(19e-9, 2e8)
    expected_phases["frame1"].put(0, 0).put(19e-9, 0)

    expected_amplitudes["frame1"].put(20e-9, 0).put(29e-9, 0)
    expected_frequencies["frame1"].put(20e-9, 1.9e8).put(29e-9, 1.9e8)
    expected_phases["frame1"].put(20e-9, 0).put(29e-9, 0)

    parser = _ApproximationParser(program=pulse_seq._program, frames=to_dict(frame))

    verify_results(parser, expected_amplitudes, expected_frequencies, expected_phases)


def test_play_arbitrary_waveforms(port):
    frame = Frame(frame_id="frame1", port=port, frequency=1e8, phase=0, is_predefined=False)
    my_arb_wf = ArbitraryWaveform([0.4 + 0.1j, -0.8 + 0.1j, 1 + 0.2j])
    pulse_seq = PulseSequence().play(frame, my_arb_wf).capture_v0(frame)

    expected_amplitudes = {"frame1": TimeSeries()}
    expected_frequencies = {"frame1": TimeSeries()}
    expected_phases = {"frame1": TimeSeries()}

    times = np.arange(0, 4e-9, port.dt)
    for t, v in zip(times, [0.4 + 0.1j, -0.8 + 0.1j, 1 + 0.2j]):
        expected_amplitudes["frame1"].put(t, v)
        expected_frequencies["frame1"].put(t, 1e8)
        expected_phases["frame1"].put(t, 0)

    parser = _ApproximationParser(program=pulse_seq._program, frames=to_dict(frame))

    verify_results(parser, expected_amplitudes, expected_frequencies, expected_phases)


def test_play_literal(port):
    frame = Frame(frame_id="frame1", port=port, frequency=1e8, phase=0, is_predefined=False)
    pulse_seq = PulseSequence()
    pulse_seq._program.play(frame=frame, waveform=[0.4 + 0.1j, -0.8 + 0.1j, 1 + 0.2j])

    expected_amplitudes = {"frame1": TimeSeries()}
    expected_frequencies = {"frame1": TimeSeries()}
    expected_phases = {"frame1": TimeSeries()}

    times = np.arange(0, 4e-9, port.dt)
    for t, v in zip(times, [0.4 + 0.1j, -0.8 + 0.1j, 1 + 0.2j]):
        expected_amplitudes["frame1"].put(t, v)
        expected_frequencies["frame1"].put(t, 1e8)
        expected_phases["frame1"].put(t, 0)

    parser = _ApproximationParser(program=pulse_seq._program, frames=to_dict(frame))

    verify_results(parser, expected_amplitudes, expected_frequencies, expected_phases)


def test_classical_variable_declaration(port):
    frame = Frame(frame_id="frame1", port=port, frequency=1e8, phase=0, is_predefined=False)
    pulse_seq = PulseSequence()
    i = IntVar(5, "i")
    pulse_seq._program.increment(i, 1)

    with pytest.raises(NotImplementedError):
        _ApproximationParser(program=pulse_seq._program, frames=to_dict(frame))


def test_play_constant_waveforms(port):
    frame = Frame(frame_id="frame1", port=port, frequency=1e8, phase=0, is_predefined=False)
    length = 6e-9
    pulse_seq = PulseSequence().play(frame, ConstantWaveform(length, 0.75))

    expected_amplitudes = {"frame1": TimeSeries()}
    expected_frequencies = {"frame1": TimeSeries()}
    expected_phases = {"frame1": TimeSeries()}

    times = np.arange(0, length, port.dt)
    values = 0.75 * np.ones_like(times)
    for t, v in zip(times, values):
        expected_amplitudes["frame1"].put(t, v)
        expected_frequencies["frame1"].put(t, 1e8)
        expected_phases["frame1"].put(t, 0)

    parser = _ApproximationParser(program=pulse_seq._program, frames=to_dict(frame))

    verify_results(parser, expected_amplitudes, expected_frequencies, expected_phases)


def test_set_scale(port):
    frame = Frame(frame_id="frame1", port=port, frequency=1e8, phase=0, is_predefined=False)
    pulse_seq = (
        PulseSequence()
        .play(frame, ConstantWaveform(10e-9, 0.66))
        .set_scale(frame, 0.5)
        .play(frame, ConstantWaveform(20e-9, 0.66))
    )

    expected_amplitudes = {"frame1": TimeSeries()}
    expected_frequencies = {"frame1": TimeSeries()}
    expected_phases = {"frame1": TimeSeries()}

    times = np.arange(0, 10e-9, port.dt)
    values = 0.66 * np.ones_like(times)
    for t, v in zip(times, values):
        expected_amplitudes["frame1"].put(t, v)
        expected_frequencies["frame1"].put(t, 1e8)
        expected_phases["frame1"].put(t, 0)

    shift_time = 10e-9
    times = np.arange(0, 20e-9, port.dt)
    values = 0.33 * np.ones_like(times)
    for t, v in zip(times, values):
        expected_amplitudes["frame1"].put(shift_time + t, v)
        expected_frequencies["frame1"].put(shift_time + t, 1e8)
        expected_phases["frame1"].put(shift_time + t, 0)

    parser = _ApproximationParser(program=pulse_seq._program, frames=to_dict(frame))

    verify_results(parser, expected_amplitudes, expected_frequencies, expected_phases)


def test_play_gaussian_waveforms(port):
    frame1 = Frame(frame_id="frame1", port=port, frequency=1e8, phase=0, is_predefined=False)

    # First gaussian with zero_at_edges=False
    gaussian_wf_ZaE_False = GaussianWaveform(
        length=1e-8, sigma=1.69e-9, amplitude=1.0, zero_at_edges=False
    )
    pulse_seq = PulseSequence().play(frame1, gaussian_wf_ZaE_False)

    times = np.arange(0, 1e-8, port.dt)
    values = np.array(
        [
            complex(0.012568049266111367),
            complex(0.06074792381889125),
            complex(0.20688853998666168),
            complex(0.49645839613814063),
            complex(0.839403382587157),
            complex(1.0),
            complex(0.839403382587157),
            complex(0.49645839613814063),
            complex(0.20688853998666168),
            complex(0.06074792381889125),
        ],
        dtype=np.complex128,
    )

    expected_amplitudes = {"frame1": TimeSeries()}
    expected_frequencies = {"frame1": TimeSeries()}
    expected_phases = {"frame1": TimeSeries()}

    for t, v in zip(times, values):
        expected_amplitudes["frame1"].put(t, v)
        expected_frequencies["frame1"].put(t, 1e8)
        expected_phases["frame1"].put(t, 0)

    # Add a delay between the two waveforms
    pulse_seq.delay(frame1, 2e-8)
    expected_amplitudes["frame1"].put(1e-8, 0).put(29e-9, 0)
    expected_frequencies["frame1"].put(1e-8, 1e8).put(29e-9, 1e8)
    expected_phases["frame1"].put(1e-8, 0).put(29e-9, 0)

    # Second gaussian with zero_at_edges=True
    gaussian_wf_ZaE_True = GaussianWaveform(
        length=1e-8, sigma=1.69e-9, amplitude=1.0, zero_at_edges=True
    )
    pulse_seq.play(frame1, gaussian_wf_ZaE_True)

    times = np.arange(0, 1e-8, port.dt)
    values = np.array(
        [
            complex(0.0),
            complex(0.04879310874736347),
            complex(0.1967938049565093),
            complex(0.49004931075238933),
            complex(0.8373593063365198),
            complex(1.0),
            complex(0.8373593063365196),
            complex(0.49004931075238906),
            complex(0.1967938049565092),
            complex(0.048793108747363416),
        ],
        dtype=np.complex128,
    )

    shift_time = 30e-9
    for t, v in zip(times, values):
        expected_amplitudes["frame1"].put(t + shift_time, v)
        expected_frequencies["frame1"].put(t + shift_time, 1e8)
        expected_phases["frame1"].put(t + shift_time, 0)

    parser = _ApproximationParser(program=pulse_seq._program, frames=to_dict(frame1))

    verify_results(parser, expected_amplitudes, expected_frequencies, expected_phases)


def test_play_drag_gaussian_waveforms(port):
    frame1 = Frame(frame_id="frame1", port=port, frequency=1e8, phase=0, is_predefined=False)

    expected_amplitudes = {"frame1": TimeSeries()}
    expected_frequencies = {"frame1": TimeSeries()}
    expected_phases = {"frame1": TimeSeries()}

    drag_gaussian_wf_ZaE_False = DragGaussianWaveform(
        length=1e-8, sigma=1.69e-9, amplitude=1.0, beta=1e-9, zero_at_edges=False
    )
    drag_gaussian_wf_ZaE_True = DragGaussianWaveform(
        length=1e-8, sigma=1.69e-9, amplitude=1.0, beta=1e-9, zero_at_edges=True
    )
    pulse_seq = (
        PulseSequence()
        .play(frame1, drag_gaussian_wf_ZaE_False)
        .delay(frame1, 2e-8)
        .play(frame1, drag_gaussian_wf_ZaE_True)
    )

    # First DRAG_Gaussian with zero_at_edges=False
    times = np.arange(0, 1e-8, port.dt)
    values = np.array(
        [
            complex(0.012568049266111367, 0.02200211698839566),
            complex(0.06074792381889125, 0.08507814687005531),
            complex(0.20688853998666168, 0.21731228597037397),
            complex(0.49645839613814063, 0.3476477687322857),
            complex(0.839403382587157, 0.29389845684225235),
            complex(1.0, 0.0),
            complex(0.839403382587157, -0.29389845684225235),
            complex(0.49645839613814063, -0.3476477687322857),
            complex(0.20688853998666168, -0.21731228597037397),
            complex(0.06074792381889125, -0.08507814687005531),
        ],
        dtype=np.complex128,
    )

    for t, v in zip(times, values):
        expected_amplitudes["frame1"].put(t, v)
        expected_frequencies["frame1"].put(t, 1e8)
        expected_phases["frame1"].put(t, 0)

    # Add a delay between the two waveforms
    shift_time = 10e-9
    expected_amplitudes["frame1"].put(shift_time + 0, 0).put(shift_time + 19e-9, 0)
    expected_frequencies["frame1"].put(shift_time + 0, 1e8).put(shift_time + 19e-9, 1e8)
    expected_phases["frame1"].put(shift_time + 0, 0).put(shift_time + 19e-9, 0)

    # Second DRAG_Gaussian with zero_at_edges=True
    times = np.arange(0, 1e-8, port.dt)
    values = np.array(
        [
            complex(0.0, 0.0),
            complex(0.04879310874736347, 0.0683352946288484),
            complex(0.1967938049565093, 0.20670894396888342),
            complex(0.49004931075238933, 0.34315977084303023),
            complex(0.8373593063365198, 0.2931827689284408),
            complex(1.0, 0),
            complex(0.8373593063365196, -0.293182768928441),
            complex(0.49004931075238906, -0.34315977084303023),
            complex(0.1967938049565092, -0.20670894396888334),
            complex(0.048793108747363416, -0.06833529462884834),
        ],
        dtype=np.complex128,
    )

    shift_time = shift_time + 20e-9
    for t, v in zip(times, values):
        expected_amplitudes["frame1"].put(t + shift_time, v)
        expected_frequencies["frame1"].put(t + shift_time, 1e8)
        expected_phases["frame1"].put(t + shift_time, 0)

    parser = _ApproximationParser(program=pulse_seq._program, frames=to_dict(frame1))

    verify_results(parser, expected_amplitudes, expected_frequencies, expected_phases)


def test_barrier_same_dt(port):
    frame1 = Frame(frame_id="frame1", port=port, frequency=1e8, phase=0, is_predefined=False)
    frame2 = Frame(frame_id="frame2", port=port, frequency=1e8, phase=0, is_predefined=False)
    pulse_seq = (
        PulseSequence()
        .play(frame1, ConstantWaveform(12e-9, 0.75))  # Inst1
        .barrier([frame1, frame2])  # Inst2
        .play(frame1, ConstantWaveform(16e-9, 1))  # Inst3
        .play(frame2, ConstantWaveform(8e-9, -1))  # Inst4
    )

    expected_amplitudes = {"frame1": TimeSeries(), "frame2": TimeSeries()}
    expected_frequencies = {"frame1": TimeSeries(), "frame2": TimeSeries()}
    expected_phases = {"frame1": TimeSeries(), "frame2": TimeSeries()}

    # Inst1
    shift_time_frame1 = 0
    pulse_length = 12e-9
    times = np.arange(0, pulse_length, port.dt)
    values = 0.75 * np.ones_like(times)
    for t, v in zip(times, values):
        expected_amplitudes["frame1"].put(shift_time_frame1 + t, v)
        expected_frequencies["frame1"].put(shift_time_frame1 + t, 1e8)
        expected_phases["frame1"].put(shift_time_frame1 + t, 0)

    # Inst2
    # Delay frame2 from 0ns to 11ns
    shift_time_frame2 = 0

    expected_amplitudes["frame2"].put(0, 0).put(11e-9, 0)
    expected_frequencies["frame2"].put(0, 1e8).put(11e-9, 1e8)
    expected_phases["frame2"].put(0, 0).put(11e-9, 0)

    # Inst3
    shift_time_frame1 = shift_time_frame1 + pulse_length
    pulse_length = 16e-9
    times = np.arange(0, pulse_length, port.dt)
    values = 1 * np.ones_like(times)
    for t, v in zip(times, values):
        expected_amplitudes["frame1"].put(shift_time_frame1 + t, v)
        expected_frequencies["frame1"].put(shift_time_frame1 + t, 1e8)
        expected_phases["frame1"].put(shift_time_frame1 + t, 0)

    # Inst4
    shift_time_frame2 = shift_time_frame1
    pulse_length = 8e-9
    times = np.arange(0, pulse_length, port.dt)
    values = -1 * np.ones_like(times)
    for t, v in zip(times, values):
        expected_amplitudes["frame2"].put(shift_time_frame2 + t, v)
        expected_frequencies["frame2"].put(shift_time_frame2 + t, 1e8)
        expected_phases["frame2"].put(shift_time_frame2 + t, 0)

    parser = _ApproximationParser(program=pulse_seq._program, frames=to_dict([frame1, frame2]))

    verify_results(parser, expected_amplitudes, expected_frequencies, expected_phases)


def test_barrier_no_args(port):
    frame1 = Frame(frame_id="frame1", port=port, frequency=1e8, phase=0, is_predefined=False)
    frame2 = Frame(frame_id="frame2", port=port, frequency=1e8, phase=0, is_predefined=False)
    pulse_seq = (
        PulseSequence()
        .play(frame1, ConstantWaveform(12e-9, 0.75))  # Inst1
        .barrier([])  # Inst2
        .play(frame1, ConstantWaveform(16e-9, 1))  # Inst3
        .play(frame2, ConstantWaveform(8e-9, -1))  # Inst4
    )

    expected_amplitudes = {"frame1": TimeSeries(), "frame2": TimeSeries()}
    expected_frequencies = {"frame1": TimeSeries(), "frame2": TimeSeries()}
    expected_phases = {"frame1": TimeSeries(), "frame2": TimeSeries()}

    # Inst1
    shift_time_frame1 = 0
    pulse_length = 12e-9
    times = np.arange(0, pulse_length, port.dt)
    values = 0.75 * np.ones_like(times)
    for t, v in zip(times, values):
        expected_amplitudes["frame1"].put(shift_time_frame1 + t, v)
        expected_frequencies["frame1"].put(shift_time_frame1 + t, 1e8)
        expected_phases["frame1"].put(shift_time_frame1 + t, 0)

    # Inst2
    # Delay frame2 from 0ns to 11ns
    shift_time_frame2 = 0

    expected_amplitudes["frame2"].put(0, 0).put(11e-9, 0)
    expected_frequencies["frame2"].put(0, 1e8).put(11e-9, 1e8)
    expected_phases["frame2"].put(0, 0).put(11e-9, 0)

    # Inst3
    shift_time_frame1 = shift_time_frame1 + pulse_length
    pulse_length = 16e-9
    times = np.arange(0, pulse_length, port.dt)
    values = 1 * np.ones_like(times)
    for t, v in zip(times, values):
        expected_amplitudes["frame1"].put(shift_time_frame1 + t, v)
        expected_frequencies["frame1"].put(shift_time_frame1 + t, 1e8)
        expected_phases["frame1"].put(shift_time_frame1 + t, 0)

    # Inst4
    shift_time_frame2 = shift_time_frame1
    pulse_length = 8e-9
    times = np.arange(0, pulse_length, port.dt)
    values = -1 * np.ones_like(times)
    for t, v in zip(times, values):
        expected_amplitudes["frame2"].put(shift_time_frame2 + t, v)
        expected_frequencies["frame2"].put(shift_time_frame2 + t, 1e8)
        expected_phases["frame2"].put(shift_time_frame2 + t, 0)

    parser = _ApproximationParser(program=pulse_seq._program, frames=to_dict([frame1, frame2]))

    verify_results(parser, expected_amplitudes, expected_frequencies, expected_phases)


def test_barrier_different_dt(port):
    port1 = Port(port_id="device_port_x1", dt=5e-9, properties={})
    port2 = Port(port_id="device_port_x2", dt=4e-9, properties={})
    frame1 = Frame(frame_id="frame1", port=port1, frequency=1e8, phase=0, is_predefined=False)
    frame2 = Frame(frame_id="frame2", port=port2, frequency=1e8, phase=0, is_predefined=False)
    pulse_seq = (
        PulseSequence()
        .play(frame1, ConstantWaveform(25e-9, 0.75))  # Inst 1
        .barrier([frame1, frame2])  # Inst 2
        .play(frame1, ConstantWaveform(40e-9, 1))  # Inst 3
        .play(frame2, ConstantWaveform(28e-9, -1))  # Inst 4
    )

    expected_amplitudes = {"frame1": TimeSeries(), "frame2": TimeSeries()}
    expected_frequencies = {"frame1": TimeSeries(), "frame2": TimeSeries()}
    expected_phases = {"frame1": TimeSeries(), "frame2": TimeSeries()}

    # Inst1
    shift_time_frame1 = 0
    pulse_length = 25e-9
    times = np.arange(0, pulse_length, port1.dt)
    values = 0.75 * np.ones_like(times)
    for t, v in zip(times, values):
        expected_amplitudes["frame1"].put(shift_time_frame1 + t, v)
        expected_frequencies["frame1"].put(shift_time_frame1 + t, 1e8)
        expected_phases["frame1"].put(shift_time_frame1 + t, 0)

    # Inst2
    # barrier at 40ns (first point coinciding with frame2 (every 20ns))
    shift_time_frame1 = shift_time_frame2 = 40e-9

    latest_time_frame1 = expected_amplitudes["frame1"].times()[-1]
    expected_amplitudes["frame1"].put(latest_time_frame1 + port1.dt, 0).put(
        shift_time_frame1 - port1.dt, 0
    )
    expected_frequencies["frame1"].put(latest_time_frame1 + port1.dt, 1e8).put(
        shift_time_frame1 - port1.dt, 1e8
    )
    expected_phases["frame1"].put(latest_time_frame1 + port1.dt, 0).put(
        shift_time_frame1 - port1.dt, 0
    )

    expected_amplitudes["frame2"].put(0, 0).put(shift_time_frame2 - port2.dt, 0)
    expected_frequencies["frame2"].put(0, 1e8).put(shift_time_frame2 - port2.dt, 1e8)
    expected_phases["frame2"].put(0, 0).put(shift_time_frame2 - port2.dt, 0)

    # Inst3
    pulse_length = 40e-9
    times = np.arange(0, pulse_length, port1.dt)
    values = 1 * np.ones_like(times)
    for t, v in zip(times, values):
        expected_amplitudes["frame1"].put(shift_time_frame1 + t, v)
        expected_frequencies["frame1"].put(shift_time_frame1 + t, 1e8)
        expected_phases["frame1"].put(shift_time_frame1 + t, 0)

    # Inst4
    pulse_length = 28e-9
    times = np.arange(0, pulse_length, port2.dt)
    values = -1 * np.ones_like(times)
    for t, v in zip(times, values):
        expected_amplitudes["frame2"].put(shift_time_frame2 + t, v)
        expected_frequencies["frame2"].put(shift_time_frame2 + t, 1e8)
        expected_phases["frame2"].put(shift_time_frame2 + t, 0)

    parser = _ApproximationParser(program=pulse_seq._program, frames=to_dict([frame1, frame2]))

    verify_results(parser, expected_amplitudes, expected_frequencies, expected_phases)


def test_pad_different_dt(port):
    port1 = Port(port_id="device_port_x1", dt=5e-9, properties={})
    port2 = Port(port_id="device_port_x2", dt=4e-9, properties={})
    frame1 = Frame(frame_id="frame1", port=port1, frequency=1e8, phase=0, is_predefined=False)
    frame2 = Frame(frame_id="frame2", port=port2, frequency=1e8, phase=0, is_predefined=False)

    pulse_seq = (
        PulseSequence()
        .play(frame1, ConstantWaveform(20e-9, 0.75))
        .play(frame2, ConstantWaveform(28e-9, -1))
    )

    expected_amplitudes = {"frame1": TimeSeries(), "frame2": TimeSeries()}
    expected_frequencies = {"frame1": TimeSeries(), "frame2": TimeSeries()}
    expected_phases = {"frame1": TimeSeries(), "frame2": TimeSeries()}

    # Inst1
    shift_time_frame1 = 0
    pulse_length = 20e-9
    times = np.arange(0, pulse_length, port1.dt)
    values = 0.75 * np.ones_like(times)
    for t, v in zip(times, values):
        expected_amplitudes["frame1"].put(shift_time_frame1 + t, v)
        expected_frequencies["frame1"].put(shift_time_frame1 + t, 1e8)
        expected_phases["frame1"].put(shift_time_frame1 + t, 0)

    # Inst2
    shift_time_frame2 = 0
    pulse_length = 28e-9
    times = np.arange(0, pulse_length, port2.dt)
    values = -1 * np.ones_like(times)
    for t, v in zip(times, values):
        expected_amplitudes["frame2"].put(shift_time_frame2 + t, v)
        expected_frequencies["frame2"].put(shift_time_frame2 + t, 1e8)
        expected_phases["frame2"].put(shift_time_frame2 + t, 0)

    parser = _ApproximationParser(program=pulse_seq._program, frames=to_dict([frame1, frame2]))

    verify_results(parser, expected_amplitudes, expected_frequencies, expected_phases)


@pytest.mark.parametrize(
    "literal_type, lhs, operator, rhs, expected_result",
    [
        (ast.FloatLiteral, 3, "+", 2, 5),
        (ast.FloatLiteral, 3, "-", 2, 1),
        (ast.FloatLiteral, 3, "*", 2, 6),
        (ast.FloatLiteral, 3, "/", 2, 1.5),
        (ast.FloatLiteral, 3, "%", 2, 1),
        (ast.FloatLiteral, 3, "**", 2, 9),
        (ast.FloatLiteral, 3, "<", 2, False),
        (ast.FloatLiteral, 3, ">", 2, True),
        (ast.FloatLiteral, 3, ">=", 2, True),
        (ast.FloatLiteral, 3, "<=", 2, False),
        (ast.FloatLiteral, 3, "==", 2, False),
        (ast.FloatLiteral, 3, "!=", 2, True),
        (ast.BooleanLiteral, True, "&&", False, False),
        (ast.BooleanLiteral, True, "||", False, True),
        (ast.BooleanLiteral, True, "|", False, True),
        (ast.BooleanLiteral, True, "^", False, True),
        (ast.BooleanLiteral, True, "&", False, False),
        (ast.IntegerLiteral, 3, "<<", 2, 12),
        (ast.IntegerLiteral, 3, ">>", 2, 0),
    ],
)
def test_binary_operations(literal_type, lhs, operator, rhs, expected_result):
    expression = ast.BinaryExpression(
        lhs=literal_type(lhs),
        op=ast.BinaryOperator[operator],
        rhs=literal_type(rhs),
    )
    pulse_seq = PulseSequence()
    parser = _ApproximationParser(program=pulse_seq._program, frames={})
    result = parser.visit_BinaryExpression(expression, Mock())
    assert result == expected_result


@pytest.mark.parametrize(
    "literal_type, operator, literal, expected_result",
    [
        (ast.FloatLiteral, "-", 3, -3),
        (ast.BooleanLiteral, "!", True, False),
        (ast.IntegerLiteral, "~", 3, ~3),
    ],
)
def test_unary_operations(literal_type, operator, literal, expected_result):
    expression = ast.UnaryExpression(
        op=ast.UnaryOperator[operator], expression=literal_type(literal)
    )
    pulse_seq = PulseSequence()
    parser = _ApproximationParser(program=pulse_seq._program, frames={})
    result = parser.visit_UnaryExpression(expression, Mock())
    assert result == expected_result


@pytest.mark.xfail(raises=ValueError)
def test_duration_literal():
    literal = Mock()
    pulse_seq = PulseSequence()
    parser = _ApproximationParser(program=pulse_seq._program, frames={})
    parser.visit_DurationLiteral(literal, Mock())


def verify_results(results, expected_amplitudes, expected_frequencies, expected_phases):
    for frame_id in expected_amplitudes.keys():
        assert _all_close(results.amplitudes[frame_id], expected_amplitudes[frame_id], 1e-10)
        assert _all_close(results.frequencies[frame_id], expected_frequencies[frame_id], 1e-10)
        assert _all_close(results.phases[frame_id], expected_phases[frame_id], 1e-10)


def to_dict(frames: Union[Frame, List]):
    if not isinstance(frames, List):
        frames = [frames]
    frame_dict = dict()
    for frame in frames:
        frame_dict[frame.id] = frame
    return frame_dict
