""" Testing module for the generator and devices classes"""

import pickle
import numpy as np
import pytest
from c3.generator.devices import LO, AWG, Mixer, Response, DigitalToAnalog, VoltsToHertz
from c3.generator.generator import Generator
from c3.signal.gates import Instruction
from c3.signal.pulse import Envelope, Carrier
from c3.c3objs import Quantity
import c3.libraries.envelopes as env_lib

sim_res = 100e9  # Resolution for numerical simulation
awg_res = 2e9  # Realistic, limited resolution of an AWG

lo = LO(name="lo", resolution=sim_res, outputs=1)
awg = AWG(name="awg", resolution=awg_res, outputs=1)
dac = DigitalToAnalog(name="dac", resolution=sim_res, inputs=1, outputs=1)
resp = Response(
    name="resp",
    rise_time=Quantity(value=0.3e-9, min_val=0.05e-9, max_val=0.6e-9, unit="s"),
    resolution=sim_res,
    inputs=1,
    outputs=1,
)
mixer = Mixer(name="mixer", inputs=2, outputs=1)
v_to_hz = VoltsToHertz(
    name="v_to_hz",
    V_to_Hz=Quantity(value=1e9, min_val=0.9e9, max_val=1.1e9, unit="Hz/V"),
    inputs=1,
    outputs=1,
)
generator = Generator(
    devices={
        "LO": lo,
        "AWG": awg,
        "DigitalToAnalog": dac,
        "Response": resp,
        "Mixer": mixer,
        "VoltsToHertz": v_to_hz,
    },
    chains={
        "d1": ["LO", "AWG", "DigitalToAnalog", "Response", "Mixer", "VoltsToHertz"]
    },
)

t_final = 7e-9  # Time for single qubit gates
sideband = 50e6
gauss_params_single = {
    "amp": Quantity(value=0.5, min_val=0.4, max_val=0.6, unit="V"),
    "t_final": Quantity(
        value=t_final, min_val=0.5 * t_final, max_val=1.5 * t_final, unit="s"
    ),
    "sigma": Quantity(
        value=t_final / 4, min_val=t_final / 8, max_val=t_final / 2, unit="s"
    ),
    "xy_angle": Quantity(
        value=0.0, min_val=-0.5 * np.pi, max_val=2.5 * np.pi, unit="rad"
    ),
    "freq_offset": Quantity(
        value=-sideband - 3e6, min_val=-56 * 1e6, max_val=-52 * 1e6, unit="Hz 2pi"
    ),
    "delta": Quantity(value=-1, min_val=-5, max_val=3, unit=""),
}
gauss_env_single = Envelope(
    name="gauss",
    desc="Gaussian comp for single-qubit gates",
    params=gauss_params_single,
    shape=env_lib.gaussian_nonorm,
)


lo_freq_q1 = 5e9 + sideband
carrier_parameters = {
    "freq": Quantity(value=lo_freq_q1, min_val=4.5e9, max_val=6e9, unit="Hz 2pi"),
    "framechange": Quantity(value=0.0, min_val=-np.pi, max_val=3 * np.pi, unit="rad"),
}
carr = Carrier(
    name="carrier", desc="Frequency of the local oscillator", params=carrier_parameters
)

RX90p_q1 = Instruction(name="RX90p", t_start=0.0, t_end=t_final, channels=["d1"])
RX90p_q1.add_component(gauss_env_single, "d1")
RX90p_q1.add_component(carr, "d1")

tstart = RX90p_q1.t_start
tend = RX90p_q1.t_end
chan = "d1"

with open("test/generator_data.pickle", "rb") as filename:
    data = pickle.load(filename)


@pytest.mark.unit
def test_LO() -> None:
    lo_sig = lo.process(RX90p_q1, "d1")
    assert (
        lo_sig["inphase"].numpy() - data["lo_sig"]["values"][0].numpy() < 1e-12
    ).all()
    assert (
        lo_sig["quadrature"].numpy() - data["lo_sig"]["values"][1].numpy() < 1e-12
    ).all()
    assert (lo_sig["ts"].numpy() == data["lo_sig"]["ts"].numpy()).all()


@pytest.mark.unit
def test_AWG() -> None:
    awg_sig = awg.process(RX90p_q1, "d1")
    assert (
        awg_sig["inphase"].numpy() - data["awg_sig"]["inphase"].numpy() < 1e-12
    ).all()
    assert (
        awg_sig["quadrature"].numpy() - data["awg_sig"]["quadrature"].numpy() < 1e-12
    ).all()


@pytest.mark.unit
def test_DAC() -> None:
    dac_sig = dac.process(RX90p_q1, "d1", data["awg_sig"])
    assert (
        dac_sig["inphase"].numpy() - data["dig_to_an_sig"]["inphase"].numpy() < 1e-12
    ).all()
    assert (
        dac_sig["quadrature"].numpy() - data["dig_to_an_sig"]["quadrature"].numpy()
        < 1e-12
    ).all()


@pytest.mark.unit
def test_Response() -> None:
    resp_sig = resp.process(RX90p_q1, "d1", data["dig_to_an_sig"])
    assert (
        resp_sig["inphase"].numpy() - data["resp_sig"]["inphase"].numpy() < 1e-12
    ).all()
    assert (
        resp_sig["quadrature"].numpy() - data["resp_sig"]["quadrature"].numpy() < 1e-12
    ).all()


@pytest.mark.unit
def test_mixer() -> None:
    # For compatiblity with old dataset we need to wrap the LO output
    lo_signal = {
        "inphase": data["lo_sig"]["values"][0],
        "quadrature": data["lo_sig"]["values"][1],
        "ts": data["lo_sig"]["ts"],
    }
    mixed_sig = mixer.process(RX90p_q1, "d1", lo_signal, data["resp_sig"])
    assert (mixed_sig["values"].numpy() - data["mixer_sig"].numpy() < 1e-12).all()


@pytest.mark.unit
def test_v2hz() -> None:
    mixer_sig = {"values": data["mixer_sig"], "ts": data["lo_sig"]["ts"]}
    final_sig = v_to_hz.process(RX90p_q1, "d1", mixer_sig)
    assert (final_sig["values"].numpy() - data["v2hz_sig"].numpy() < 1).all()


@pytest.mark.integration
def test_full_signal_chain() -> None:
    full_signal = generator.generate_signals(RX90p_q1)
    assert (
        full_signal["d1"]["values"].numpy()
        - data["full_signal"][0]["d1"]["values"].numpy()
        < 1
    ).all()
