import numpy as np
from enum import Enum
import socket
import struct

probe = Enum('probe', [('10mV', 0), 
                       ('20mV', 1), 
                       ('50mV', 2), 
                       ('100mV', 3), 
                       ('200mV', 4), 
                       ('500mV', 5), 
                       ('1V', 6), 
                       ('2V', 7), 
                       ('5V', 8),
                       ('10V', 9), 
                       ('20V', 10), 
                       ('50V', 11), 
                       ('100V', 12), 
                       ('200V', 13)])



# # # # # # # 3 # # # # #  #
# CONFIGURATION PARAMETERS #
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

# Sample rate in MHz
sample_rate = np.float32(10.0)
# Total measurement time in seconds
total_measure_time = np.float32(0.01)

# Get the number of samples by multiplying sample rate by total measurement time
num_samples = np.uint32(sample_rate * total_measure_time * 1e6)

# Number of channels
nchannels = np.uint32(2)
# Channel ranges in volts, represented as an array of uint8
channel_ranges = np.array([probe['1V'].value, probe['500mV'].value], dtype=np.uint8)

# Number of expected signals
num_trigger = np.uint32(64)
# Trigger direction: 1 for rising edge, 0 for falling edge
trig_dirrection = np.int32(1)
# Trigger pre-measurement time in percentage of the total measurement time
trig_premeasure = np.uint32(100)
# Trigger channel, 0 for channel 1, 1 for channel 2, etc.
trig_channel = np.uint8(0)
# Trigger threshold (0 is 0V and 32767 is max of channel range)
trig_threshold = np.int16(32767)

# Auto trigger time in milliseconds
auto_trigger_time = np.int16(10000)

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #



# # # # # # #  #
# MAIN PROGRAM #
# # # # # # #  # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
client_socket.settimeout(5.0)
client_socket.connect(('localhost', 5002))

# Open device
msg = struct.pack('<BB', 0xAA, 0x01)
client_socket.sendall(msg)

# Set sample rate
msg = struct.pack('<BBI', 0xAA, 0x07, np.uint32(sample_rate * 1e6))
client_socket.sendall(msg)

# Set number of samples
msg = struct.pack('<BBI', 0xAA, 0x17, num_samples)
client_socket.sendall(msg)

# Configure channels
msg = struct.pack('<BBIsBiHh', 0xAA, 0x09, nchannels, channel_ranges.tobytes(), trig_channel, trig_dirrection, trig_threshold, auto_trigger_time)
client_socket.sendall(msg)

# Set pre-measure time
msg = struct.pack('<BBI', 0xAA, 0x28, trig_premeasure)
client_socket.sendall(msg)

# Set trigger number
msg = struct.pack('<BBI', 0xAA, 0x19, num_trigger)
client_socket.sendall(msg)

# Start measurement
msg = struct.pack('<BB', 0xAA, 0x1B)
client_socket.sendall(msg)