#!/usr/bin/env python3

#**************************************************************************
#   Copyright (C) 2009-2023 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, expanduser
from os import umask, popen, system, getenv
import datetime

program = 'oms2v2d'
version = '1.19'
verdate = '20230318'
author = 'Walter Brisken, Mark Wainright'

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('  --sim')
	print('  -S       generate .v2d file for vdifsim\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')
	print('If environment variable VLBA_STATION_FORMAT is set,')
	print('the file at that location will be read and it will')
	print('be used to override the data format for that station.')
	print('Each line of this file should be of the form:')
	print('  <stn> <fmt>')
	print('Where <stn> is the station code (typically 2 letters) and')
	print('<fmt> is the format designator, e.g., VDIF/5032/2 .\n')
	exit(0)

vdifsimDotFile = expanduser('~/.vdifsim')

def loadvdifsimDotFile():
    try:
        data = open(vdifsimDotFile, 'r').readlines()
        print('Found %s' % vdifsimDotFile)
        computeMachines = []
        antennaMachines = {}
        for d in data:
            s = d.strip().split('#')[0].split()
            if len(s) == 1:
                computeMachines.append(s[0])
            elif len(s) >= 3:
                antennaMachines[s[0].upper()] = s[2]
        return computeMachines, antennaMachines
    except IOError:
        print('Note: %s not found' % vdifsimDotFile)
        return None, None

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

# read the dbs.start file and return a dictionary of DBS start date-times
def readDBSStart():
	dbsStart = {}

	difxBase = getenv('DIFX_BASE')
	if difxBase == None:
		return dbsStart
	dbsStartPath = difxBase + '/dbs.start'
	
	try:
		dbsStartLines = open(dbsStartPath).readlines()
		if len(dbsStartLines) > 0:
			for line in dbsStartLines:
				if line[0] != '#':
					pair = line.split()
					dbsStart[pair[0]] = pair[1]
	except FileNotFoundError:
		# continue if file not found
		pass

	return dbsStart

# check if DBS start is earlier than project info start and if format is Mark5B
def dbsFormatChange(site, dbsStart, oms):
	
	siteStart = dbsStart[site]

	# get PROJECT_INFO START date-time
	projectInfoStart = oms['PROJECT_INFO'][0]['START'][0]

	# get STATION_SETUP FORMAT
	stationSetupFormat = oms['STATION_SETUP'][0]['FORMAT'][0]

	if stationSetupFormat == 'MARK5B' and siteStart < projectInfoStart:
		return True
	else:
		return False

# 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, sim, datacopyFormat, evlbiType, nStream, profile, dbsStart, needClockEOP, 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'

	projectMJD = 0

	stations = []
	for s in oms['STATION_INFO']:
		stations.append(s['STATION_ID'][0])
		if projectMJD == 0:
			mjd0 = datetime.datetime(1858, 11, 17, 0, 0)
			mjdday = s['START'][0].split('/')[0].split('-')
			yr = int(mjdday[0])
			doy = int(mjdday[1])
			projectMJD = (datetime.datetime(yr, 1, 1) + datetime.timedelta(doy-1) - mjd0).days
			
	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')

	if sim:
		compute_machines, antenna_machines = loadvdifsimDotFile()
		if compute_machines == None or antenna_machines == None:
			print('Sim mode relies on a ~/.vdifsim file existing.')
			exit(0)

		machineString = ','.join(compute_machines)
		out.write('machines = %s\n' % machineString)
		out.write('nCore = %d\n' % (len(compute_machines) - 1))
		out.write('nThread = 12\n\n')
		out.write('singleScan = True\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')

	fmt = {}
	a = getenv('VLBA_STATION_FORMAT')
	if a:
		data = open(a).readlines()
		for d in data:
			s = d.strip().split()
			if len(s) < 2:
				continue
			if len(s) > 2:
				startMJD = float(s[2])
			else:
				startMJD = 0.0
			if startMJD < mjdNow():
				fmt[s[0].upper()] = s[1]

	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 sim:
			filelistFile = '%s.%s.filelist' % (prefix.lower(), s.lower())
			try:
				machine = antenna_machines[s.upper()]
			except KeyError:
				print('Station %s is not properly defined in %s.  Make sure both path and default machine are included.\n' % (s.upper(), vdifsimDotFile))
				exit(0)
			extra += 'filelist=%s machine=%s ' % (filelistFile, machine)
			if needClockEOP:
				extra += 'clockOffset=0 '
			
		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

		if s in fmt:
			f = 'format=' + fmt[s] + ' '
			if len(dsData) > 0:
				for i in range(len(dsData)):
					dsData[i] += f
			else:
				dsData.append(f)

		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 += ' '

		# check if DBS changed Mark5B to VDIF
		if len(dbsStart) > 0 and dbsFormatChange(s, dbsStart, oms) == True:
			# DBS changed record format to VDIF5032
			extra += 'format=VDIF5032 '

		out.write('ANTENNA %s { toneSelection=smart %s}\n' % (s, extra))
		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')

	if needClockEOP and projectMJD > 0:
		out.write('\n# NOTE! The following EOP values are totally fictitious!  Do not use for science!\n')
		for i in range(0, 5):
			M = projectMJD - 2 + i
			out.write('EOP %d { tai_utc=37 ut1_utc=0.236958 xPole=0.10597 yPole=0.53906 }\n' % M)

	out.close()



# main below here

umask(2)

print('')

inFile = ''
force = False
stop = False
profile = False
nrt = False	# near real time mode
filebased = False
sim = False
datacopyFormat = None
evlbiType = None
needClockEOP = False
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 ['-S', '--sim']:
			sim = 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.obs'):
	vexFile = prefix + '.vex.obs'
elif isfile(prefix + '.vex.preobs'):
	vexFile = prefix + '.vex.preobs'
elif isfile(prefix + '.vex'):
	if nrt or (evlbiType != None):
		vexFile = prefix + '.vex.preobs'
	elif sim:
		vexFile = prefix + '.vex'
		needClockEOP = True
	else:
		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)

dbsStart = readDBSStart()

writev2d(oms, v2dFile, prefix, nrt, filebased, sim, datacopyFormat, evlbiType, nStream, profile, dbsStart, needClockEOP, 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')