#!/usr/bin/env python

"""
Condense a raw sample file generated by USPR, removing idle periods.
"""

import os
import argparse
import scipy
import itertools
import struct

WINDOW = 80


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('file', help="Raw sample file")
    parser.add_argument('--thres', type=int, default=1000, help="Power threshold")
    parser.add_argument('--out', help="Output file")
    parser.add_argument('--skip', type=int, default=int(0.1*20e6),
                        help="Inital samples to skip (USRP stabilization period)")

    args = parser.parse_args()

    wave = scipy.fromfile(args.file, dtype=scipy.int16)
    raw_samples = [complex(i, q) for i, q in
                   itertools.izip(wave[::2], wave[1::2])]
    print "%d raw samples" % (len(raw_samples))

    if args.out is None:
        args.out = '_condensed'.join(os.path.splitext(args.file))

    fh = open(args.out, 'wb')

    state = 'idle'
    count = 0
    countdown = 0
    for i in xrange(0, len(raw_samples), WINDOW):
        if i < args.skip:
            continue

        if state == 'idle':
            if any([abs(c.real) > args.thres for c in raw_samples[i:i+WINDOW]]):
                state = 'trigger'
                countdown = 80*10
                count += 1

        if state == 'trigger':
            if countdown > 0:
                do_write = True
            elif any([abs(c.real) > args.thres for c in raw_samples[i:i+WINDOW]]):
                do_write = True
            else:
                do_write = False
                state = 'idle'

            if do_write:
                for c in raw_samples[i:i+WINDOW]:
                    fh.write(struct.pack('<hh', int(c.real), int(c.imag)))
                countdown -= WINDOW

    fh.close()
    print "%d packets" % (count)
    print "Output written to %s" % (args.out)


if __name__ == '__main__':
    main()