import socket


class RPSCPI:
    def connect(self, host):
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.connect((host, 5000))
        self.sock_f = self.sock.makefile()

    def close(self):
        self.sock_f.close()
        self.sock.close()

    def sendmsg(self, msg):
        self.sock.send(msg.encode() + b"\r\n")

    def recvmsg(self):
        return self.sock_f.readline().strip()

    def trigger(self):
        self.sendmsg("ACQ:START")
        self.sendmsg("ACQ:TRIG NOW")
        while True:
            self.sendmsg("ACQ:TRIG:STAT?")
            if self.recvmsg() == "TD":
                break

    def get_data(self, channel):
        self.sendmsg("ACQ:SOUR{}:DATA?".format(channel))
        return self.recvmsg()[1:-1]

def main():
    rp = RPSCPI()
    rp.connect("192.168.1.104")
    try:
        rp.trigger()
        y1_raw = rp.get_data(1)
        y2_raw = rp.get_data(2)

        with open('rp_y1_raw.bin', 'wb') as f:
            f.write(y1_raw.encode('utf-8'))
            print("Succesfully written y1 raw string from RP.")
        with open('rp_y2_raw.bin', 'wb') as f:
            f.write(y2_raw.encode('utf-8'))
            print("Succesfully written y2 raw string from RP.")

        # DEBUGGING
        y1 = [float(i) for i in y1_raw.split(',')]
        y2 = [float(i) for i in y2_raw.split(',')]
        with open('rp_y1_raw.txt', 'w') as f:
            for i in y1:
                f.write(str(i) + '\n')
            print("[DEBUG] Succesfully written y1 human-readable data.")
        with open('rp_y2_raw.txt', 'w') as f:
            for i in y2:
                f.write(str(i) + '\n')
            print("[DEBUG] Succesfully written y2 human-readable data.")

    finally:
        rp.close()

if __name__ == "__main__":
    main()