#!/usr/bin/env python3

# Note: this utility can run under python2.7 or python3

program = 'vdifspec'
author  = 'Walter Brisken'
version = '0.7'
verdate = '20210417'

from os import system
from sys import exit, argv

def printVersion():
	print('%s ver. %s  %s  %s' % (program, version, author, verdate))

def usage():
	print('')
	printVersion()
	print('')
	print('A VDIF spectrometer for multi-thread VDIF data.  Uses vmux')
	print('and m5spec to do the heavy lifting.\n')
	print('Usage : %s <infile> <frame size> <data rate> <threadlist> <nchan> <nint> <outfile> [<nbit> [<offset> [<bchan> [<echan>]]]]\n' % program)
	print('  <infile> is the name of the VDIF file\n')
	print('  <frame size> is the size of each input VDIF frame, inc. header (e.g., 5032)\n')
	print('  <data rate> is the stream data rate (Mbps)\n')
	print('  <threadlist> is a comma-separated list of threads to process\n')
	print('  <nchan> is the number of spectral channels to make per baseband channel\n')
	print('  <nint> is the number of FFT frames to spectrometize\n')
	print('  <outfile> is the name of the output file\n')
	print('  <nbit> is number of bits per sample (default is 2)\n')
	print('  <offset> optionally jump into input file by this many bytes\n')
	print('  <bchan> and <echan> select range of channels to store in output file\n')

if '--version' in argv:
	printVersion()
	exit()

if len(argv) < 8:
	usage()
	exit()

extraArgs = ''
infile = argv[1]
framesize = int(argv[2])
datarate = int(argv[3])
threadlist = argv[4]
nt = len(threadlist.split(','))
nchan = int(argv[5])
nint = int(argv[6])
outfile = argv[7]
if len(argv) > 8:
	nbit = int(argv[8])
else:
	nbit = 2
if len(argv) > 9:
	offset = int(argv[9])
else:
	offset = 0
if len(argv) > 10:
	extraArgs += ' -bchan=%s' % argv[10]
if len(argv) > 11:
	extraArgs += ' -echan=%s' % argv[11]

nc = 1
while(nc < nt):
	nc *= 2

outputpayload = nc*(framesize-32)
framespersecond = (datarate*1000000)//(nt*(framesize-32)*8)
outputdatarate = outputpayload*framespersecond*8//1000000

cmd = 'vmux %s %d %d %s - %d | m5spec %s - VDIF_%d-%d-%d-%d %d %d %s' % (infile, framesize, framespersecond, threadlist, offset, extraArgs, outputpayload, outputdatarate, nc, nbit, nchan, nint, outfile) 

print('Executing: %s' % cmd)
system(cmd)