#!/usr/bin/env python
# LeCrunch
# Copyright (C) 2010 Anthony LaTorre 
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import os
import sys
import time
import socket
import struct
import h5py
from lecroy import LeCroyScope
from config import get_settings

def fetch(filename, events):
    """
    Fetch and save waveform traces from the oscilloscope.

    Args:
        - filename: str
            Filename to store traces in (in hdf5 format).
        - events: int
            Number of triggered events to save in `filename`.
    """
    scope = LeCroyScope('scope01.hep.upenn.edu')
    scope.connect()

    # turn off the display
    scope.send('display off')
    scope.check_last_command()

    # clear the output queue
    scope.clear()

    # get active channels
    channels = scope.getchannels()

    # get scope configuration
    settings = get_settings(scope)

    # get wave descriptors for each channel
    # important to do this before queue is primed!
    wavedesc = {}
    for channel in channels:
        wavedesc[channel] = scope.getwavedesc(channel)

    # prime the output queue
    for i in range(10):
        scope.trigger()
        for channel in channels:
            scope.send('c%i:wf? dat1' % channel)

    # open up the output file
    f = h5py.File(filename, 'w')
    # set scope configuration
    for command, setting in settings.items():
        f.attrs[command] = setting
    dataset = {}
    for channel in channels:
        samples = wavedesc[channel]['wave_array_count']
        chunks = (max(1,min(100, events//100)), samples)
        dataset[channel] = f.create_dataset('channel%i' % channel,
                                            (events, samples),
                                            dtype=wavedesc[channel]['dtype'],
                                            chunks=chunks,
                                            compression='gzip')
        for key, value in wavedesc[channel].items():
            try:
                dataset[channel].attrs[key] = value
            except ValueError:
                pass

    # start a timer
    time0 = time.time()

    try:
        for i in range(events):
            print '\rSaving event: %i' % (i+1),
            sys.stdout.flush()

            try:
                scope.trigger()
                for channel in channels:
                    dataset[channel][i] = scope.getwaveform(channel,
                                                            wavedesc[channel])
            except (socket.error, struct.error, RuntimeError, OverflowError):
                print "\nFailed to receive waveform %i" % (i+1)
                # clear the output queue
                scope.clear()
                # reprime the queue
                for i in range(10):
                    scope.trigger()
                    for channel in channels:
                        scope.send('c%i:wf? dat1' % channel)
                continue
        print
    except KeyboardInterrupt:
        print
        print 'Resizing datasets...'
        for channel in channels:
            dataset[channel].resize((i, wavedesc[channel]['wave_array_count']))
        raise KeyboardInterrupt
    finally:
        f.close()
        scope.clear()
        scope.send('display on')
        scope.check_last_command()
        scope.close()

        elapsed = time.time() - time0

        print 'Completed %i events in %.3f seconds.' % (i+1, elapsed)
        print 'Averaged %.5f seconds per acquisition.' % (elapsed/(i+1))
        print "Wrote to file '%s'." % filename

if __name__ == '__main__':
    import optparse

    usage = "usage: %prog <filename/prefix> [-n] [-r]"
    parser = optparse.OptionParser(usage, version="%prog 0.1.0")
    parser.add_option("-n", type="int", dest="events",
                      help="number of events to store per run", default=1000)
    parser.add_option("-r", type="int", dest="runs",
                      help="number of runs", default=1)
    parser.add_option("--time", action="store_true", dest="time",
                      help="append time string to filename", default=False)
    (options, args) = parser.parse_args()

    if len(args) < 1:
        sys.exit(parser.format_help())
    
    if options.events < 1 or options.runs < 1:
        sys.exit("Please specify a number >= 1 for number of events/runs")

    if options.runs == 1 and not options.time:
        fetch(args[0], options.events)
    else:
        import time
        import string

        for i in range(options.runs):
            timestr = string.replace(time.asctime(time.localtime()), ' ', '-')
            filename = args[0] + '_' + timestr + '.hdf5'
            print '-' * 65
            print 'Saving to file %s' % filename
            print '-' * 65

            try:
                fetch(filename, options.events)
            except KeyboardInterrupt:
                break