# Automated loading and housekeeping of WSRT data.
# Original version: mkargo 20120102
# Based on eMERLIN_load.py


####################################################################################
# Put the data directories you want somewhere logical and copy this script there.
# The script searches recursively through directories it finds and loads 
# everything with a ".UVF" extention, so delete/move anything you do not wish to load.
#
# Once the script has run, your data will be loaded and processed ready for flagging.
#
# If the script fails, the first thing to do is run 'source /aips/LOGIN.CSH'.
####################################################################################

# Procedure:
# load *UVF using fitld
# use PUTHEAD to set SORTORD to 'TB' in each (the data are really in this sort order but many tasks want to know this in the header)
# use PUTHEAD to set CRVAL2 to -1 in each (this redefines the polarization products which were measured)
# run INDXR with CPARM=0,0,1 on each (makes NX and basic CL tables) with one minute gain table entries (important for pol. calibration)
# run APCAL on each (makes SN table from the TY and GC tables which contain info on system temperature variations during the obs.)
# make a list of KLASS attributes and remove duplicates (if there are multiple configurations)
# use DBCON to concatenate (two at a time) all calibrator and source fields in time sequence (this only retains the SN table from the very first file in the sequence)
# use INDXR on the result
# use TACOP to copy all additional SN tables to the combined file
# use PRTAB to check that all SN tables have a consistent day number
# use TABED if necessary to edit those SN tables which are completely on a new UT day, by adding one to the UT day number
# use CLCAL to smooth, merge and apply the SN tables for Tsys calibration, using for example SAMPTYP 'MWF' and BPARM 0.2 to smooth over the odd glitch, producing CL.2 

# Import useful things

import os, re, time, datetime, sys, fnmatch
from os.path import join, getsize
from datetime import date
#import Utilities
from collections import deque
from AIPS import AIPS
from AIPSTask import AIPSTask, AIPSList
from AIPSData import AIPSUVData, AIPSImage, AIPSCat
from Wizardry.AIPSData import AIPSUVData as WizAIPSUVData


print "------------------------------------------------------------"
print "  Load WSRT data from a local disk and do some houekeeping"
print "------------------------------------------------------------"
print "Your data should be in subdirectories within your current" 
print "working directory. If this is not the case, quit now."
print "Your data will be FITLDd, MSORTd, INDXRd and DBCONd."



# Get AIPS user number and disk info
print "Enter your AIPS number: ",
usernum = raw_input()
if usernum.isdigit() :
	usernum = int(usernum)
else :
	print "Does not compute."
	sys.exit()
print "Enter the AIPS disk to use: ",
indisk = raw_input()
if indisk.isdigit() :
	indisk = int(indisk)
else :
	print "Does not compute."
	sys.exit()
outdisk = indisk

print "Thank you, proceeding. We suggest that you have a nice hot cup of tea."


# Set up AIPS things

AIPS.userno = usernum
fitld = AIPSTask('FITLD')
rename = AIPSTask('RENAME')
zap = AIPSTask('ZAP')
msort = AIPSTask('MSORT')
msort.sort = 'TB'
indxr = AIPSTask('INDXR')
apcal = AIPSTask('APCAL')
tacop = AIPSTask('TACOP')
dbcon = AIPSTask('DBCON')
clcal = AIPSTask('CLCAL')
tabed = AIPSTask('TABED')

# AIPS is chatty. Supress most messages on the terminal but log everything.
AIPSTask.msgkill = -4
AIPS.log = open('WSRT_load.log', 'w')							

# Check we have an empty AIPS disk before we start
pca = AIPSCat(indisk)
if len(pca[indisk]) != 0 :
	print "You appear to already have file(s) on your AIPS disk, quitting."
	sys.exit()


# Create list of files to load, then do stuff with them

datadir = os.getcwd()					# search directories below where we are now
os.listdir(datadir)
print "Loading from:", datadir
mydirlist = [f for f in os.listdir(datadir)		# list sub-directories
	if os.path.isdir(os.path.join(datadir, f))]
mydirlist.sort()					# time order the list
thisdir = 0
nfiles = 0
filelist = []
for dir in mydirlist:					# for each sub-directory...
	thisdir = thisdir + 1
	numdirs = len(dir)
	mylist = [f for f in os.listdir(dir)		# make a list of files to load
	if os.path.isfile(os.path.join(dir, f))]
	mylist.sort()
	for fitsfile in mylist:				# for each file...
		if fnmatch.fnmatch(fitsfile, '*.UVF'):
			descriptor = "./" + dir + "/" + fitsfile
			print "Loading " + descriptor
			fitld.datain = descriptor
			outn = 'TMP'
			fitld.digicor = -1
			fitld.douvcomp = -1
			fitld.outdata = AIPSUVData(outn,'',outdisk,0)
			fitld.go()				# Load the data



pca = AIPSCat(indisk)
print "Your AIPS catalog looks like this:"
print AIPSCat(indisk)
print "#######################################"
print "Loading complete. Tidying up the files."
print "#######################################"

numfiles = len(pca[indisk])
i = 0
filelist = []

for fitsfil in pca[indisk]:
	# reset the filename based on the source name in the SU table
	datafile = AIPSUVData(fitsfil.name, fitsfil.klass, indisk, fitsfil.seq)
	sourcelist = datafile.table('SU',1)
	outn = datafile.sources[0]
	if len(outn)>12 :
		outn = outn[0:12]		# catch long source names
	inname = outn
	freq = datafile.header.crval[2]/1e9
	if freq > 2:
		band = 'C_BAND'
	else:
		if datafile.header.naxis[2] == 512:
			band = 'L_LINE'
		else:
			band = 'L_CONT'

	print "Renaming:", datafile, "as", outn,band
	datafile.rename(outn,band,0)
	i = i + 1


pca = AIPSCat(indisk)
print pca
i = 0
queueNAME = deque([])
queueKLAS = deque([])
queueSEQ  = deque([])

for fitsfil in pca[indisk]:
		datafile = AIPSUVData(fitsfil.name, fitsfil.klass, indisk, fitsfil.seq)
		print "Working on file:", datafile

		# set up lists for dbcon
		queueNAME.append(fitsfil.name)
		queueKLAS.append(fitsfil.klass)
		queueSEQ.append(fitsfil.seq)

		# PUTHEAD each file - not actually using puthead, have to convert to Wizardry format first
		uvwiz = WizAIPSUVData(datafile)
		uvwiz.header.sortord = 'TB'
		uvwiz.header.crval[1] = -1
		uvwiz.header.update()
		datafile = AIPSUVData(uvwiz)

		# msort each file, writing back to same file
		msort.indata = datafile
		msort.outdata = datafile
		msort.go()

		# index each file
		indxr.indata = datafile
		indxr.cparm[1:] = 0,0,1,0
		indxr.go()

		# run apcal to generate SN#1 for each file
		apcal.indata = datafile
		apcal.go()

		i = i + 1


print "All extensions", queueKLAS
extensions = []
extensions = set(queueKLAS)



print "################################################################################"
print "Proceeding with DBCON. Unique classes:", extensions
print "################################################################################"


for mode in extensions:
	tmpNAME = deque([])
	tmpKLAS = deque([])
	tmpSEQ  = deque([])
	i = 0
	j = 0
	for j in range(len(queueNAME)):
		if queueKLAS[j] == mode:
			tmpNAME.append(queueNAME[j])
			tmpKLAS.append(queueKLAS[j])
			tmpSEQ.append(queueSEQ[j])
	print "For", mode, "queue:", tmpNAME, tmpKLAS, tmpSEQ

	j = 1
	while len(tmpNAME)>1 :
		print "Combining " + tmpNAME[0] + '.' + tmpKLAS[0] + '.' + format(tmpSEQ[0]) + " and " + tmpNAME[1] + '.' + tmpKLAS[1] + '.' + format(tmpSEQ[1])
		uvdata1 = AIPSUVData(tmpNAME.popleft(), tmpKLAS.popleft(), indisk, tmpSEQ.popleft())
		uvdata2 = AIPSUVData(tmpNAME.popleft(), tmpKLAS.popleft(), indisk, tmpSEQ.popleft())
		dbcon.indata = uvdata1
		dbcon.in2data = uvdata2
		dbcon.outname = mode
		dbcon.outclass = 'DBCON'
		dbcon.outdisk = outdisk
		dbcon.go()

		indxr.indata = AIPSUVData(dbcon.outname, dbcon.outclass, indisk, dbcon.outseq)
		indxr.go()

		tmpNAME.append(mode)
		tmpKLAS.append('DBCON')
		tmpSEQ.append(j)

		j = j + 1

	print "Final file:", tmpNAME[0], tmpKLAS[0], indisk, tmpSEQ[0]

	# Run MSORT on the final DBCONd file
	print "Final sort and index for:", mode
	data = AIPSUVData(tmpNAME[0], tmpKLAS[0], indisk, tmpSEQ[0])
	msort.indata = data
	msort.outdata = data
	msort.go()
	indxr.indata = data
	indxr.go()
	outdata = data


	# Copy over the SN tables from each input file
	outdata.zap_table('SN', -1)
	for i in range(len(queueNAME)):
		datafile = AIPSUVData(queueNAME[i], queueKLAS[i], indisk, queueSEQ[i])
		if queueKLAS[i] == mode:
			print "Copying SN#1 from", datafile, "to", outdata
			tacop.indata = datafile
			tacop.outdata = outdata
			tacop.inext = 'SN'
			tacop.inver = 1
			tacop.ncount = 1
			tacop.outver = 0
			tacop.go()


	# Check if day numbers in the SN tables are consistant
	print "Checking for date rollover in SN tables."
	for i in range(len(outdata.tables)) :
		if 'SN' in outdata.tables[i][1] :
			ver = outdata.tables[i][0]
	ntables = ver
	#foreach SN table from 1 to ntables:
	#	read every line and check the date - should be zero or one
	#	if we find a one at any point, modify all subsequent tables to add a 1
	rollover = 0
	lastrow = 0
	for i in range(ver):
		sntable=outdata.table('SN',i)
		if (rollover == 0):
			for row in sntable:
				if (row.antenna_no == 1):
					print "Examining timestamps in table", i+1
					if (row.time > 1) or (row.time<lastrow):
						rollover = i
						print "Found one in table ", i+1
						break
					else:
						lastrow = row.time
	
	#rollover = i
	if rollover > 0:
		print "Date rollover detected! We need to modify every table higher than", rollover+1
		thing = ver - rollover
		print "thing, rollover", thing, range(rollover, ver)
		for i in range(rollover+1, ver+1):
#			print "!!!!!!!!!!!!!!!!!! Modifying table ", i
			tabed.indata = outdata
			tabed.inext = 'SN'
			tabed.inver = i
			tabed.outdata = outdata
			tabed.outvers = i
			tabed.bcount = 1
			tabed.ecount = 0
			tabed.optype = 'ADD'
			tabed.aparm[1:] = 1,0
			tabed.keyvalue[1:] = 1,0
			tabed.go()
	else:
		print "No data rollover issues, skipping TABED."







# Tidy up: delete original and most of DBCON files
print "Finished! Tidying up your AIPS catalogue."
pca = AIPSCat(indisk)
j = 0
for fitsfil in pca[indisk]:
	if fitsfil.klass in extensions :
		print "Zapping: " + fitsfil.name + '.' + fitsfil.klass + '.' + format(fitsfil.seq)
		uvdata1 = AIPSUVData(fitsfil.name, fitsfil.klass, indisk, fitsfil.seq)
		uvdata1.zap()
	if fitsfil.klass == 'DBCON' :
		if fitsfil.name.isdigit() :
			print "Zapping: " + fitsfil.name + '.' + fitsfil.klass + '.' + format(fitsfil.seq)
			uvdata1 = AIPSUVData(fitsfil.name, fitsfil.klass, indisk, fitsfil.seq)
			uvdata1.zap()
	j = j + 1




print "CLCALing your data."
pca = AIPSCat(indisk)
i = 0
for fitsfil in pca[indisk]:
	datafile = AIPSUVData(fitsfil.name, fitsfil.klass, indisk, fitsfil.seq)
	clcal.indata = datafile
	clcal.opcode = 'CALI'
	clcal.interpol = '2PT'
	clcal.samptype = 'MWF'
	clcal.bparm[1:] = 0.2,0.2,0.2,0.2,0.2,0.2,0.2,0.2,0.2,0.2
	clcal.snver = 0
	clcal.inver = 0
	clcal.gainver = 1
	clcal.gainuse = 2
	clcal.go()
	






print "Your AIPS catalog (user number: " + format(usernum) + ") now looks like this:"
print AIPSCat(indisk)

print "Your data are now DBCONd ready for processing."
print "You should probably edit them before calibrating."