from htra_api import *
import matplotlib.pyplot as plt

###Open Device###
Status = 0          #Function return value
Device = c_void_p() #Memory address of the current device
DevNum = c_int(0)   #Specifies device number

BootProfile = BootProfile_TypeDef() #Boot configuration structure, including physical interface, power supply, etc.
BootInfo = BootInfo_TypeDef()       #Boot information structure, including device info, USB speed, etc.

BootProfile.DevicePowerSupply = DevicePowerSupply_TypeDef.USBPortAndPowerPort #Use both USB data port and independent power port
BootProfile.PhysicalInterface = PhysicalInterface_TypeDef.USB                 #Use USB interface for data transmission

#Configure ETH interface for Ethernet devices
#BootProfile.PhysicalInterface = PhysicalInterface_TypeDef.ETH #Use Ethernet for data transmission
#BootProfile.ETH_IPVersion = IPVersion_TypeDef.IPv4
#BootProfile.ETH_RemotePort = 5000
#BootProfile.ETH_ReadTimeOut = 10000
#BootProfile.ETH_IPAddress[0] = 192
#BootProfile.ETH_IPAddress[1] = 168
#BootProfile.ETH_IPAddress[2] = 1
#BootProfile.ETH_IPAddress[3] = 100

Status = dll.Device_Open(pointer(Device),DevNum,pointer(BootProfile),pointer(BootInfo)) #Open device
if(Status == 0):
    print("Device is opened successfully")
else:
    print("Return other errors Status = {:d}".format(Status))

###Configure IQS###
IQS_ProfileIn = IQS_Profile_TypeDef()  #IQS input configuration, including start frequency, stop frequency, RBW, reference level, etc.
IQS_ProfileOut = IQS_Profile_TypeDef() #IQS output configuration
StreamInfo = IQS_StreamInfo_TypeDef()  #Information on IQ data under current settings, including bandwidth, single-channel sampling rate, etc.
IQStream = IQStream_TypeDef()          #Stores IQ data packets, including IQ data and configuration information

dll.IQS_ProfileDeInit(pointer(Device),pointer(IQS_ProfileIn)) #Initialize parameters for IQS mode

IQS_ProfileIn.CenterFreq_Hz = 1e9                           #Set center frequency
IQS_ProfileIn.RefLevel_dBm = 0                              #Set reference level
IQS_ProfileIn.DecimateFactor = 2                            #Set decimation factor
IQS_ProfileIn.DataFormat = DataFormat_TypeDef.Complex16bit  #Set IQ data format
IQS_ProfileIn.TriggerSource = IQS_TriggerSource_TypeDef.Bus #Set trigger source to internal bus trigger
IQS_ProfileIn.BusTimeout_ms = 5000
IQS_ProfileIn.TriggerMode = TriggerMode_TypeDef.FixedPoints
IQS_ProfileIn.TriggerLength = 16242                         #Set the number of points collected per trigger event

Status = dll.IQS_Configuration(pointer(Device),pointer(IQS_ProfileIn),pointer(IQS_ProfileOut),pointer(StreamInfo)) #Apply IQS mode configuration
if(Status == 0):
    print("Configuration applied successfully")
else:
    print("Configuration failed, error code{:d}".format(Status))

###Retrieve Data###
AlternIQStream_data = (c_int16 * (StreamInfo.StreamSamples * 2))()  #Create array to store IQ data

#Enable DSP functionality
DSP = c_void_p()
dll.DSP_Open(pointer(DSP)) #Open DSP

#Open DSP
IQToSpectrumIn = DSP_FFT_TypeDef()  #Configure FFT mode parameters
IQToSpectrumOut = DSP_FFT_TypeDef() #Feedback of actual FFT mode parameters applied
TracePoints = c_uint32(0)           #Number of frequency spectrum points after FFT

dll.DSP_FFT_DeInit(pointer(IQToSpectrumIn)) #Initialize FFT mode parameters

IQToSpectrumIn.Calibration = 0                                             #Enable or disable calibration (0 = off, other values = on)
IQToSpectrumIn.DetectionRatio = 1                                          #Set detection ratio
IQToSpectrumIn.TraceDetector = TraceDetector_TypeDef.TraceDetector_PosPeak #Set detection method
IQToSpectrumIn.FFTSize = StreamInfo.StreamSamples                          #Set FFT size
IQToSpectrumIn.Intercept = 1                                               #Set interception ratio
IQToSpectrumIn.SamplePts = StreamInfo.StreamSamples                        #Set number of sampling points
IQToSpectrumIn.WindowType = Window_TypeDef.FlatTop                         #Set window type

RBWRatio = c_double(0) #This parameter returns the RBW ratio，RBW = RBWRatio * StreamInfo.IQSampleRate。
dll.DSP_FFT_Configuration(pointer(DSP),pointer(IQToSpectrumIn),pointer(IQToSpectrumOut),pointer(TracePoints),pointer(RBWRatio))
Frequency = (c_double * TracePoints.value)()    #Array to store frequency data
PowerSpec_dBm = (c_float * TracePoints.value)() #Array to store power spectrum data

#Using pyplot for plotting
fig, ax = plt.subplots()
line, = ax.plot([], [], label='Spectrum') #Initialize an empty plot line
ax.set_title('IQSToSpectrum')
ax.set_xlabel('Frequency')
ax.set_ylabel('PowerSpec_dBm')

ax.set_xlim(IQS_ProfileIn.CenterFreq_Hz-StreamInfo.IQSampleRate/2, IQS_ProfileIn.CenterFreq_Hz+StreamInfo.IQSampleRate/2)
ax.set_ylim(-160, 25)
ax.legend()
# Keep the plot open after the loop ends (interactive mode)
plt.ion()
plt.show()

while(True):

    Status = dll.IQS_BusTriggerStart(pointer(Device)) #Call IQS_BusTriggerStart to trigger the device. If the trigger source is external trigger, this function is not needed.

    for i in range(0,StreamInfo.PacketCount):
        dll.IQS_GetIQStream_PM1(pointer(Device),pointer(IQStream)) #Get IQ data packet, trigger info, I-channel max value, and its array index

        if i != StreamInfo.PacketCount - 1:  #Not the last packet
          samples_to_copy = StreamInfo.PacketSamples * 2
        else: #For the last packet, determine the number of samples to copy
          samples_to_copy = (StreamInfo.StreamSamples % StreamInfo.PacketSamples) * 2 if StreamInfo.StreamSamples % StreamInfo.PacketSamples != 0 else StreamInfo.PacketSamples * 2

        AlternIQStream_type = c_int16 * (samples_to_copy)                             #Define array type to extract data from IQStream.AlternIQStream
        source_data_ptr = cast(IQStream.AlternIQStream, POINTER(AlternIQStream_type)) #Pointer type conversion
        source_data = source_data_ptr.contents                                        #Get data and store in source_data
        for j in range(samples_to_copy):                                              #Loop to concatenate data into AlternIQStream_data 
            AlternIQStream_data[i * (StreamInfo.PacketSamples * 2) + j] = source_data[j]

    raw_data_ptr = pointer(AlternIQStream_data)                     #Create pointer pointing to concatenated data
    IQStream.AlternIQStream = cast(raw_data_ptr, POINTER(c_void_p)) #Convert raw data pointer to void* type pointer

    dll.DSP_FFT_IQSToSpectrum(pointer(DSP), pointer(IQStream), Frequency, PowerSpec_dBm) #Execute IQ-to-spectrum conversion function
    line.set_xdata(Frequency)                                                            #Plot spectrum
    line.set_ydata(PowerSpec_dBm)

    ax.relim()
    ax.autoscale_view()
    plt.draw()
    plt.pause(0.01)

Status = dll.IQS_BusTriggerStop(pointer(Device)) #Call IQS_BusTriggerStop to stop triggering the device
###Close the device###
Status = dll.Device_Close(pointer(Device))
if(Status == 0):
    print("Device closed successfully")
else:
    print("Device failed to close")


