#!/usr/bin/env python

#**************************************************************************
#   Copyright (C) 2008-2013 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$
#
#============================================================================

from sys import argv, stdout, exit
from string import split, strip, find
import gzip

program = 'vlog'
version = 1.5
verdate = '20130122'
author  = 'Walter Brisken'

def usage():
	print '\n%s ver. %s   %s %s\n' % (program, version, verdate, author)
	print 'A program to preprocess the cal files to simiplfy difx2fits.\n'
	print 'Usage: %s [options] <TSM file> [<antenna list>]\n' % (argv[0])
	print '  options can include:\n'
	print '    --help'
	print '    -h         print this help info and quit\n'
	print '    --mjd'
	print '    -m         use mjd rather than doy timestamps\n'
	print '  <TSM file> is <project>cal.vlba or <project>cal.vlba.gz\n'
	print '  <antenna list> is a comma separated list of antennas with'
	print '          no spaces, e.g., FD,GB,Y.  Default -- all ants.\n'
	exit(0)


def doy2mjd(year, doy):
	NYR = int(year)
	NYRD = int(doy)
	# L2 IS -1 IF LEAP YEAR, -2 IF NOT. 
	L2 = NYR/4-(NYR+7)/4-NYR/100+(NYR+99)/100+NYR/400-(NYR+399)/400
	NYR1 = NYR-1
	MJD = doy-678576+365*NYR1+NYR1/4-NYR1/100+NYR1/400
	return MJD

def calctime(t, startyear, startdoy, useMJD):
	if not useMJD or t > 50000.0:
		return t
	else:
		if int(t) < startdoy:
			return doy2mjd(startyear+1, t)
		else:
			return doy2mjd(startyear, t)


def getlist(sections, header):
	s = split(header)
	sec = s[2]
	ant = s[-2]
	if not sections.has_key(sec):
		sections[sec] = {}
	S = sections[sec]
	if not S.has_key(ant):
		S[ant] = []
	return S[ant]

def chop(filename):
	sections = {}
	current = []
	year = 0
	doy = 0
	if filename[-3:] == '.gz':
		lines = gzip.open(filename, 'r').readlines()
	else:
		lines = open(filename, 'r').readlines()
	for D in lines:
		d = strip(D)
		if year == 0 and find(d, '! For UT timerange:') == 0:
			startday = split(d)[4]
			year = int(startday[:4])
			doy = int(split(startday, '/')[-1])

		if find(d, ' -----') > 0:
			current = getlist(sections, d)
		elif len(d) > 0:
			current.append(d)
	return sections, year, doy

def parsetime(str):
	s = split(str, '-')
	t = split(s[1], ':')
	return int(s[0]) + float(t[0])/24.0 + float(t[1])/1440.0 + \
		float(t[2])/86400.0

def makeWR(sections, antenna, startyear, startdoy, useMJD):
	if not sections.has_key('Weather'):
		return None
	if not sections['Weather'].has_key(antenna):
		return None
	WeatherData = sections['Weather'][antenna]
	WR = []
	if useMJD:
		dateFormat = 'MJD'
	else:
		dateFormat = 'D.O.Y.'
	WR.append('# Ant %s temp(C) pressure(mbar) dewpoint(C) windspeed(m/s) winddir(degrees) rain(cm) gust(m/s)' % dateFormat)
	for w in WeatherData:
		if w[0] in ['*', '!']:
			continue
		s = split(w)
		if len(s) < 8:
			continue
		t = parsetime(s[0])
		WR.append('%9.7f %s %s %s %s %s %s %s' % (calctime(t, startyear, startdoy, useMJD), s[1], s[2], s[3], s[4], s[5], s[6], s[7]))
	return WR

def bbcflag(flags, bbc, t1, t2, chan, reason, startyear, startdoy, useMJD):
	for b in bbc:
		T1 = max(t1, b[0])
		T2 = min(t2, b[1])
		nBBC = len(b) - 2
		if T1 < T2:
			for i in range(nBBC):
				if b[i+2][0][0:-1] == chan:
					flags.append('%9.7f %9.7f %d %s' % (calctime(T1, startyear, startdoy, useMJD), calctime(T2, startyear, startdoy, useMJD), i, reason))

def makeFG(sections, BBCs, antenna, startyear, startdoy, useMJD):
	if not sections.has_key('Edit'):
		return None
	if not sections['Edit'].has_key(antenna):
		return None
	bbc = BBCs[antenna]
	FlagData = sections['Edit'][antenna]
	FG = []
	if useMJD:
		dateFormat = 'MJD'
	else:
		dateFormat = 'D.O.Y.'
	FG.append('# Ant %s.start %s.end RecChan reason' % (dateFormat, dateFormat))
	key = "ant_name='%s'" % antenna
	for f in FlagData:
		s = split(f)
		if s[0] == key:
			t = split(split(s[1], '=')[1], ',')
			t1 = parsetime('%s-%s:%s:%s' % (t[0], t[1], t[2], t[3]))
			t2 = parsetime('%s-%s:%s:%s' % (t[4], t[5], t[6], t[7]))
			p = find(f, 'reason=')
			p1 = p + 9
			p2 = p1 + find(f[p1:], "'")
			reason = f[p1-2:p2+1]
			r = split(reason[1:-1])
			if r[0] == 'channel' and r[2] == 'bbc':
				bbcflag(FG, bbc, t1, t2, r[1], reason, startyear, startdoy, useMJD)
			else:
				FG.append('%9.7f %9.7f -1 %s' % (calctime(t1, startyear, startdoy, useMJD), calctime(t2, startyear, startdoy, useMJD), reason) )
	return FG


def countTones(ToneTable, StateTable):
	polList  = []
	bandnums = {}
	polCount = {}
	nTone  = 0
	nBand  = 0
	nPol   = 0
	nState = 0
	# [chan, sb, pol, freq; polnum, bandnum, tonenum]
	toneKeys = ToneTable.keys()
	toneKeys.sort()
	# [chan, sb, pol, nState; polnum, bandnum]
	stateKeys = StateTable.keys()
	stateKeys.sort()

	# first populate polarization list so we can force proper order
	for key in toneKeys:
		pol = ToneTable[key][2]
		if not pol in polList:
			polList.append(pol)
			nPol += 1
	for key in stateKeys:
		pol = StateTable[key][2]
		if not pol in polList:
			polList.append(pol)
			nPol += 1
	
	# if both RCP and LCP, make sure RCP is first
	polList.sort()
	polList.reverse()

	for key in toneKeys:
		T = ToneTable[key]
		pol = T[2]
		chanpol = T[0]+T[1]+T[2]
		chan = T[0]
		polnum = polList.index(pol)
		if chanpol in bandnums:
			bandnums[chanpol][1] += 1
			if bandnums[chanpol][1] >= nTone:
				nTone = bandnums[chanpol][1] + 1
		else:
			if pol in polCount:
				polCount[pol] += 1
				if polCount[pol] >= nBand:
					nBand = polCount[pol] + 1
			else:
				polCount[pol] = 0
			if nBand == 0:
				nBand = 1
			bandnums[chanpol] = [polCount[pol], 0]
			if nTone == 0:
				nTone = 1

	for key in stateKeys:
		if not key in ToneTable:
			S = StateTable[key]
			pol = S[2]
			chanpol = S[0]+S[1]+S[2]
			polnum = polList.index(pol)
			if chanpol in bandnums:
				bandnums[chanpol][1] += 1
			else:
				if pol in polCount:
					polCount[pol] += 1
					if polCount[pol] >= nBand:
						nBand = polCount[pol] + 1
				else:
					polCount[pol] = 0
				if nBand == 0:
					nBand = 1
				bandnums[chanpol] = [polCount[pol], 0]

	toneIndex = {}

	for key in toneKeys:
		T = ToneTable[key]
		pol = T[2]
		chanpol = T[0]+T[1]+T[2]
		T.append(polList.index(pol))	# polnum
		T.append(bandnums[chanpol][0])	# bandnum
		if chanpol in toneIndex:
			toneIndex[chanpol] += 1
		else:
			toneIndex[chanpol] = 0
		T.append(toneIndex[chanpol])	# tonenum

	
	for key in stateKeys:
		S = StateTable[key]
		if S[3] > nState:
			nState = S[3]
		pol = S[2]
		chanpol = S[0]+S[1]+S[2]
		S.append(polList.index(pol))	# polnum
		S.append(bandnums[chanpol][0])	# bandnum

	return nPol, nBand, nTone, nState

def bbc2recChan(bbcNum, t, bbcMap):
	for b in bbcMap:
		if t >= b[0] and t < b[1]:
			nBBC = len(b) - 2
			for i in range(nBBC):
				if b[i+2][1] == bbcNum:
					return ('%d' % i), nBBC
	return '-1', 0
	

def makePH(sections, BBCs, antenna, startyear, startdoy, useMJD):
	if not sections.has_key('PulseCal'):
		return None
	if not sections['PulseCal'].has_key(antenna):
		return None
	PulseCalData = sections['PulseCal'][antenna]
	bbc = BBCs[antenna]

	PH = []
	if useMJD:
		dateFormat = 'MJD'
	else:
		dateFormat = 'D.O.Y.'
	PH.append('# Ant %s dur(days) C.Cal nPol nBand nTone nState nRecChan (RecChan, freq, mag, phase(degrees))[nPol][nBand][nTone] (RecChan, stateCount[nState])[nPol][nBand];  RecChan is recorder channel num (0-based)' % dateFormat)

	action = 0	# 0 == nothing
			# 1 = ToneTable
			# 2 = pulsecal data
			# 3 = write row

			# [bandnum, polnum, tonenum] are zero-based
	ToneTable = {}	# [chan, sb, pol, freq; polnum, bandnum, tonenum]
	StateTable = {}	# [chan, sb, pol, nState; polnum, bandnum]
	ToneData = []
	StateData = []
	nBBC = 0
	lastKey = None

	for row in PulseCalData:
		if row[0] == '!':
			continue
		s = split(row)
		if len(s) == 0:
			continue

		if s[0] == '/' and action == 1:	# process ToneTable
			nPol, nBand, nTone, nState = \
				countTones(ToneTable, StateTable)
		elif s[0] == 'FREQUENCY':  # start (re)population of ToneTable
			action = 1
			ToneTable = {}
			StateTable = {}
		elif s[0] == 'PULSE-CAL':  # start data row(s)
			action = 2
		elif action == 1 and len(s) == 8:
			key = s[0][1:-1]
			if s[2] == "'U'" or s[2] == "'L'":
				if float(s[4]) < 1.0:
					StateTable[key] = [ \
						s[1], s[2][1:-1], \
						s[3][1:-1], 1 << int(s[5])]
				else:
					ToneTable[key] = [ \
						s[1], s[2][1:-1], \
						s[3][1:-1], s[4]]
				lastKey = key
		elif action == 2 and len(s) >= 4:
			key = s[2][1:-1]
			if key == lastKey:
				action = 3
			if key == 'CC':
				ToneData = []
				for i in range(nPol*nBand*nTone):
					initial = ['0']*4
					initial[0] = -1
					ToneData.append(initial)
				StateData = []
				for i in range(nPol*nBand):
					zeros = ['0']*nState
					StateData.append(['-1', zeros])
				t = parsetime(s[0]+'-'+s[1])
				ccal = s[3]
				if ccal == '-Inf' or ccal == 'Inf':
					ccal = '999.9'
				dt = float(s[4])/86400.0
			elif key in ToneTable:
				T = ToneTable[key]
				pol  = T[4]
				band = T[5]
				tone = T[6]
				index = tone + nTone*(band + nBand*pol)
				recChan, nBBC = bbc2recChan(int(T[0]), t, bbc)
				try:
					ToneData[index] = [recChan, T[3], s[3], s[4]]
				except IndexError:
					print 'ToneData index error: T=', T, 's=', s
					print 'index=', index, 'array size=', len(ToneData)
					print 'Full line from file: ', row
					print 'nPol=%d nBand=%d nTone=%d' % (nPol, nBand, nTone)
					print 'ToneTable=',ToneTable
					print 'StateTable=',StateTable,'\n'
			elif key in StateTable:
				S = StateTable[key]
				pol  = S[4]
				band = S[5]
				for state in range(S[3]):
					index = band + nBand*pol
					try:
						StateData[index][0], nBBC = \
							bbc2recChan(int(S[0]), t, bbc)
						StateData[index][1][state] = \
							s[3+state]
					except IndexError:
						print 'StateData index error: S=', S, 's=', s
						print 'index=', index, 'array size=', len(StateData)
						print 'Full line from file: ', row
						print 'nPol=%d nBand=%d nTone=%d\n' % (nPol, nBand, nTone)
		if action == 3:
			line = '%10.7f %9.7f %s %d %d %d %d %d' % \
				(calctime(t, startyear, startdoy, useMJD), dt, ccal, nPol, nBand, nTone, nState, nBBC)
			for T in ToneData:
				line += ('  %s %s %s %s' % (T[0], T[1], T[2], T[3]))
			for S in StateData:
				if len(S[1]) > 0:
					line += ('  %s' % S[0])
					for s in S[1]:
						line += (' %s' % s)
			PH.append(line)
			action = 2
	return PH

def makeTS(sections, antenna, startyear, startdoy, useMJD):
	if not sections.has_key('Tsys'):
		return None, None
	if not sections['Tsys'].has_key(antenna):
		return None, None
	TsysData = sections['Tsys'][antenna]

	TS = []
	if useMJD:
		dateFormat = 'MJD'
	else:
		dateFormat = 'D.O.Y.'
	TS.append('# ant %s dur(days) nRecChan (tsys, bandName)[nRecChan]' % dateFormat);

	recvName = []
	nRecChan = 0
	timerange = (0.0, 0.0)
	BBCmap = []
	BBCrow = [0, 0]
	t0 = 0
	t1 = 0

	for row in TsysData:
		s = split(row)
		# check for a comment line with vital info
		if s[0] == '!' and len(s) == 11 and \
				s[3] in ['A', 'B', 'C', 'D'] and \
				s[6] in ['U', 'L'] and s[8][-1] in ['M', 'K']:
			# calculate net side band, rather than use the supplied
			# BBC sideband:
			if s[9][0] == '-':
				s[6] = 'L'
			else:
				s[6] = 'U'
			if s[1] == '1':	# reset on channel 1
				recvName = []
				nRecChan = 0
				if t1 < t0-350 and t0 >= 365:
					t1 += int(t0)	# handle leap day
				BBCrow[1] = t0
				BBCrow = [t0, t1]
				BBCmap.append(BBCrow)
			recvName.append(s[2])
			BBCrow.append([s[5]+s[6], int(s[1]), s[4][0]])
			nRecChan += 1
		elif s[0] == '!' and len(s) == 5:
			tr = split(s[4], '/')
			if len(tr) == 2 and s[1] == antenna:
				t0 = parsetime(tr[0])
				t1 = parsetime(tr[1])
				BBCrow[1] = t1
				if BBCrow[1] < BBCrow[0]-350 and BBCrow[0] >= 365:
					BBCrow[1] += int(BBCrow[0])
		elif row[0] == 'TSYS':
			pass
		elif row[0] != '!' and len(s) == 4+nRecChan:
			t = parsetime(s[0]+'-'+s[1]+':0')
			dt = 0.0
			line = "%10.7f %9.7f %d" % (calctime(t, startyear, startdoy, useMJD), dt, nRecChan)
			for i in range(nRecChan):
				line += ('  %s %s' % (s[2+i], recvName[i]))
			TS.append(line)

	return TS, BBCmap

def writetable(table, filename, ants, tablename):
	f = open(filename, 'w')
	stdout.write('%s:' % tablename)
	for a in ants:
		if table.has_key(a):
			t = table[a]
			stdout.write(' %s' % a)
			for row in t:
				if row[0] == '#':
					f.write('%s\n' % row)
				else:
					f.write('%s %s\n' % (a, row))
	f.close()
	stdout.write('\n')

def findantennas(sections):
	ants = []
	for s in sections:
		for a in sections[s].keys():
			if not a in ants:
				ants.append(a)
	ants.sort()
	return ants

# -- main() -----------------

sections = {}
ants = []
useMJD = False
antArg = None

for a in argv[1:]:
	if a[0] == '-':
		if a in ['-h', '--help']:
			usage()
		elif a in ['-m', '--mjd']:
			useMJD = True
		else:
			print 'unrecongized option %s' % a
			exit(0)
	elif len(sections) == 0:
		sections, startyear, startdoy = chop(a)
	elif antArg == None:
		antArg = a
	else:
		usage()

if len(sections) == 0:
	usage()

if antArg != None:
	ants = split(antArg, ',')
else:
	ants = findantennas(sections)

if len(ants) == 0:
	usage()

print 'It looks like the observation started on day %d of year %d' % (startdoy, startyear)

print 'Processing antennas:', ants

BBCs = {}
TSs  = {}
PHs  = {}
WRs  = {}
FGs  = {}
for a in ants:
	ts, bbc = makeTS(sections, a, startyear, startdoy, useMJD)
	if ts != None:
		BBCs[a] = bbc
		TSs[a] = ts
	wx = makeWR(sections, a, startyear, startdoy, useMJD)
	if wx != None:
		WRs[a] = wx

for a in ants:
	if BBCs.has_key(a):
		fg = makeFG(sections, BBCs, a, startyear, startdoy, useMJD)
		if fg == None:
			continue
		FGs[a] = fg
		ph = makePH(sections, BBCs, a, startyear, startdoy, useMJD)
		if ph != None:
			PHs[a] = ph

writetable(FGs, 'flag',    ants, 'FG')
writetable(TSs, 'tsys',    ants, 'TS')
writetable(WRs, 'weather', ants, 'WR')
writetable(PHs, 'pcal',    ants, 'PH')