from rigol_dho_lib.rigol import RigolOsc
import time
from rigol_dho_lib import rigol
from siglent_sdg import siglent
import numpy as np

import matplotlib.pyplot as plt

SINE_FREQ = 100e3

CYCLE_COUNT = 1000

DIV_COUNT = 10


def setWindowSize(osc, cycles, frequency):
    period = 1 / frequency * cycles / DIV_COUNT
    osc.setTimescale(period)


def dft(data, time_array, freq):
    carrier = np.exp(2j * np.pi * freq * time_array)
    x = carrier * data
    return sum(x) / len(carrier) * 2


# uv
def main():

    try:
        gen = siglent.SiglentGen("TCPIP::10.112.1.17::INSTR")
    except ():
        print("Can't connent!")

    try:
        osc = rigol.RigolOsc("TCPIP::10.112.1.19::INSTR")
    except ():
        print("Can't connent!")

    mem_depth = osc.getMemoryDepth()

    osc.setPoints(mem_depth)

    setWindowSize(osc, CYCLE_COUNT, SINE_FREQ)

    osc.channels[1].setVScale(0.5)
    gen.channels[siglent.ChannelID.CH1].apply_sine(SINE_FREQ, 10, 0)
    gen.channels[siglent.ChannelID.CH1].set_output(True)

    time.sleep(1)

    osc.single()

    # time.sleep(0.2)
    channel_A_data = osc.channels[0].getWaveform()
    channel_B_data = osc.channels[1].getWaveform()

    time_array = osc.channels[0].genTimeArray(channel_A_data)

    v_a = dft(channel_A_data, time_array, SINE_FREQ)

    v_o = dft(channel_B_data, time_array, SINE_FREQ)

    # plt.plot(time_array, channel_A_data)
    # plt.plot(time_array, channel_B_data)

    # plt.show()

    # plt.plot(time_array, np.real(x))
    # plt.plot(time_array, np.imag(x))

    print(f"v_a: {v_a}")
    print(f"v_o: {v_o}")

    R0 = 10
    z = R0 * (v_a - v_o) / v_o

    print(f"impedance: {z}")

    return


if __name__ == "__main__":
    main()