#!/usr/bin/env python

from sys import argv, exit
from string import split, strip, upper, lower
from glob import glob
from os import getenv, environ, chdir, system, getcwd, rmdir, umask, listdir, remove, popen
from os.path import isfile, isdir, getsize
from time import gmtime, time

#
program = 'makemark4'
version = '0.2'
author  = 'David Gordon'
verdate = '20140105'

def setAntennaCodes(vexfile, passname, verbose):
	# Adapted from single_code.c in difx2mark4
	stationCodeTable = [ \
		["K", "Kk"], \
		["b", "Br"], \
		["f", "Fd"], \
		["h", "Hn"], \
		["k", "Kp"], \
		["l", "La"], \
		["m", "Mk"], \
		["n", "Nl"], \
		["o", "Ov"], \
		["p", "Pt"], \
		["s", "Sc"], \
		["G", "Gb"], \
		["c", "Cm"], \
		["H", "Hb"], \
		["H", "Ho"], \
		["T", "Ts"], \
		["U", "Ur"], \
		["S", "Sv"], \
		["S", "Nt"], \
		["I", "Ma"], \
		["P", "Oh"], \
		["X", "On"], \
		["V", "Wz"], \
		["N", "Ny"], \
		["B", "Bd"], \
		["B", "Mc"], \
		["M", "Mc"], \
		["y", "Y" ], \
		["r", "Ro"], \
		["w", "Wb"], \
		["Y", "Yg"], \
		["Y", "Yb"], \
		["Y", "Ys"], \
		["y", "Yg"], \
		["y", "Yb"], \
		["y", "Ys"], \
		["F", "Ft"], \
		["f", "Ft"], \
		["J", "Ht"], \
		["A", "Ke"], \
		["O", "Tc"], \
		["C", "Ts"], \
		["Z", "Zc"], \
		["r", "Zc"], \
		["w", "Wf"], \
		["j", "Hh"], \
		["j", "HH"], \
		["s", "Sh"], \
		["S", "Sh"] \
	]

	options='ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'
	start = 0

	antennadata = popen('vexpeek %s' % vexfile).readlines()
	if(len(antennadata) < 2):
		return
	antennas = []
	for d in antennadata[1:]:
		ant = split(d)[0]
		if len(ant) == 1:
			antennas.append(upper(ant[0]))
		else:
			antennas.append(upper(ant[0]) + lower(ant[1]))

	outfile = '%s/%s.stationcodes' % (getcwd(), passname)
	out = open(outfile, 'w')
	done = []
	used = []

	for k, s in stationCodeTable:
		if s in antennas:
			done.append(s)
			used.append(k)
			out.write('%s xx\n' % k)	# W__ is this needed?
			out.write('%s %s\n' % (k, s))
	for a in antennas:
		if a not in done:
			while options[start] in used:
				start += 1
			used.append(options[start])
			done.append(a)
			out.write('%s xx\n' % options[start])
			out.write('%s %s\n' % (options[start], a))
			
	out.close()
	environ['HOPS_STATION_CODE'] = outfile


def genDateStr(t):
	d = gmtime(t)
	return '%02d%02d%02dT%02d%02d%02d' %  (d.tm_year % 100, d.tm_mon, d.tm_mday, d.tm_hour, d.tm_min, d.tm_sec)


def usage(prog):
	print '\n%s ver %s  %s  %s' % (program, version, author, verdate)
	print '\nUsage:  %s [options] <exp name>' % prog
	print '\n <exp name> is the lower case experiment name'
	print '\nOptions can include:'
	print '  --verbose'
	print '  -v         Send more output to the screen\n'
	print '  --help'
	print '  -h         Print this help information and quit\n'
	exit(0)

def parsekv(str):
	kv = {}
	ss = split(str)
	for s in ss:
		p = split(s, '=')
		if len(p) != 2:
			print 'Error parsing key=value statement: %s', s
			return {}
		kv[p[0]] = p[1]
#		kv = p[1]
	return kv

def getmjd(t):
	return 40587.0 + t/86400.0

def splitobscode(exper):
	obsSeg = ''
	proposal = exper[:]

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

	return proposal, obsSeg

class JobList:
	def __init__(self, filename,verbose):
		self.filename = filename
		data = open(filename).readlines()
		self.jobs = []
		self.mjds = []
		mstart = []
		mstop  = []
		self.npsrbins = 0
		self.nsrcfiles = 1
		if len(data) < 2:
			print 'Malformed .joblist file %s' % filename
			exit(0)

		self.kw = parsekv(data[0])

		if len(self.kw) < 1:
			print 'Malformed .joblist file %s line 1' % filename
			exit(0)
#  Get job name and start/stop MJD's
		for n in range(1, len(data)):
			d = split(data[n], ' ')[0]
			self.jobs.append(d)
			d1 = split(data[n], ' ')[1]
			d2 = split(data[n], ' ')[2]
			mstart.append(d1)
			mstop.append(d2)
#  Find project start and stop times
		fstart = [float(s) for s in mstart]
		self.xstart = min(fstart)
		fstop  = [float(x) for x in mstop ]
		self.xstop  = max(fstop )
#		print ' min = %f  max = %f ' % (self.xstart, self.xstop)



	def makemark4(self, queuedir, exp, joblistfile, verbose):

		jobdir = getcwd()  # save current jobs directory

#  Check that the mark4 listfile doesn't already exist.
		mark4listfile = '%s/%s.mark4list' % (jobdir, self.kw['pass'])
		if isfile(mark4listfile):
			print "Mark4 listing file " + mark4listfile + " already exists; not proceeding!"
			print "Remove this file if the mark4 files listed within have already been"
			print "archived or are not required"
			exit(0)
		if verbose:
			print '\n mark4 list file will be: %s ' % mark4listfile

#
		projdir = queuedir + '/' + self.kw['exper']

		if not isdir(projdir):	# check for directory existence
			print '\n difx queue directory %s does not exist ' % projdir
			print ' Enter new directory or Return to terminate '
			projdir = ''
			projdir = raw_input()
			if len(projdir) < 1: 
				exit(0)
		if not isdir(projdir):
			print '\n difx queue directory %s does not exist, terminating ' % projdir
			exit(0)
		print '\n Moving to project queue directory %s ' % projdir
		chdir(projdir)	# move to difx files directory

# Check for existence of the 1234 directory. It should not already exist.
#		projdir = getcwd()
		mk4dir = projdir + '/1234'
		if isdir(mk4dir):
			ld = listdir(mk4dir)
			if len(ld) > 0:
				print '\nError: mark4 directory %s already exists and is not empty. ' % mk4dir
				print ' Either remove or rename it and try again %s.' % mk4dir
				print ' '
				exit(0)
		if verbose:
			print '\n Directory for mark4 output will be %s ' % mk4dir

# Check that the metadata file doesn't already exist.
#  If you change the name of the metadata file here, you also
#  need to change it in difxarch.
		metadatafile = projdir + '/' + self.kw['exper'] + '.mark4.metadata.txt'
		if isfile(metadatafile):
			print "Mark4 metadata file " + metadatafile + " already exists; not proceeding!"
			print "Remove this file if the mark4 files have already been"
			print "archived or are not required"
			exit(0)
		if verbose:
			print ' metadata file will be: %s ' % metadatafile

#		Find observing band(s)
		bandcodestring = []
		if self.kw.has_key('vex'):
			vexfile = self.kw['vex']
		else:
			vexfile = projdir + '/' + self.kw['pass'] + '.vex.obs'
		if not isfile(vexfile):
			vexfile = projdir + '/' + self.kw['pass'] + '.vex.preobs'
		if not isfile(vexfile):
			vexfile = projdir + '/' + self.kw['pass'] + '.vex'
		if isfile(vexfile):
			bandcodestring = strip(popen('vexpeek %s --bands 2> /dev/null' % vexfile).readline())
			if verbose:
				print '\n Observing band(s): %s ' % bandcodestring
			setAntennaCodes(vexfile, self.kw['pass'], verbose)
		else:
			bandcodestring = []


		njobs = 0
		nwarn = 0
		warnlist = []

		print ' '
		if not verbose:
			print ' Processing jobs:',
		nm1 = len(self.jobs) - 1
		for n in range(0, len(self.jobs)):
			msgfile = "msg." + self.jobs[n]
			cmd = 'difx2mark4 -v --override-version -e 1234 ' + self.jobs[n] + ' &> ' + msgfile
			if verbose:
				print cmd
			else:
				print self.jobs[n],
			system(cmd)	# run difx2mark4 and output to the 1234 directory
			njobs += 1
#
#			clgr = "grep \'Warning -- not all input files converted!\' "
			clgr = "grep \'0 of 1 scans converted!\' " + msgfile + " > warning.msg"
#			print clgr
			system(clgr)	# grep for failed difx2mark4 conversions

			f3 = open("warning.msg","r")
			f3line = f3.readline()
			if f3line:
				nwarn += 1
				warnlist.append(self.jobs[n])	# add to list of failed conversions
#			else:
#				remove(msgfile)	# delete message file if successful
			f3.close

		print '\n\n difx2mark4 finished, %d jobs processed. %d errors.' % (njobs, nwarn)
#		printout list of jobs with warning messages
		if len(warnlist) > 0:
			print "\n Errors on the following jobs: %s " % warnlist

		clp = "cp *vex* 1234/. "
		if verbose:
			print '\n Copying vex.obs file to %s ' % mk4dir
		system(clp)	# copy the vex.obs file to the 1234 directory

		t = time()
		datestr = genDateStr(t)
		mk4filename = 'VLBA_' + self.kw['exper'] + '_' + self.kw['pass'] + '_BIN0_SRC0_0_' +  datestr + '.mark4.tar'
		cmd = 'tar -cvf ' + mk4filename + ' 1234/* > mark4_tar_listing'
		print '\n Tar\'ing mark4 files to file %s ' % mk4filename
		if verbose:
			print '  tar listing in file \'mark4_tar_listing\' '
		system(cmd)

# make sure anyone in the difx group can remove the file later
		groupId = getenv('DIFX_GROUP_ID')
#		print "groupId %s " % groupId
		cmd = 'chown -R :%s %s' % (groupId, mk4filename)
		print '\n Changing permissions on file %s ' % mk4filename
		if verbose:
			print cmd
		system(cmd)

#   Compress mark4 tar file
		cmd = 'gzip ' + mk4filename 
		print '\n Compressing mark4 file %s ' % mk4filename
		if verbose:
			print cmd
		system(cmd)
		mk4filenamegzip = mk4filename + '.gz'
		size = getsize(mk4filenamegzip)/1000000.0
		size1= getsize(mk4filenamegzip)/1000.0

# Write the metadata file
#		metadatafile = self.kw['exper'] + '.mark4.metadata'
#		print ' metadatafile = %s ' % metadatafile 

#		Find experiment segment
		segm = []
		segm = splitobscode(self.kw['exper'])
		if verbose:
			print '\n Experiment segment: %s ' % segm[1]

		meta = open(metadatafile, 'w')

		meta.write('PROJECT_CODE     = %s\n' % (self.kw['exper']))
		meta.write('SEGMENT          = %s\n' % segm[1])
		meta.write('STARTTIME        = %13.7f\n' % self.xstart)
		meta.write('STOPTIME         = %13.7f\n' % self.xstop)
		meta.write('TELESCOPE        = VLBA\n') 
		meta.write('ARCH_FORMAT      = mark4\n') 
		meta.write('DATA_TYPE        = raw\n') 
		meta.write('ARCH_FILE        = %s\n' % mk4filenamegzip )
		meta.write('FILE_SIZE        = %f\n' % size1)
		meta.write('RAW_PROJECT_CODE = %s\n' % self.kw['exper'])
		meta.write('OBS_BANDS        = %s\n' % bandcodestring) 

		meta.close() 

		print '\n Writing metadata file: %s ' % metadatafile
		print ' But it may need hand editing. Check before archiving.'
		if len(bandcodestring) < 1:
			print '\n !! Observing band(s) were not found. !! \n  '

#  Return to original jobs directory and write the *.mark4list file
		if verbose:
			print "\n Returning to jobs directory: %s to make the mark4 list file and metadata file " % jobdir
		chdir(jobdir) 

		out = open(mark4listfile, 'w')
		print '\n Writing mark4 list file: %s ' % mark4listfile

		mjd = getmjd(t)
		out.write('exper=%s  pass=%s  jobs=%s  mjd=%9.7f  DiFX=%s  difx2mark4=0' % \
			(self.kw['exper'], self.kw['pass'], joblistfile, getmjd(t), self.kw['DiFX']))
		if 'label' in self.kw:
			out.write('  label=%s' % self.kw['label'])
		out.write('\n')

#		nDigit = countDigits(indices[-1][0])
#		nBinDigits = countDigits(binrange)
#		nSrcDigits = countDigits(self.nsrcfiles)

#		i[0] = 0
#		i[1] = 0
#		i[2] = 0
#		fn = '%s.%d.bin%04d.source%04d.mark4' % (self.kw['exper'], i[0], i[1], i[2])

		fn = self.kw['exper'] + '.0.bin0000.source0000.mark4'
		out.write('%s %4.2f %s\n' % (mk4filenamegzip, size, fn))

		out.close()


		print '\n Writing metadata file: %s ' % metadatafile
		print ' But it may need hand editing. Check before archiving.'
		if len(bandcodestring) < 1:
			print '\n !! Observing band(s) were not found. !! \n  '


		if len(warnlist) > 0:
			print '\n makemark4 finished but there were %d difx2mark4 errors!!! \n' % nwarn
		else:
			print '\n makemark4 finished. No difx2mark4 errors. \n'



# main starts here

umask(02)


verbose = False
joblistfile =''

args = []
for a in argv[1:]:
	if a[0] == '-':
		if a in ('-h', '--help'):
			usage(argv[0])
		if a in ('-v', '--verbose'):
			verbose = True
	elif  isfile(a + '.joblist'):
			joblistfile = a + '.joblist'
if len(joblistfile) < 1:
	usage(argv[0])


print "\n Program makemark4. "
print "\n Converts difx correlator output to Mark4 format."
print " Should be run after makefits when requested. "
# print " Must be run in the jobs directory. "

if verbose:
	print "\n joblist file: %s " % joblistfile

# Make sure joblist file exists in current directory
if not isfile(joblistfile):
	print '\n Cannot find joblist file: %s ' % joblistfile
	print 'Quitting! '
	exit(0)

difxQueueBase = getenv('DIFX_QUEUE_BASE')
if difxQueueBase == None:
	print 'Error: env var DIFX_QUEUE_BASE is not set.  Cannot proceed.'
	exit(0)
if verbose:
	print '\n difxQueueBase: %s ' % difxQueueBase

jl = JobList(joblistfile,verbose)

OK = jl.makemark4(difxQueueBase, a, joblistfile, verbose)


print ""
print ""