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

"""AutoQASM tests for parameter support."""

import numpy as np
import pytest

import autoqasm as aq
from autoqasm import pulse
from autoqasm.instructions import cnot, cphaseshift, gpi, h, measure, ms, rx, rz, x
from autoqasm.simulator import McmSimulator
from braket.circuits import FreeParameter
from braket.devices import LocalSimulator
from braket.tasks.local_quantum_task import LocalQuantumTask


def _test_parametric_on_local_sim(program: aq.Program, inputs: dict[str, float]) -> np.ndarray:
    device = LocalSimulator(backend=McmSimulator())
    task = device.run(program, shots=100, inputs=inputs)
    assert isinstance(task, LocalQuantumTask)
    assert isinstance(task.result().measurements, dict)
    return task.result().measurements


@pytest.fixture
def simple_parametric():
    @aq.main
    def simple_parametric(theta):
        rx(0, theta)
        measure(0)

    return simple_parametric


@pytest.fixture
def simple_parametric_fp():
    @aq.main
    def simple_parametric():
        rx(0, FreeParameter("theta"))
        measure(0)

    return simple_parametric


def test_simple_parametric(simple_parametric):
    """Test a program with a parameter can be serialized."""

    expected = """OPENQASM 3.0;
input float theta;
qubit[1] __qubits__;
rx(theta) __qubits__[0];
bit __bit_0__;
__bit_0__ = measure __qubits__[0];"""
    assert simple_parametric.build().to_ir() == expected


def test_simple_parametric_fp(simple_parametric_fp):
    """Test a program with an explicit free parameter can be serialized."""

    expected = """OPENQASM 3.0;
input float theta;
qubit[1] __qubits__;
rx(theta) __qubits__[0];
bit __bit_0__;
__bit_0__ = measure __qubits__[0];"""
    assert simple_parametric_fp.build().to_ir() == expected


def test_sim_simple(simple_parametric):
    measurements = _test_parametric_on_local_sim(simple_parametric, {"theta": 0})
    assert 1 not in measurements["__bit_0__"]
    measurements = _test_parametric_on_local_sim(simple_parametric, {"theta": np.pi})
    assert 0 not in measurements["__bit_0__"]


def test_sim_simple_fp(simple_parametric_fp):
    measurements = _test_parametric_on_local_sim(simple_parametric_fp, {"theta": 0})
    assert 1 not in measurements["__bit_0__"]
    measurements = _test_parametric_on_local_sim(simple_parametric_fp, {"theta": np.pi})
    assert 0 not in measurements["__bit_0__"]


@pytest.fixture
def multi_parametric():
    @aq.main
    def multi_parametric(alpha, theta):
        rx(0, alpha)
        rx(1, theta)
        c = measure([0, 1])  # noqa: F841

    return multi_parametric


def test_multiple_parameters(multi_parametric):
    """Test that multiple free parameters all appear in the processed program."""

    expected = """OPENQASM 3.0;
input float alpha;
input float theta;
bit[2] c;
qubit[2] __qubits__;
rx(alpha) __qubits__[0];
rx(theta) __qubits__[1];
bit[2] __bit_0__ = "00";
__bit_0__[0] = measure __qubits__[0];
__bit_0__[1] = measure __qubits__[1];
c = __bit_0__;"""
    assert multi_parametric.build().to_ir() == expected


def test_sim_multi_param(multi_parametric):
    measurements = _test_parametric_on_local_sim(multi_parametric, {"alpha": np.pi, "theta": 0})
    assert all(val == "10" for val in measurements["c"])
    measurements = _test_parametric_on_local_sim(multi_parametric, {"alpha": 0, "theta": np.pi})
    assert all(val == "01" for val in measurements["c"])


def test_repeat_parameter():
    """Test that programs can use the same parameter multiple times."""

    @aq.main
    def parametric(alpha, theta):
        rx(0, alpha)
        rx(1, theta)
        cnot(0, 1)
        rx(0, theta)
        rx(1, alpha)

    expected = """OPENQASM 3.0;
input float alpha;
input float theta;
qubit[2] __qubits__;
rx(alpha) __qubits__[0];
rx(theta) __qubits__[1];
cnot __qubits__[0], __qubits__[1];
rx(theta) __qubits__[0];
rx(alpha) __qubits__[1];"""
    assert parametric.build().to_ir() == expected


@pytest.mark.xfail(
    reason="Regression: FreeParameters create FloatVar with _needs_declaration=False"
)
def test_parameter_in_subroutine():
    """Test that parameters in subroutines are declared appropriately."""
    # TODO (#816): the openqasm generated here isn't strictly valid
    # (cannot close over non-const global variables)
    # https://openqasm.com/language/scope.html#subroutine-and-gate-scope

    @aq.subroutine
    def rx_alpha(qubit: int):
        rx(qubit, FreeParameter("alpha"))

    @aq.main(num_qubits=3)
    def parametric():
        rx_alpha(2)

    expected = """OPENQASM 3.0;
def rx_alpha(int[32] qubit) {
    rx(alpha) __qubits__[qubit];
}
input float[64] alpha;
qubit[3] __qubits__;
rx_alpha(2);"""
    assert parametric.build().to_ir() == expected


def test_captured_parameter():
    """Test that a parameter declared in a larger scope is captured
    and functions correctly.
    """
    alpha = FreeParameter("alpha")

    @aq.main
    def parametric():
        rz(0, alpha)
        rx(1, alpha)

    expected = """OPENQASM 3.0;
input float alpha;
qubit[2] __qubits__;
rz(alpha) __qubits__[0];
rx(alpha) __qubits__[1];"""
    assert parametric.build().to_ir() == expected


def test_multi_angle_gates():
    """Test that FreeParameters work with gates that take multiple inputs."""
    qubit_0 = 2
    theta = 0.5

    @aq.main(num_qubits=5)
    def parametric(phi: float):
        ms(0, qubit_0, phi, phi, theta)

    expected = """OPENQASM 3.0;
input float phi;
qubit[5] __qubits__;
ms(phi, phi, 0.5) __qubits__[0], __qubits__[2];"""
    assert parametric.build().to_ir() == expected


def test_sim_multi_angle():
    theta = 0

    @aq.main
    def parametric(phi: float):
        ms(0, 1, phi, phi, theta)

    _test_parametric_on_local_sim(
        parametric,
        inputs={"phi": np.pi},
    )


def test_parameters_passed_as_main_arg():
    """Test that parameters work when passed as input values."""

    @aq.main
    def parametric(my_phi: float):
        cphaseshift(0, 1, my_phi)

    expected = """OPENQASM 3.0;
input float my_phi;
qubit[2] __qubits__;
cphaseshift(my_phi) __qubits__[0], __qubits__[1];"""
    assert parametric.build().to_ir() == expected


def test_simple_subroutine_arg():
    """Test that parameters work when passed as input values."""

    @aq.subroutine
    def silly_rz(theta: float):
        rz(0, theta)

    @aq.main(num_qubits=1)
    def parametric(alpha):
        silly_rz(alpha)

    expected = """OPENQASM 3.0;
def silly_rz(float[64] theta) {
    rz(theta) __qubits__[0];
}
input float alpha;
qubit[1] __qubits__;
silly_rz(alpha);"""
    assert parametric.build().to_ir() == expected


def test_parameters_passed_as_subroutine_arg():
    """Test that parameters work when passed as input values."""

    @aq.subroutine
    def silly_ms(qubit_0: int, phi: float, theta: float):
        ms(0, qubit_0, phi, phi, theta)

    @aq.main(num_qubits=5)
    def parametric(alpha):
        silly_ms(1, alpha, 0.707)
        silly_ms(3, 0.5, FreeParameter("beta"))

    expected = """OPENQASM 3.0;
def silly_ms(int[32] qubit_0, float[64] phi, float[64] theta) {
    ms(phi, phi, theta) __qubits__[0], __qubits__[qubit_0];
}
input float alpha;
input float beta;
qubit[5] __qubits__;
silly_ms(1, alpha, 0.707);
silly_ms(3, 0.5, beta);"""
    assert parametric.build().to_ir() == expected


def test_sim_subroutine_arg():
    @aq.subroutine
    def rx_theta(theta: float):
        rx(0, theta)

    @aq.main
    def parametric(theta):
        rx_theta(theta)
        measure(0)

    measurements = _test_parametric_on_local_sim(parametric, {"theta": np.pi})
    assert 0 not in measurements["__bit_0__"]


def test_parametric_gate_args():
    """Test that gates can be used with parameters."""

    @aq.gate
    def rx_theta(q: aq.Qubit, theta: float):
        rx(q, theta)

    @aq.main(num_qubits=3)
    def parametric():
        rx_theta(2, FreeParameter("θ"))

    expected = """OPENQASM 3.0;
gate rx_theta(theta) q {
    rx(theta) q;
}
input float θ;
qubit[3] __qubits__;
rx_theta(θ) __qubits__[2];"""
    assert parametric.build().to_ir() == expected


def test_parametric_pulse_cals():
    """Test that pulse calibrations work with free parameters."""

    @aq.gate_calibration(implements=rx, target="$1")
    def cal_1(theta: float):
        pulse.delay("$1", theta)

    @aq.main
    def my_program(a):
        rx("$1", a)

    expected = """OPENQASM 3.0;
input float a;
defcal rx(angle[32] theta) $1 {
    delay[theta * 1s] $1;
}
rx(a) $1;"""
    qasm = my_program.build().with_calibrations(cal_1).to_ir()
    assert qasm == expected


def test_bind_parameters():
    """Test binding FreeParameters to concrete values."""

    @aq.main
    def parametric(alpha: float):
        rx(0, alpha)
        measure(0)

    unbound_expected = """OPENQASM 3.0;
input float alpha;
qubit[1] __qubits__;
rx(alpha) __qubits__[0];
bit __bit_0__;
__bit_0__ = measure __qubits__[0];"""

    bound_template = """OPENQASM 3.0;
float alpha = {};
qubit[1] __qubits__;
rx(alpha) __qubits__[0];
bit __bit_0__;
__bit_0__ = measure __qubits__[0];"""

    assert parametric.build().to_ir() == unbound_expected
    bound_prog = parametric.build().make_bound_program({"alpha": 0.5})
    # Original program unchanged
    assert parametric.build().to_ir() == unbound_expected
    assert bound_prog.to_ir() == bound_template.format(0.5)
    # Can rebind
    bound_prog = parametric.build().make_bound_program({"alpha": 0.432143})
    assert bound_prog.to_ir() == bound_template.format(0.432143)


def test_multi_bind_parameters():
    """Test binding FreeParameters to concrete values."""

    @aq.subroutine
    def sub(alpha: float, theta: float):
        rx(0, alpha)
        rx(1, theta)
        cnot(0, 1)
        rx(0, theta)
        rx(1, alpha)

    @aq.subroutine
    def rx_alpha(qubit: int):
        rx(qubit, FreeParameter("alpha"))

    @aq.main(num_qubits=3)
    def parametric(alpha: float, beta: float):
        sub(alpha, beta)
        rx_alpha(2)

    bound_prog = parametric.build().make_bound_program({"alpha": 0.5, "beta": 1.5})

    expected = """OPENQASM 3.0;
def sub(float[64] alpha, float[64] theta) {
    rx(alpha) __qubits__[0];
    rx(theta) __qubits__[1];
    cnot __qubits__[0], __qubits__[1];
    rx(theta) __qubits__[0];
    rx(alpha) __qubits__[1];
}
def rx_alpha(int[32] qubit_) {
    rx(alpha) __qubits__[qubit_];
}
float alpha = 0.5;
float beta = 1.5;
qubit[3] __qubits__;
sub(alpha, beta);
rx_alpha(2);"""
    assert bound_prog.to_ir() == expected


def test_partial_bind():
    """Test binding some but not all FreeParameters."""

    @aq.subroutine
    def rx_alpha(qubit: int, theta: float):
        rx(qubit, theta)

    @aq.main(num_qubits=3)
    def parametric(alpha: float, beta: float):
        rx_alpha(2, alpha)
        rx_alpha(2, beta)

    bound_prog = parametric.build().make_bound_program({"beta": np.pi})

    expected = """OPENQASM 3.0;
def rx_alpha(int[32] qubit_, float[64] theta) {
    rx(theta) __qubits__[qubit_];
}
input float alpha;
float beta = 3.141592653589793;
qubit[3] __qubits__;
rx_alpha(2, alpha);
rx_alpha(2, beta);"""
    assert bound_prog.to_ir() == expected


def test_binding_pulse_parameters():
    """Test binding programs with parametric pulse instructions."""

    @aq.gate_calibration(implements=rx, target="$1")
    def cal_1(theta: float):
        pulse.delay("$1", theta)

    @aq.main
    def my_program(a):
        rx("$1", a)

    qasm1 = my_program.build().with_calibrations(cal_1).make_bound_program({"a": 0.6}).to_ir()
    qasm2 = my_program.build().make_bound_program({"a": 0.6}).with_calibrations(cal_1).to_ir()
    assert qasm1 == qasm2

    expected = """OPENQASM 3.0;
float a = 0.6;
defcal rx(angle[32] theta) $1 {
    delay[theta * 1s] $1;
}
rx(a) $1;"""
    assert expected == qasm1


def test_bind_empty_program():
    """Test that binding behaves well on empty programs."""

    @aq.main
    def empty_program():
        pass

    qasm = empty_program.build().to_ir()
    bound_program1 = empty_program.build().make_bound_program({}).to_ir()
    bound_program2 = empty_program.build().make_bound_program({"alpha": 0.5}).to_ir()
    assert qasm == bound_program1 == bound_program2


def test_strict_parameter_bind():
    """Test make_bound_program with strict set to True."""

    @aq.main
    def parametric(alpha: float):
        rx(0, alpha)
        measure(0)

    template = """OPENQASM 3.0;
float alpha = {};
qubit[1] __qubits__;
rx(alpha) __qubits__[0];
bit __bit_0__;
__bit_0__ = measure __qubits__[0];"""

    bound_prog = parametric.build().make_bound_program({"alpha": 0.5}, strict=True)
    assert bound_prog.to_ir() == template.format(0.5)


def test_strict_parameter_bind_failure():
    """Test make_bound_program with strict set to True."""

    @aq.main
    def parametric(alpha: float):
        rx(0, alpha)
        measure(0)

    with pytest.raises(
        aq.errors.ParameterNotFoundError, match="No parameter in the program named: beta"
    ):
        parametric.build().make_bound_program({"beta": 0.5}, strict=True)


def test_duplicate_variable_name_fails():
    """Test using a variable and FreeParameter with the same name."""

    @aq.main
    def parametric_explicit():
        alpha = aq.FloatVar(1.2)  # noqa: F841
        rx(0, FreeParameter("alpha"))

    with pytest.raises(RuntimeError, match="conflicting variables with name alpha"):
        parametric_explicit.build()

    @aq.main
    def parametric(alpha):
        alpha = aq.FloatVar(1.2)
        rx(0, alpha)

    with pytest.raises(RuntimeError, match="conflicting variables with name alpha"):
        parametric.build()


def test_binding_variable_fails():
    """Test that trying to bind a variable that isn't declared as a FreeParameter fails."""

    @aq.main
    def parametric():
        alpha = aq.FloatVar(1.2)  # noqa: F841

    with pytest.raises(
        aq.errors.ParameterNotFoundError, match="No parameter in the program named: beta"
    ):
        parametric.build().make_bound_program({"beta": 0.5}, strict=True)


def test_compound_condition():
    """Test parameters used in greater than conditional statements."""

    @aq.main
    def parametric(val: float):
        threshold = 0.9
        if val > threshold or val >= 1.2:
            x(0)
        measure(0)

    expected = """OPENQASM 3.0;
input float val;
qubit[1] __qubits__;
bool __bool_0__;
__bool_0__ = val > 0.9;
bool __bool_1__;
__bool_1__ = val >= 1.2;
bool __bool_2__;
__bool_2__ = __bool_0__ || __bool_1__;
if (__bool_2__) {
    x __qubits__[0];
}
bit __bit_3__;
__bit_3__ = measure __qubits__[0];"""
    assert parametric.build().to_ir() == expected


def test_lt_condition():
    """Test parameters used in less than conditional statements."""

    @aq.main
    def parametric(val: float):
        if val < 0.9:
            x(0)
        if val <= 0.9:
            h(0)
        measure(0)

    expected = """OPENQASM 3.0;
input float val;
qubit[1] __qubits__;
bool __bool_0__;
__bool_0__ = val < 0.9;
if (__bool_0__) {
    x __qubits__[0];
}
bool __bool_1__;
__bool_1__ = val <= 0.9;
if (__bool_1__) {
    h __qubits__[0];
}
bit __bit_2__;
__bit_2__ = measure __qubits__[0];"""
    assert parametric.build().to_ir() == expected


def test_parameter_in_predicate_in_subroutine():
    """Test parameters used in conditional statements."""

    @aq.subroutine
    def sub(val: float):
        threshold = 0.9
        if val > threshold:
            x(0)

    @aq.main
    def parametric(val: float):
        sub(val)
        measure(0)

    expected = """OPENQASM 3.0;
def sub(float[64] val) {
    bool __bool_0__;
    __bool_0__ = val > 0.9;
    if (__bool_0__) {
        x __qubits__[0];
    }
}
input float val;
qubit[1] __qubits__;
sub(val);
bit __bit_1__;
__bit_1__ = measure __qubits__[0];"""
    assert parametric.build().to_ir() == expected


def test_eq_condition():
    """Test parameters used in conditional equals statements."""

    @aq.main
    def parametric(threshold: int):
        if threshold == 1:
            h(0)
        elif threshold == 2:
            x(0)
        else:
            pass
        measure(0)

    expected = """OPENQASM 3.0;
input int[32] threshold;
qubit[1] __qubits__;
bool __bool_0__;
__bool_0__ = threshold == 1;
if (__bool_0__) {
    h __qubits__[0];
} else {
    bool __bool_1__;
    __bool_1__ = threshold == 2;
    if (__bool_1__) {
        x __qubits__[0];
    }
}
bit __bit_2__;
__bit_2__ = measure __qubits__[0];"""
    assert parametric.build().to_ir() == expected


def test_sim_conditional_stmts():
    @aq.main
    def main(basis: int):
        if basis == 1:
            h(0)
        else:
            x(0)
        c = measure(0)  # noqa: F841

    measurements = _test_parametric_on_local_sim(main, {"basis": 0})
    assert all(val == 1 for val in measurements["c"])
    measurements = _test_parametric_on_local_sim(main, {"basis": 1})
    assert 1 in measurements["c"] and 0 in measurements["c"]


def test_sim_comparison_stmts():
    @aq.main
    def main(basis: float):
        if basis > 0.5:
            x(0)
        c = measure(0)  # noqa: F841

    measurements = _test_parametric_on_local_sim(main, {"basis": 0.5})
    assert all(val == 0 for val in measurements["c"])
    measurements = _test_parametric_on_local_sim(main, {"basis": 0.55})
    assert all(val == 1 for val in measurements["c"])


def test_param_neq():
    """Test parameters used in conditional not equals statements."""

    @aq.main
    def parametric(val: int):
        if val != 1:
            h(0)
        measure(0)

    expected = """OPENQASM 3.0;
input int[32] val;
qubit[1] __qubits__;
bool __bool_0__;
__bool_0__ = val != 1;
if (__bool_0__) {
    h __qubits__[0];
}
bit __bit_1__;
__bit_1__ = measure __qubits__[0];"""
    assert parametric.build().to_ir() == expected


def test_param_or():
    """Test parameters used in conditional `or` statements."""

    @aq.main
    def parametric(alpha: float, beta: float):
        if alpha or beta:
            rx(0, alpha)
            rx(0, beta)
        measure(0)

    expected = """OPENQASM 3.0;
input float alpha;
input float beta;
qubit[1] __qubits__;
bool __bool_0__;
__bool_0__ = alpha || beta;
if (__bool_0__) {
    rx(alpha) __qubits__[0];
    rx(beta) __qubits__[0];
}
bit __bit_1__;
__bit_1__ = measure __qubits__[0];"""
    assert parametric.build().to_ir() == expected


def test_param_and():
    """Test parameters used in conditional `and` statements."""

    @aq.main
    def parametric(alpha: float, beta: float):
        if alpha and beta:
            rx(0, alpha)
        measure(0)

    expected = """OPENQASM 3.0;
input float alpha;
input float beta;
qubit[1] __qubits__;
bool __bool_0__;
__bool_0__ = alpha && beta;
if (__bool_0__) {
    rx(alpha) __qubits__[0];
}
bit __bit_1__;
__bit_1__ = measure __qubits__[0];"""
    assert parametric.build().to_ir() == expected


def test_param_and_float():
    """Test parameters used in conditional `and` statements."""
    beta = 1.5

    @aq.main
    def parametric(alpha: float):
        if alpha and beta:
            rx(0, alpha)
        measure(0)

    expected = """OPENQASM 3.0;
input float alpha;
qubit[1] __qubits__;
bool __bool_0__;
__bool_0__ = alpha && 1.5;
if (__bool_0__) {
    rx(alpha) __qubits__[0];
}
bit __bit_1__;
__bit_1__ = measure __qubits__[0];"""
    assert parametric.build().to_ir() == expected


def test_param_not():
    """Test parameters used in conditional `not` statements."""

    @aq.main
    def parametric(val: int):
        if not val:
            h(0)
        measure(0)

    expected = """OPENQASM 3.0;
input int[32] val;
qubit[1] __qubits__;
bool __bool_0__;
__bool_0__ = !val;
if (__bool_0__) {
    h __qubits__[0];
}
bit __bit_1__;
__bit_1__ = measure __qubits__[0];"""
    assert parametric.build().to_ir() == expected


def test_parameter_binding_conditions():
    """Test that parameters can be used in conditions and then bound."""

    @aq.main
    def parametric(val: float):
        if val == 1:
            x(0)
        measure(0)

    template = """OPENQASM 3.0;
float val = {};
qubit[1] __qubits__;
bool __bool_0__;
__bool_0__ = val == 1;
if (__bool_0__) {{
    x __qubits__[0];
}}
bit __bit_1__;
__bit_1__ = measure __qubits__[0];"""
    bound_prog = parametric.build().make_bound_program({"val": 0})
    assert bound_prog.to_ir() == template.format(0)
    bound_prog = parametric.build().make_bound_program({"val": 1})
    assert bound_prog.to_ir() == template.format(1)


def test_parameter_expressions():
    """Test expressions of free parameters with numeric literals."""

    @aq.main
    def parametric(theta):
        expr = 2 * theta
        gpi(0, expr)

    @aq.main
    def parametric_fp():
        expr = 2 * FreeParameter("theta")
        gpi(0, expr)

    expected_1 = """OPENQASM 3.0;
input float theta;
qubit[1] __qubits__;
gpi(2 * theta) __qubits__[0];"""

    expected_2 = """OPENQASM 3.0;
input float theta;
qubit[1] __qubits__;
gpi(2.0 * theta) __qubits__[0];"""
    assert parametric.build().to_ir() == expected_1
    assert parametric_fp.build().to_ir() == expected_2


def test_sim_expressions():
    @aq.main
    def parametric(phi):
        rx(0, 2 * phi)
        measure(0)

    measurements = _test_parametric_on_local_sim(parametric, {"phi": np.pi / 2})
    assert 0 not in measurements["__bit_0__"]


def test_multi_parameter_expressions():
    """Test expressions of multiple free parameters."""

    @aq.main
    def parametric(alpha, theta):
        expr = alpha * theta
        gpi(0, expr)

    expected = """OPENQASM 3.0;
input float alpha;
input float theta;
qubit[1] __qubits__;
gpi(alpha * theta) __qubits__[0];"""
    assert parametric.build().to_ir() == expected


def test_bound_parameter_expressions():
    """Test expressions of free parameters bound to specific values."""

    @aq.main
    def parametric(phi):
        rx(0, 2 * phi)

    expected = """OPENQASM 3.0;
float phi = 1.5707963267948966;
qubit[1] __qubits__;
rx(2 * phi) __qubits__[0];"""
    assert parametric.build().make_bound_program({"phi": np.pi / 2}).to_ir() == expected


def test_partially_bound_parameter_expressions():
    """Test expressions of free parameters partially bound to specific values."""

    @aq.main
    def parametric(prefactor, theta):
        expr = prefactor * theta
        gpi(0, expr)

    expected = """OPENQASM 3.0;
float prefactor = 3;
input float theta;
qubit[1] __qubits__;
gpi(prefactor * theta) __qubits__[0];"""
    assert parametric.build().make_bound_program({"prefactor": 3}).to_ir() == expected


def test_subroutine_parameter_expressions():
    """Test expressions of free parameters passed to subroutines."""

    @aq.subroutine
    def rotate(theta: float):
        rx(0, 3 * theta)

    @aq.main
    def parametric(alpha):
        rotate(2 * alpha)

    expected = """OPENQASM 3.0;
def rotate(float[64] theta) {
    rx(3 * theta) __qubits__[0];
}
input float alpha;
qubit[1] __qubits__;
rotate(2 * alpha);"""
    assert parametric.build().to_ir() == expected


def test_gate_parameter_expressions():
    """Test expressions of free parameters passed to custom gates."""

    @aq.gate
    def rotate(q: aq.Qubit, theta: float):
        rx(q, 3 * theta)

    @aq.main
    def parametric(alpha):
        rotate(0, 2 * alpha)

    expected = """OPENQASM 3.0;
gate rotate(theta) q {
    rx(3 * theta) q;
}
input float alpha;
qubit[1] __qubits__;
rotate(2 * alpha) __qubits__[0];"""
    assert parametric.build().to_ir() == expected


def test_conditional_parameter_expressions():
    """Test expressions of free parameters contained in conditional statements."""

    @aq.main
    def parametric(phi):
        if 2 * phi > np.pi:
            h(0)
        measure(0)

    expected = """OPENQASM 3.0;
input float phi;
qubit[1] __qubits__;
bool __bool_0__;
__bool_0__ = 2 * phi > 3.141592653589793;
if (__bool_0__) {
    h __qubits__[0];
}
bit __bit_1__;
__bit_1__ = measure __qubits__[0];"""
    assert parametric.build().to_ir() == expected


def test_parameter_expressions_range_index():
    """Test expressions of free parameters contained in a range index."""

    @aq.main
    def two_n_hs(n: int):
        for _ in aq.range(2 * n):
            h(0)
        measure(0)

    expected = """OPENQASM 3.0;
input int[32] n;
qubit[1] __qubits__;
for int _ in [0:2 * n - 1] {
    h __qubits__[0];
}
bit __bit_0__;
__bit_0__ = measure __qubits__[0];"""
    assert two_n_hs.build().to_ir() == expected
    _test_parametric_on_local_sim(two_n_hs, {"n": 3})
