#!/usr/bin/env python3

#**************************************************************************
#   Copyright (C) 2009-2020 by Walter Brisken                             *
#                                                                         *
#   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, write to the                         *
#   Free Software Foundation, Inc.,                                       *
#   59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             *
#**************************************************************************

#===========================================================================
# SVN properties (DO NOT CHANGE)
#
# $Id$
# $HeadURL: $
# $LastChangedRevision$
# $Author$
# $LastChangedDate$
#
#============================================================================

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

from sys import exit, argv
from os.path import isfile, isdir
from os import umask, popen, system

program = 'oms2v2d'
version = '1.16'
verdate = '20210428'
author = 'Walter Brisken'

def usage(pgm):
	print('%s ver. %s  %s  %s\n' % (program, version, author, verdate))
	print('Program to take sched\'s .oms file and produce a skeleton .v2d file\n')
	print('Usage: %s [options] <oms file>\n' % pgm)
	print('options can include\n')
	print('  --help')
	print('  -h       print this help info and quit\n')
	print('  --force')
	print('  -f       force operation, even if output file exits\n')
	print('  --verbose')
	print('  -v       be more verbose in execution (use -v -v for even more!)\n')
	print('  --nrt')
	print('  -n       use near-real-time options: .preobs and filelists\n')
	print('  --datacopy')
	print('  -d       look for files in /home/datacopy-* (VLBA only, assumes VDIF)\n')
	print('  --datacopy5b')
	print('           look for files in /home/datacopy-* (VLBA only, assumes Mark5B)\n')
	print('  --file')
	print('  -F       make .v2d file for file-based correlation\n')
	print('  --evlbi')
	print('  -e       make .v2d file for vlitebuf-based evlbi\n')
	print('  --split')
	print('  -s       make .v2d file for vlitebuf-based evlbi w/ 1032 byte frames\n')
	print('  --profile')
	print('  -p       generate .v2d for profile mode\n')
	print('  --dual')
	print('  -D       configure for two datastreams\n')
	print('  --quad')
	print('  -Q       configure for four datastreams\n')
	print('  --oct')
	print('  -O       configure for eight datastreams\n')
	print('  --quiet')
	print('  -q       be less verbose\n')
	print('<oms file> is a .oms file produced by sched\n')
	exit(0)


def readOMS(filename, verbose):
	if verbose > 1:
		print('Reading: %s' % filename)
	data = open(filename).readlines()
	section = []
	current = {}
	n = 0
	minbw = 1.0e9
	maxbw = 0.0
	spaces = ''
	for D in data:
		n = n + 1
		d = D.strip()
		if(len(d) < 1):
			continue
		p = d.find("=")
		if p < 0:
			key = '=COMMENT'
			value = d.strip()
		else:
			key = d[0:p].strip()
			if len(key) < 1:
				print('Illegal line [%d] : %s' % (n, d))
				exit(0)
			value = d[p+1:].strip()
		if key == 'BEGIN':
			section.append( (value, current) )
			last = current
			current = {}
			if not value in last:
				last[value] = []
			last[value].append(current)
			if verbose > 2:
				print('%sBEGIN %s' % (spaces, value))
			spaces = spaces + '  '
		elif key == 'END':
			spaces = spaces[:-2]
			if len(section) == 0:
				print('Error: line %d unmatched END: %s' % (n, d))
				exit(0)
			oldvalue, current = section.pop()
			if verbose > 2:
				print('%sEND %s' % (spaces, value))
			if oldvalue != value:
				print('END with value %s does not match previous BEGIN with value %s, line %d' % (value, oldvalue, n))
				exit(0)
		elif key == 'BANDWIDTH':
			if verbose > 3:
				print('%s%s = %s' % (spaces, key, value))
			for v in value.strip().split():
				f = float(v)
				if f > maxbw:
					maxbw = f
				if f < minbw:
					minbw = f
		else:
			if verbose > 3:
				print('%s%s = %s' % (spaces, key, value))
			if not key in current:
				current[key] = []
			for v in value.strip().split():
				current[key].append(v)

	if len(section) > 0:
		print('Error: unmatched BEGIN')
		exit(0)
	
	current['MINBW'] = minbw
	current['MAXBW'] = maxbw

	if verbose > 1:
		print('Min channel bandwidth = %f MHz' % minbw)
		print('Max channel bandwidth = %f MHz' % maxbw)

	return current

# return list of (host, filesystem)
def getDatacopyInfo():
	data = popen('dcadmin listfs', 'r')
	info = []
	for d in data:
		s = d.strip().split()
		if len(s) > 2:
			if s[-1] == 'Enabled':
				info.append( (s[1], s[2]) )

	return info

def getDatacopyAntennaSetup(datacopyInfo, prefix, stn, datacopyFormat):
	filelist = '%s.%s.filelist' % (prefix.lower(), stn.lower())
	if datacopyFormat == 'VDIF':
		cmd = 'vsum -s `dcls %s %s` > %s' % (prefix, stn, filelist)
	elif datacopyFormat == 'Mark5B':
		cmd = 'm5bsum -s `dcls %s %s` > %s' % (prefix, stn, filelist)
	print('Executing: %s' % cmd)
	system(cmd)

	# get first line of file
	try:
		firstfile = (open(filelist).readline().split())[0]
	except IndexError:
		return ''

	for machine, filesystem in datacopyInfo:
		if filesystem in firstfile:
			return 'filelist=%s machine=%s ' % (filelist, machine)

	return None

def getMachineFromFilelist(filelistFile):
	try:
		line1 = open(filelistFile).readline()
		s = line1.strip().split()
		if s[0] == '#' and s[1] == 'machine':
			return s[2]
	except:
		return None

# This determines if new VLBA station positions are in use
# If so, assume sched 11.6 (or newer) is in use and that new model should be used
def isNew(fn):
	if not isfile(fn) and fn[-4:] == '.obs':
		fn = fn[:-4]
	if not isfile(fn):
		print('Warning: vex file not found; assuming new delay model.\n')
		return True
		
	cmd = 'vexpeek -c %s' % fn
	data = popen(cmd).readlines()
	n = 0
	ants = ['BR', 'FD', 'HN', 'KP', 'LA', 'MK', 'NL', 'OV', 'PT', 'SC']
	for d in data:
		s = d.strip().split()
		if len(s) == 8:
			epoch = float(s[7])
			if epoch > 58848.5 and s[0] in ants:
				n += 1
	return (n > 0)

def dropzeros(s):
	ss = s.strip('0')
	if ss[-1] == '.':
		ss += '0'
	if ss[0] == '.':
		ss = '0' + ss
	return ss

def writev2d(oms, v2dFile, prefix, nrt, filebased, datacopyFormat, evlbiType, nStream, profile, verbose):
	machineMap = {'BR':'swc001', 'FD':'swc002', 'HN':'swc003', 'KP':'swc004', 'LA':'swc005', 'MK':'swc006', 'NL':'swc007', 'OV':'swc008', 'PT':'swc009', 'SC':'swc010', 'Y':'swc011', 'GB':'swc012'}
	corr = oms['CORRELATOR_PARAMETERS'][0]

	if verbose > 0:
		print('Writing: %s' % v2dFile)
	out = open(v2dFile, 'w')

	if datacopyFormat != None:
		datacopyInfo = getDatacopyInfo()

	bw = oms['MAXBW']
	if bw != oms['MINBW']:
		print('Warning: different bandwidths used within this file.')
		print('Using the maximum bandwidth to set spectral resolution.')

	try:
		nFFT = int(corr['FFT_SIZE'][0])
	except KeyError:
		print('Warning: FFT SIZE not provided.  Assuming 256.')
		nFFT = 256
	try:
		specAvg = int(corr['SPECTRAL_AVERAGE'][0])
	except KeyError:
		specAvg = 256//nFFT
		if specAvg == 0:
			specAvg = 1
	
	if profile:
		doPolar = 'False'
	elif 'POLARIZATION' in corr:
		if corr['POLARIZATION'][0] == 'NO':
			doPolar = 'False'
		else:
			doPolar = 'True'

	stations = []
	for s in oms['STATION_INFO']:
		stations.append(s['STATION_ID'][0])
	stations.sort()
	stnString = stations[0]
	for s in stations[1:]:
		stnString += ', ' + s

	out.write('# base .v2d file generated by %s version %s operating on file %s\n\n' % (program, version, omsFile))
	out.write('vex = %s\n\n' % vexFile)
	out.write('antennas = %s\n\n' % stnString)
	out.write('maxLength=3200\n\n')

	if filebased:
		out.write('delayModel = difxcalc+met\n\n')
	elif isNew(vexFile):
		out.write('delayModel = difxcalc\n\n')

	if nrt or filebased or (datacopyFormat != None) or (evlbiType != None):
		machineString = 'swc000, swc001, swc002, swc003, swc004, swc005, swc006, swc007, swc008, swc009, swc010'
		out.write('machines = %s\n' % machineString)
		out.write('nCore = 10\n')
		out.write('nThread = 12\n\n')
		out.write('singleScan = True\n\n')

	if profile:
		out.write('# profile mode selected:\n')
		out.write('mode = profile\n')
		out.write('startSeries = 100\n')
		out.write('minSubarray = 1\n\n')

	for s in oms['SOURCE_INFO']:
		srcParams = ''
		if 'CAL_CODE' in s:
			if len(s['CAL_CODE']) > 0:
				c = s['CAL_CODE'][0]
				if len(c) > 0:
					srcParams = srcParams + ' calCode = ' + c
		out.write('SOURCE %s {%s }\n' % (s['SOURCE_NAME'][0], srcParams))
	out.write('\n')

	for s in stations:
		extra = ''
		nds = 0
		dsData = []
		if nrt:
			extra += 'filelist=%s_%s.filelist machine=%s ' % (prefix.lower(), s.lower(), machineMap[s])
		elif filebased:
			filelistFile = '%s.%s.filelist' % (prefix.lower(), s.lower())
			machine = getMachineFromFilelist(filelistFile)
			if machine == None:
				machine = machineMap[s]
			if evlbiType == 'split':
				dsInfo = 'filelist=%s machine=%s ' % (filelistFile, machine)
				if evlbiType == 'split':
					dsInfo += 'frameSize=1032 '
				dsData.append(dsInfo)
			else:
				extra += 'filelist=%s machine=%s ' % (filelistFile, machine)
		elif evlbiType != None:
			for d in range(nStream):
				dsInfo = 'file=/tmp/vlitebuf_%02d/%%now.vdif%% machine=%s ' % (d, machineMap[s])
				if evlbiType == 'split':
					dsInfo += 'frameSize=1032 '
				dsData.append(dsInfo)
		elif datacopyFormat != None:
			datacopyAntennaSetup = getDatacopyAntennaSetup(datacopyInfo, prefix, s, datacopyFormat)
			if datacopyAntennaSetup == None:
				print('Warning: data for antenna %s not found.  Skipping.  You will need to update the antenna list in the .v2d file.' % s)
				out.write('# ANTENNA %s -- Data not found\n' % s)
				continue
			extra += datacopyAntennaSetup
		for ds in dsData:
			out.write('DATASTREAM %s%d { %s}\n' % (s, nds, ds))
			if nds == 0:
				extra += 'datastreams='
			else:
				extra += ','
			extra += '%s%d' % (s, nds)
			nds += 1
		if nds > 0:
			extra += ' '
		out.write('ANTENNA %s { toneSelection=smart %s}\n' % (s, extra))
		if nds > 0:
			out.write('\n')
	out.write('\n')

	out.write('SETUP default\n')
	out.write('{\n')

	if 'TIME_AVERAGE' in corr:
		out.write('  tInt = %s\n' % dropzeros(corr['TIME_AVERAGE'][0]))

	if specAvg != -1 and nFFT !=-1:
		out.write('  fftSpecRes = %s\n' % dropzeros('%10.8f' % (bw/(nFFT//2))))
		out.write('  specRes = %s\n' % dropzeros('%10.8f' % (bw/(nFFT//(2*specAvg)))))
	if 'POLARIZATION' in corr:
		out.write('  doPolar = %s\n' % doPolar)
	if nFFT > 256:
		out.write('  numBufferedFFTs = 10\n')
	out.write('  maxNSBetweenACAvg = 2000000\n')
	out.write('}\n')
	out.close()



# main below here

umask(2)

print('')

inFile = ''
force = False
stop = False
profile = False
nrt = False	# near real time mode
filebased = False
datacopyFormat = None
evlbiType = None
verbose = 1
nStream = 1
for a in argv[1:]:
	if a[0] == '-':
		if a in ['-h', '--help']:
			usage(argv[0])
		elif a in ['-f', '--force']:
			force = True
		elif a in ['-v', '--verbose']:
			verbose += 1
		elif a in ['-q', '--quiet']:
			verbose -= 1
		elif a in ['-n', '--nrt']:
			nrt = True
		elif a in ['-F', '--file']:
			filebased = True
		elif a in ['-d', '--datacopy']:
			datacopyFormat = 'VDIF'
		elif a in ['--datacopy5b']:
			datacopyFormat = 'Mark5B'
		elif a in ['-e', '--evlbi']:
			evlbiType = 'vlite'
		elif a in ['-s', '--split']:
			evlbiType = 'split'
		elif a in ['-p', '--profile']:
			profile = True
		elif a in ['-D', '--dual']:
			nStream = 2
		elif a in ['-Q', '--quad']:
			nStream = 4
		elif a in ['-O', '--oct']:
			nStream = 8
		else:
			print('Unknown command line option: %s' % a)
			stop = True
	else:
		if inFile == '':
			inFile = a
		else:
			print('Extra command line arg. given: %s' % a)
			stop = True

if inFile == '':
	print('No input .oms file provided.')
	stop = True

if stop:
	print('Please run with -h to get usage instructions.\n')
	exit(0)


if inFile[-4:] == '.oms':
	prefix = inFile[:-4]
else:
	prefix = inFile

omsFile = prefix + '.oms'
v2dFile = prefix + '.v2d'


if not isfile(omsFile):
	print('oms2v2d: Input oms file %s not found.  Quitting.\n' % omsFile)
	exit(0)

if isfile(prefix + '.vex.preobs'):
	vexFile = prefix + '.vex.preobs'
elif isfile(prefix + '.vex'):
	if nrt or (evlbiType != None):
		vexFile = prefix + '.vex.preobs'
	else:
		vexFile = prefix + '.vex.obs'
elif isfile(prefix + '.vex.obs'):
	vexFile = prefix + '.vex.obs'
else:
	print('No vex file found here.  It should be called %s.vex or %s.vex.preobs.' % (prefix, prefix))
	if force:
		if nrt:
			vexFile = prefix + '.vex.preobs'
		else:
			vexFile = prefix + '.vex.obs'
		print('Assuming vex file is called %s.\n' % (prefix + '.vex'))
	else:
		print('Run with --force to continue anyway.\n')
		exit(0)

if isfile(v2dFile):
	if force:
		if verbose > 0:
			print('Warning: overwriting exisiting file %s.\n' % v2dFile)
	else:
		print('Error: output file %s exists.  Run with --force to continue.\n' % v2dFile)
		exit(0)

oms = readOMS(omsFile, verbose)

writev2d(oms, v2dFile, prefix, nrt, filebased, datacopyFormat, evlbiType, nStream, profile, verbose)

if evlbiType != None:
	print('\nNOTE: if running in eVLBI mode for VLBA, consider using fiberLock.py to ensure routine network monitornig is turned off\n\n')