from c3.signal.pulse import Envelope, EnvelopeNetZero
from c3.c3objs import Quantity as Qty
from c3.libraries import envelopes
import numpy as np
import pytest

flux_env = Envelope(
    name="flux",
    desc="Flux bias for tunable coupler",
    shape=envelopes.rect,
    params={
        "amp": Qty(value=0.1, unit="V"),
        "t_final": Qty(value=10e-9, unit="s"),
        "freq_offset": Qty(value=0, min_val=0, max_val=1, unit="Hz 2pi"),
        "xy_angle": Qty(value=0, min_val=0, max_val=np.pi, unit="rad"),
    },
)

flux_env_netzero = EnvelopeNetZero(
    name="flux",
    desc="Flux bias for tunable coupler",
    shape=envelopes.rect,
    params={
        "amp": Qty(value=0.1, unit="V"),
        "t_final": Qty(value=5e-9, unit="s"),
        "freq_offset": Qty(value=0, min_val=0, max_val=1, unit="Hz 2pi"),
        "xy_angle": Qty(value=0, min_val=0, max_val=np.pi, unit="rad"),
    },
)


@pytest.mark.unit
def test_envelope() -> None:
    ts = np.arange(0, 12e-9, 1e-9)
    shape = flux_env.get_shape_values(ts)
    assert np.all(
        shape.numpy()
        == np.array([1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0])
    )


@pytest.mark.unit
def test_envelope_netzero() -> None:
    ts = np.arange(0, 12e-9, 1e-9)
    shape = flux_env_netzero.get_shape_values(ts)
    assert np.all(
        shape.numpy()
        == np.array([1.0, 1.0, 1.0, 1.0, 1.0, 0.0, -1.0, -1.0, -1.0, -1.0, -1.0, -0.0])
    )
