#!/bin/env python

# TODO - change USNO processing to check if schedule only contains MK and PT antennas

from sys import argv, exit, stdin
from os import system, chdir, getcwd, environ, popen, stat, path, remove, listdir
from os.path import isfile, isdir
from string import replace, strip, split, lower, upper, find
from glob import glob
from time import time
import datetime

program = 'queueVex'
version = 1.14
verdate = '20130424'
author  = 'Walter Brisken'

verbose = 1

observTxFile = 'observ.tx'
ccScriptDir = '/export/home/cc/vlba/scripts'
prestartTime = 300	# seconds
poststopTime = 30	# seconds
difxPath = '/home/swc/DiFX-trunk/bin'
observTxHistoryDir = '/home/vlba/metadata/observtxhistory'
observTxHistoryDays = 7
linesToShow = 10
keepDays = 45

# ------------------

def usage():
	print '\n%s ver. %s  %s  %s\n' % (program, version, verdate, author)
	print 'Usage: %s [options] <vexfile1> [vexfile2] ... [ant1] [ant2] ...\n' % argv[0]
	print 'options can be:\n'
	print '  -h or --help                print help info and quit.'
	print '  -v or --verbose             make more verbose.'
	print '  -q or --quiet               make less verbose.'
	print '  -f or --force               proceed without question in some cases.'
	print '  -p or --priority            interrupt a conflicting obs.'
	print '  -u or --usno                convert scripts to use USNO devices'
	print '  -e or --edit                allow time to edit the files'
	print '  -n or --nocopy              don\'t copy files to stations'
	print '  -d or --dbefile <filename>  use provided personality file for DBE\n'
	print ''
	print 'One or more vex files can be supplied.  All supplied vex files will'
	print 'be handled separately in the order listed.'
	print ''
	print 'Optional list of antenna codes can be used to restrict which antennas'
	print 'to operate on.  If no antennas are listed, all antennas in the vex'
	print 'file(s) will be assumed.'



mjd0 = datetime.datetime(1858, 11, 17, 0, 0)


def deleteOld(dir, days):
	n = 0
	now = time()
	print ''
	for f in listdir(dir):
		fn = path.join(dir, f)
		age = (now - stat(fn).st_mtime)/86400.0	# in days
		if age > days:
			if path.isfile(fn) and f[0:6] == 'observ':
				print 'Removing %s because it is %3.1f days old' % (fn, age)
				remove(fn)
				n += 1
	print '%d old observ.tx files from %s were deleted because they were older than %d days\n' % (n, dir, days)

def saveObservTxFile(srcfile, station):
	todayMJD = time()/86400.0 + 40587

	destfile = path.join(observTxHistoryDir, 'observ.tx_%s_%12.6f' % (station, todayMJD))

	cmd = 'cp %s %s' % (srcfile, destfile)
	if verbose > 0:
		print 'Executing command: %s' % cmd
	system(cmd)

def execute(cmd, nostdout=0):
	if verbose > 0:
		print 'About to execute: %s' % cmd
	if nostdout == 0:
		v = system(cmd)
	else:
		v = -1
		popen(cmd).readlines()
	
	return v

def vexPeek(vexFile):
	
	cmd = '%s/vexpeek %s' % (difxPath, vexFile)
	if verbose > 0:
		print 'Executing command: %s' % cmd

	p = popen(cmd)
	data = p.readlines()

	if len(data) == 0:
		return 'Error', 'Error', 'Error'
	
	obsCode = upper(strip(data[0]))
	obsSeg = ''

	if obsCode[0:5] == 'ERROR':
		return 'Error', 'Error', 'Error'

	if len(obsCode) > 3:
		if obsCode[0].isalpha() and obsCode[1].isalpha() and obsCode[2].isdigit():
			for i in range(3, len(obsCode)):
				if obsCode[i].isalpha():
					obsSeg = obsCode[i:]
					obsCode = obsCode[0:i]
					break
		if obsCode[0].isalpha() and obsCode[1].isdigit():
			for i in range(2, len(obsCode)):
				if obsCode[i].isalpha():
					obsSeg = obsCode[i:]
					obsCode = obsCode[0:i]
					break
		

	stationTimes = {}
	for d in data[1:]:
		s = split(strip(d))
		stationTimes[upper(s[0])] = [float(s[1])-prestartTime/86400.0, float(s[2])+poststopTime/86400.0]

	print 'This is experiment %s %s' % (obsCode, obsSeg)

	return obsCode, obsSeg, stationTimes

def mjd2datetime(mjd):
	mon = ['JAN', 'FEB', 'MAR', 'APR', 'MAY', 'JUN', 'JUL', 'AUG', 'SEP', 'OCT', 'NOV', 'DEC']
	dt = datetime.timedelta(int(mjd), int((mjd - int(mjd))*86400.0 + 0.5))
	t = mjd0 + dt
	return '%04d%s%02d %02dh%02dm%02ds' % (t.year, mon[t.month-1], t.day, t.hour, t.minute, t.second)

# take string date (d) and time (t), as formatted above, and return MJD
def vlba2mjd(d, t):
	monnames = ['JAN', 'FEB', 'MAR', 'APR', 'MAY', 'JUN', 'JUL', 'AUG', 'SEP', 'OCT', 'NOV', 'DEC']
	yr = int(d[0:4])
	monname = d[4:7]
	mn = 0
	for m in range(len(monnames)):
		if monname == monnames[m]:
			mn = m+1
	da = int(d[7:9])
	hr = int(t[0:2])
	mi = int(t[3:5])
	se = int(t[6:8])
	mjd = (datetime.datetime(yr, mn, da, 0, 0) - mjd0).days + hr/24.0 + mi/1440.0 + se/86400.0

	return mjd

def sortmetric(x):
        y = split(x)
	return vlba2mjd(y[1], y[2])

# modify observ.tx contents
def addObservation(orig, expt, startMJD, stopMJD, station, overlaps):
	startDate = mjd2datetime(startMJD)
	stopDate = mjd2datetime(stopMJD)
	todayMJD = time()/86400.0 + 40587
	epsilon = 5/86400.0	# extra gap to add between experiments in overlap situations

	observations = []

	for o in orig:
		v = split(strip(o), '>')
		s = split(v[0])
		if len(s) < 5 or len(v) < 2:
			continue
		expt0 = s[0]
		if expt0 == expt:
			continue	# don't keep old entry for same expt
		startDate0 = s[1]
		startMJD0 = vlba2mjd(s[1], s[2])
		stopMJD0 = vlba2mjd(s[3], s[4])
		station0 = v[1]

		# discard old observations
		if stopMJD0 < todayMJD-keepDays:
			continue

		# check for overlap with another observation
		if startMJD < stopMJD0 and stopMJD > startMJD0:	# some overlap condition exists!

			overlapLength = min(stopMJD0, stopMJD) - max(startMJD0, startMJD)

			overlaps.append('%s overlaps with %s by %4.2f hours at station %s' % (expt, expt0, (overlapLength*24), station))
			if startMJD0 < startMJD - epsilon:	# original project starts first
				newobs = '%s %s %s > %s\n' % (expt0, mjd2datetime(startMJD0), mjd2datetime(startMJD - epsilon), station0)
				observations.append(newobs)
			if stopMJD0 > stopMJD + epsilon:		# original projects ends last
				newobs = '%s %s %s > %s\n' % (expt0, mjd2datetime(stopMJD + epsilon), mjd2datetime(stopMJD0), station0)
				observations.append(newobs)
		else:	# no overlap, so this is easy!
			observations.append(o)
	newobs = '%s %s %s > %s\n' % (expt, startDate, stopDate, station)
	observations.append(newobs)

	# sort by start time
	observations.sort(key=sortmetric)

	return observations

# -----------------

vexFiles = []
antennas = []
priority = 0
usnoConvert = False
edit = False
force = False
nocopy = False
haveFile = False
file = ''
options = ''

for a in argv[1:]:
	if haveFile:
		file = a
		haveFile = False
	elif a[0] == '-':
		if a in ['-h', '--help']:
			usage()
			exit(1)
		elif a in ['-p', '--priority']:
			priority += 1
		elif a in ['-v', '--verbose']:
			verbose += 1
		elif a in ['-q', '--quiet']:
			verbose -= 1
		elif a in ['-u', '--usno']:
			usnoConvert = True
		elif a in ['-f', '--force']:
			force = True
		elif a in ['-e', '--edit']:
			edit = True
		elif a in ['-n', '--nocopy']:
			nocopy = True
		elif a in ['-d', '--dbefile']:
			haveFile = True
		else:
			print 'Unknown option "%s"\n' % a
			exit(1)
	elif len(a) > 3:
		vexFiles.append(a)
	else:
		antennas.append(lower(a))

if len(vexFiles) == 0:
	usage()
	exit(0)

if len(vexFiles) > 1:
	print 'The following vex files will be queued:'
	for vexFile in vexFiles:
		print '  %s' % vexFile
	if not force:
		print 'If this is not what you intend, use ^C now to stop.'
		print 'Otherwise press Enter to continue.'
		stdin.readline()

cwd = getcwd()

if find(upper(cwd), 'MARK5C') < 0:
	if isdir(cwd+'/MARK5C') > 0:
		if not force:
			print '\nWarning: %s/MARK5C is presumably where you want to run this.' % cwd
			print 'Try again with -f to force the operation'
		else:
			print '\nContinuing because of -f (or --force) even though %s/MARK5C' % cwd
			print 'seems more likely to be where you want to run this.'
	elif isdir(cwd+'/../MARK5C') > 0:
		if not force:
			print '\nWarning: %s/../MARK5C is presumably where you want to run this.' % cwd
			print 'Try again with -f to force the operation'
		else:
			print '\nContinuing because of -f (or --force) even though %s/MARK5C' % cwd
			print 'seems more likely to be where you want to run this.'
	else:
		print '\nYou are running %s from a non-Mark5C directory.' % program
		print 'I hope this is what you want!\n'
		options = options + " --mark5A"
	
for vexFile in vexFiles:
	if not isfile(vexFile):
		print 'Error: cannot find file %s' % vexFile
		exit(0)


tmpDir = '/tmp/' + replace(vexFiles[0], '/', ',') + '.workdir'

if isfile(tmpDir) or isdir(tmpDir):
	print 'Error: cannot proceed as temporary directory %s already exists.  Delete and try again.' % tmpDir
	exit(0)

execute('mkdir -p %s' % tmpDir)

chdir(tmpDir)

overlaps = []

# remove saved obser.tx files older than observTxHistoryDays
deleteOld(observTxHistoryDir, observTxHistoryDays)

for vexFile in vexFiles:
	if vexFile[0] != '/':
		vexFile = '%s/%s' % (cwd, vexFile)

	obsCode, obsSeg, stationTimes = vexPeek(vexFile)

	stns = stationTimes.keys()
	stns.sort()
	print stns
	if stns == ['MK', 'PT'] and not usnoConvert:
		print 'Warning: This looks like a USNO experiment because it has only MK and PT, but the --usno flag was not provided.'
		print 'Press enter to continue or Ctrl-C to stop this process'
		stdin.readline()

	v = execute('%s/vex2script %s %s' % (difxPath, vexFile, options))
	if v != 0:
		print 'An error occurred running vex2script on file %s.  Stopping.' % vexFile
		exit(0)

	scripts = glob('*.??.py')
	scripts.sort()

	if file != '':
		for script in scripts:
			execute("mv %s %s.tmp" % (script, script))
			execute("""sed -e "s/dbe0[ ]*=[ ]*RDBE(0,[ ]*'pfb')/dbe0 = RDBE(0, 'pfb', '%s')/" -e "s/dbe0[ ]*=[ ]*RDBE(0,[ ]*'ddc')/dbe0 = RDBE(0, 'ddc', '%s')/" <%s.tmp >%s""" % (file, file, script, script))
			execute("rm %s.tmp" % script)

	if edit:
		print '\nNow you can edit the .py files in %s before they are sent to the stations.' % tmpDir
		print '\nPress enter to continue'
		stdin.readline()

	for script in scripts:
		s = split(script, '.')
		station = lower(s[-2])
		if len(antennas) > 0 and not station in antennas:
			print '** Skipping antenna %s from vex file %s' % (upper(station), vexFile)
			continue
		expt = upper(s[0])
		print '\n==========================================================================='
		if station == 'gb':
			cc = 'gb-cc.gb.nrao.edu'

	# GB not yet handled per the request of GB staff

			#print 'Warning: files for GB are not yet handled by this script'
			#print '\n==========================================================================='
			continue
		else:
			cc = '%s-cc.vlba.nrao.edu' % station
		
		upperStation = upper(station)
		timeRange = stationTimes[upperStation]
		startDate = mjd2datetime(timeRange[0])
		stopDate = mjd2datetime(timeRange[1])
		print 'STATION %s : observing %s from %s to %s' % (upperStation, expt, startDate, stopDate)
		print '===========================================================================\n'

		if usnoConvert and (upperStation == 'MK' or upperStation == 'PT'):
			print 'USNO Experiment %s' % (expt)
			print 'Modifying script file for %s to use USNO equipment.' % upperStation
			execute("mv %s %s.tmp" % (script, script))
			execute("""sed -e "s/recorder0[ ]*=[ ]*Mark5C('-1')/recorder0 = Mark5C('usno')/" -e "s/dbe0[ ]*=[ ]*RDBE(0,[ ]*'pfb'/dbe0 = RDBE(2, 'pfb'/" -e "s/subarray.set4x4Switch('1/subarray.set4x4Switch('2/" <%s.tmp >%s""" % (script, script))
			execute("rm %s.tmp" % script)
			print 'Done.\n'

		print 'Sending %s to %s' % (script, cc)
		execute('scp vlbamon@%s:%s/%s . 2>&1' % (cc, ccScriptDir, observTxFile), 1)
		if not isfile(observTxFile):
			execute('touch %s' % observTxFile)

		saveObservTxFile(observTxFile, station)

		observations = open(observTxFile).readlines()
		observations = addObservation(observations, expt, timeRange[0], timeRange[1], station, overlaps)
		f = open(observTxFile, 'w')
		for o in observations:
			f.write(o)
		f.close()

		print 'Contents of %s at %s (only last %d lines shown):' % (observTxFile, upperStation, linesToShow)
		system('tail -n %d %s' % (linesToShow, observTxFile))
		print ''
		if nocopy:
			print '*** Not copying %s and %s to vlbamon@%s:%s 2>&1' % (script, observTxFile, cc, ccScriptDir)
		else:
			execute('scp %s %s vlbamon@%s:%s 2>&1' % (script, observTxFile, cc, ccScriptDir), 1)
		execute('rm %s' % observTxFile)

chdir(cwd)

execute('cp %s/*.py .' % tmpDir)
execute('rm -rf %s' % tmpDir)

if len(overlaps) > 0:
	print '\n*** WARNING ***\n'
	print 'There were scan overlaps at %d stations:' % len(overlaps)
	for o in overlaps:
		print '  %s' % o
	print ''