import sys, os
import time
import gpib

import gpib
from prologix import GPIB
class GPIBplotter(GPIB):
    _replies={
        "OE": "0",
        "OH": "0,0,10000,7500",
        "OI": "7470A",
        "OP": "0,0,10000,7500",
        "OO": "0,1,0,0,0,0,0,0",
        "OF": "40,40",
        "OS": "24",
        }
    def __init__(self, device="/dev/ttyUSB0", baudrate=115200, timeout=0.1,
                 address=5):
        super(GPIBplotter, self).__init__(device, baudrate, timeout, address, mode=0)

    def plotStarted(self):
        pass

    def load_plot(self, wait_timeout=0):
        """
        Make a full plot process.

        'wait_timeout' is the first timeout. Same semantic as the
        'timeout' property of a serial connection (if set to None, it
        will block until the GPIB device actually perform a plot
        command.)
        """
        res = ""
        i=0
        replies = self._replies.copy()
        if wait_timeout is not None and not isinstance(wait_timeout, (float, int)):
            raise TypeError, "wait_timeout (%s:%s) has wrong type"%(repr(wait_timeout), type(wait_timeout))
        if wait_timeout<0:
            raise ValueError
        
        self._cnx.timeout = wait_timeout
        firstloop = True
        newdata = False
        while i<self._retries:
            l = self._cnx.readline().strip()
            if firstloop:
                self._cnx.timeout = self._timeout
                firstloop = False
                self.plotStarted()
            if l == "":
                if i == 0:# > (self._retries/2):
                    # ie. we just received new stuffs (i is reset in the else block)
                    for k, v in replies.items():
                        # check wether we should reply smthg
                        eres = res.replace('\n', '').strip()
                        if eres.endswith(k) or eres.endswith(k+';') or eres.endswith(k+';OE'):
                            self._cnx.write("%s"%v)
                            if k == "OS":
                                replies[k] = "16"
                            break
                    self._cnx.write('\r')
                    #time.sleep(0.1)
                i += 1
            else:
                if not res:
                    print "Plotting..."
                res += l + '\n'
                i = 0
                #time.sleep(0.1)
        if res:
            print "DONE (received %d characters)"%len(res)
        return res

if __name__ == '__main__':
    import optparse
    opt = optparse.OptionParser('A simple HP7470A GPIB plotter emulator for USB-GPIB bundle (ProLogix)')
    opt.add_option('-f', '--filename', default=None,
                   dest='filename',
                   help='Output filename. If not set, write to stdout')
    opt.add_option('-d', '--device', default='/dev/ttyUSB0',
                   dest='device',
                   help='Device of the RS232 connection (default: /dev/ttyUSB0)',
                   )
    opt.add_option('-a', '--address', default=0,
                   dest='address',
                   help='GPIB address of the device',
                   )
    opt.add_option('-l', '--loop', default=False,
                   action="store_true",
                   dest="loop",
                   help="Continuously wait for new plots. If set, filename must be set and files will be created with _n at the end")
    opt.add_option('-v', '--verbose', default=False,
                   action="store_true",
                   dest="verbose",
                   help="Verbose mode",)
    
    options, argv = opt.parse_args(sys.argv)

    if options.loop and not options.filename:
        opt.error('If loop is set, you *must* provide a filename')

    if options.filename:        
        outf = open(options.filename, "w")
    else:
        outf = sys.stdout


    try:
        plotter = GPIBplotter(device=options.device, address=int(options.address),
                              timeout=0.06)
    except (gpib.SerialException, gpib.ConnectionError), e:
        sys.stderr.write('Connection error:\n\t' + '\n\t'.join([str(x) for x in e.args]) + '\n')
        sys.stderr.write('Check your parameters\n')
        sys.exit(1)
    if options.verbose:
        sys.stderr.write('connection established\n')

    loop = True
    nloop = 0
    while loop:
        plot = plotter.load_plot(wait_timeout=0.1)
        if options.verbose:
            sys.stderr.write('.')
        if plot:
            outf.write(plot)
            if options.verbose:
                sys.stderr.write('\n')
            sys.stderr.write('Received a new plot (written to %s)\n'%outf.name) 
            if not options.loop:
                loop = False
            else:
                nloop += 1
                fname, ext = os.path.splitext(options.filename)
                outf = open(fname + "_%d"%nloop + ext, 'w')
                
        
    


