import socket
import numpy as np
from time import sleep


class SCPI_HTTP(object):

    def __init__(self, ip_address, port):
        self.socket = socket.create_connection((ip_address, port))
        self.socket.settimeout(2)  # s

    def communicate(self, message):
        self.socket.sendall(message)
        data_segments = []
        while True:
            try:
                data = self.socket.recv(2**12)
            except socket.timeout:
                break
            data_segments.append(data)
            if len(data) == 0 or data[-1] == '\n':
                break
        return ''.join(data_segments)

    def get_points_in_memory_count(self):
        return int(self.communicate(':acquire:mdepth?'))

    def get_channel_data(self, channel_i):
        data_points = self.get_points_in_memory_count()
        self.communicate(':waveform:source channel{0};:waveform:mode raw;:waveform:format ascii'.format(channel_i))
        sleep(0.5)
        data_all = []
        segment_size = 15625
        start_i = 1
        while True:
            stop_i = start_i + segment_size - 1
            self.communicate(':waveform:start {0};:waveform:stop {1}'.format(start_i, stop_i))
            data = self.communicate(':waveform:data?')
            data_all += [float(point) for point in data[11:-1].split(',')]
            start_i += segment_size
            if stop_i >= data_points:
                break
        return np.array(data_all)

    def get_sample_rate(self):
        return float(self.communicate(':acquire:srate?'))

        

if __name__ == '__main__':
    dev = SCPI_HTTP('192.168.2.99', 5555)
    channels_data = [dev.get_channel_data(i+1) for i in range(2)]
    dt = 1.0 / dev.get_sample_rate()
    time_axis = np.arange(channels_data[0].shape[0]) * dt
    channels_data.insert(0, time_axis)
    np.savetxt('data_all', np.vstack(channels_data).T)