# 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.

"""Tests for the pulse control module."""

import textwrap

import pytest

import autoqasm as aq
from autoqasm.instructions import rx
from autoqasm.pulse import (
    barrier,
    capture_v0,
    delay,
    play,
    set_frequency,
    set_phase,
    set_scale,
    shift_frequency,
    shift_phase,
)
from braket.parametric import FreeParameter
from braket.pulse import ArbitraryWaveform, Frame, Port

PORT = Port(port_id="device_port_x0", dt=1e-9, properties={})
FRAME1 = Frame(frame_id="predefined_frame_1", frequency=2e9, port=PORT, phase=0, is_predefined=True)
FRAME2 = Frame(frame_id="predefined_frame_2", frequency=1e6, port=PORT, is_predefined=True)
WAVEFORM = ArbitraryWaveform([complex(1, 0.4), 0, 0.3, complex(0.1, 0.2)], id="arb_wf")


def test_mix_gate_pulse() -> None:
    """Test mixed usage of gates and pulses."""

    @aq.main
    def my_program():
        shift_frequency(FRAME1, 0.1234)
        rx(1, 0.1)
        play(FRAME1, WAVEFORM)

    expected = textwrap.dedent(
        """
        OPENQASM 3.0;
        cal {
            waveform arb_wf = {1.0 + 0.4im, 0, 0.3, 0.1 + 0.2im};
        }
        qubit[2] __qubits__;
        cal {
            shift_frequency(predefined_frame_1, 0.1234);
        }
        rx(0.1) __qubits__[1];
        cal {
            play(predefined_frame_1, arb_wf);
        }
        """
    ).strip()
    assert my_program.build().to_ir() == expected


def test_merge_cal_box() -> None:
    """Test subsequent cal boxes are merged."""

    @aq.main
    def my_program():
        barrier("$0")
        delay(["$3", "$4"], 0.34)

    expected = textwrap.dedent(
        """
        OPENQASM 3.0;
        cal {
            barrier $0;
            delay[340.0ms] $3, $4;
        }
        """
    ).strip()
    assert my_program.build().to_ir() == expected


@pytest.mark.parametrize(
    "instruction,qubits_or_frames,params,expected_qasm",
    [
        (
            shift_frequency,
            FRAME1,
            [0.12],
            "\ncal {\n    shift_frequency(predefined_frame_1, 0.12);\n}",
        ),
        (
            set_frequency,
            FRAME1,
            [0.12],
            "\ncal {\n    set_frequency(predefined_frame_1, 0.12);\n}",
        ),
        (
            shift_phase,
            FRAME1,
            [0.12],
            "\ncal {\n    shift_phase(predefined_frame_1, 0.12);\n}",
        ),
        (
            set_phase,
            FRAME1,
            [0.12],
            "\ncal {\n    set_phase(predefined_frame_1, 0.12);\n}",
        ),
        (
            set_scale,
            FRAME1,
            [0.12],
            "\ncal {\n    set_scale(predefined_frame_1, 0.12);\n}",
        ),
        (delay, "$3", [0.34], "\ncal {\n    delay[340.0ms] $3;\n}"),
        (delay, ["$3", "$4"], [0.34], "\ncal {\n    delay[340.0ms] $3, $4;\n}"),
        (
            delay,
            FRAME1,
            [0.34],
            "\ncal {\n    delay[340.0ms] predefined_frame_1;\n}",
        ),
        (
            delay,
            [FRAME1, FRAME2],
            [0.34],
            "\ncal {\n    delay[340.0ms] predefined_frame_1, predefined_frame_2;\n}",
        ),
        (barrier, "$3", [], "\ncal {\n    barrier $3;\n}"),
        (barrier, ["$3", "$4"], [], "\ncal {\n    barrier $3, $4;\n}"),
        (barrier, FRAME1, [], "\ncal {\n    barrier predefined_frame_1;\n}"),
        (
            barrier,
            [FRAME1, FRAME2],
            [],
            "\ncal {\n    barrier predefined_frame_1, predefined_frame_2;\n}",
        ),
        (
            play,
            FRAME1,
            [WAVEFORM],
            (
                "\ncal {\n    waveform arb_wf = {1.0 + 0.4im, 0, 0.3, 0.1 + 0.2im};"
                "\n    play(predefined_frame_1, arb_wf);\n}"
            ),
        ),
        (
            capture_v0,
            FRAME1,
            [],
            ("\nbit __bit_0__;\ncal {\n    __bit_0__ = capture_v0(predefined_frame_1);\n}"),
        ),
    ],
)
def test_pulse_control(instruction, qubits_or_frames, params, expected_qasm) -> None:
    """Test pulse control operations."""
    with aq.build_program() as program_conversion_context:
        instruction(qubits_or_frames, *params)

    assert expected_qasm in program_conversion_context.make_program().to_ir()


@pytest.mark.parametrize(
    "instruction,qubits_or_frames,params",
    [
        (barrier, "1", []),
        (barrier, 1, []),
        (barrier, ["1", "2"], []),
        (barrier, [1, 2], []),
    ],
)
def test_pulse_control_invalid_physical_qubit(instruction, qubits_or_frames, params) -> None:
    """Test pulse control operations with invalid lables for physical qubits."""
    with pytest.raises(ValueError):
        with aq.build_program():
            instruction(qubits_or_frames, *params)


def test_pulse_freeparameter() -> None:
    """Test pulse program with free parameter."""

    @aq.main
    def my_program(duration1):
        delay(["$3", "$4"], duration1)
        delay(["$3", "$4"], FreeParameter("duration2"))

    expected = textwrap.dedent(
        """
        OPENQASM 3.0;
        input float duration1;
        input float duration2;
        cal {
            delay[duration1 * 1s] $3, $4;
            delay[duration2 * 1s] $3, $4;
        }
        """
    ).strip()
    assert my_program.build().to_ir() == expected


def test_pulse_freeparameter_bound() -> None:
    """Test pulse program with freeparameter bound with values."""

    @aq.main
    def my_program(duration1):
        delay(["$3", "$4"], duration1)

    expected = textwrap.dedent(
        """
        OPENQASM 3.0;
        float duration1 = 0.123;
        cal {
            delay[duration1 * 1s] $3, $4;
        }
        """
    ).strip()
    qasm = my_program.build().make_bound_program({"duration1": 0.123}).to_ir()
    assert qasm == expected


def test_pulse_freeparameter_condition() -> None:
    """Test pulse program with freeparameter in condition."""

    @aq.main
    def my_program(duration1, duration2):
        delay(["$3", "$4"], duration1)
        if duration1 > duration2:
            delay(["$3", "$4"], duration2)

    expected = textwrap.dedent(
        """
        OPENQASM 3.0;
        input float duration1;
        input float duration2;
        cal {
            delay[duration1 * 1s] $3, $4;
        }
        bool __bool_0__;
        __bool_0__ = duration1 > duration2;
        if (__bool_0__) {
            cal {
                delay[duration2 * 1s] $3, $4;
            }
        }
        """
    ).strip()
    assert my_program.build().to_ir() == expected
