#!/bin/env python3

import usbtmc

READ_CHUNK_SIZE = 256

class SSA3023(object):
    def __init__(self):
        self.instr =  usbtmc.Instrument(0xf4ec, 0x1300)
        #TODO: Check if ID number matches
        #id_str = self.instr.ask("*IDN?")

        self.instr.write(":FORMat:TRACe:DATA ASCii")

    def get_start_freq(self):
        return float(self.instr.ask(":SENSe:FREQuency:STARt?"))

    def get_stop_freq(self):
        return float(self.instr.ask(":SENSe:FREQuency:STOP?"))

    def get_ref_level(self):
        return float(self.instr.ask(":DISPlay:WINDow:TRACe:Y:SCALe:RLEVel?"))

    def get_norm_ref_level(self):
        return float(self.instr.ask(":DISPlay:WINDow:TRACe:Y:NRLevel?"))

    def is_displayed_trace(self, num):
        if num < 1 or num > 4:
            raise ValueError("%d is not a valid trace %d", num)
        res = self.instr.ask(":TRACe%d:MODE?" % num)

        return res != "BLANk"

    def get_trace(self, num):
        if num < 1 or num > 4:
            raise ValueError("%d is not a valid trace %d", num)

        self.instr.write(":TRACe:DATA? %d" % num)
        data = b''
        while not data.endswith(b'\n'):
            chunk = self.instr.read_raw(num=READ_CHUNK_SIZE)
            data = data + chunk

        data = data.decode("utf-8").split(',')[:-1]
        data = map(lambda x: float(x), data)

        return list(data)