# Automated merging and averaging of e-MERLIN data.
# Original version: mkargo 20110717

###############################################################################
# For each file in the catalog specified by the user:
# 	query header to find number of channels
#	uvavg in time if requested
# 	avspc to requested number of channels if requested
#	bdcon files together
# Then clean up, leaving just the UVFIX files and the final DBCON database.
#
# If the script fails, the first thing to do is run 'source /aips/LOGIN.CSH'.
###############################################################################

# Changelog
# 20120106 - Neatened up the loops
# 20111024 - Added option for no spectral averaging.
# 20110928 - Added time averaging capability using UVAVG, no averaging is an 
#	option.  Also added a check for UVFIX files, exits if none are found.


##########################################################
# You shouldn't need to change anything!		 #
##########################################################


# Import useful things

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


print "-----------------------------------------------------------"
print "Merge and average e-MERLIN data already loaded and flagged."
print "-----------------------------------------------------------"
print "Your data should already be loaded, UVFIXd, MSORTd, INDXRd "
print "and flagged spectrally with SPFLG. If this is not the case,"
print "please quit now and come back here later."
print "-----------------------------------------------------------"

# 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 where your data are: ",
indisk = raw_input()
if indisk.isdigit() :
	indisk = int(indisk)
else :
	print "Does not compute."
	sys.exit()
outdisk = indisk
print "Your data will be DBCONd and averaged together."
print "Enter the desired number of channels per IF (-1 for no averaging): ",
nchan = raw_input()
if nchan=='-1' or nchan.isdigit() :
	nchan = int(nchan)
else :
	print "Does not compute.", nchan
	sys.exit()
print "If requested, your data will also be time averaged with UVAVG."
print "Enter the desired averaging time (-1 for no averaging): ",
tint = raw_input()
if tint=='-1' or tint.isdigit() :
	tint = int(tint)
else :
	print "Does not compute.", tint
	sys.exit()


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

# Set up AIPS things

today = date.today()
now = today.strftime("%Y%m%d")
list=[]
AIPS.userno = usernum
tasav = AIPSTask('TASAV')
uvavg = AIPSTask('UVAVG')
avspc = AIPSTask('AVSPC')
dbcon = AIPSTask('DBCON')
msort = AIPSTask('MSORT')
indxr = AIPSTask('INDXR')

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


pca = AIPSCat(indisk)
print "Your AIPS catalog looks like this:"
print AIPSCat(indisk)

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


# First, check we have no stale files
todo = 0
backedup = 0
for fitsfil in pca[indisk]:
	if fitsfil.klass == 'SUB SP' :
		print "Error: please remove all AVSPC files before running this script!"
		sys.exit()
	if fitsfil.klass == 'DBCON' :
		print "Error: please remove all DBCON files before running this script!"
		sys.exit()
	if fitsfil.klass == 'UVFIX' :
		todo = todo + 1
	if fitsfil.klass == 'TASAV' :
		backedup = 1
	i = i + 1


i = 0

# Check we have some files to work with
if todo == 0 :
	print "Error: no UVFIX files found! Are you sure you loaded your data?"
	sys.exit()


print "backup = ", backedup
# Backup the tables if we need to
if backedup != 1 :
	for fitsfil in pca[indisk]:
		uvdata = AIPSUVData(fitsfil.name, fitsfil.klass, indisk, fitsfil.seq)
		saved = AIPSUVData(fitsfil.name, "TASAV", indisk, fitsfil.seq)
		print "backing up ", uvdata
		tasav.indata = uvdata
		tasav.outdata = saved
		tasav.go()
		i = i + 1

i = 0

# UVAVG only if requested - note that this applies the highest flag table!
if tint != -1 :
	print "Averaging in time with UVAVG."
	for fitsfil in pca[indisk]:
		if fitsfil.klass == 'UVFIX' :
			filelist.append(fitsfil.name)
			uvdata = AIPSUVData(fitsfil.name, fitsfil.klass, indisk, fitsfil.seq)
			print "Time averaging " + fitsfil.name + '.' + fitsfil.klass + '.' + format(fitsfil.seq)
			uvavg.indata = uvdata
			uvavg.outname = fitsfil.name
			uvavg.outdisk = outdisk
			uvavg.outclass = 'UVAVG'
			uvavg.outseq = fitsfil.seq
			uvavg.doacor = 1
			uvavg.yinc = tint
			#uvavg.flagver = -1   # uncomment to avoid applying flag tables at this point
			uvavg.go()

			# Copy over flag tables
			#tacop.indata = uvdata
			#tacop.outname = pca[indisk][i]["name"]
			#tacop.outclass = 'UVAVG'
			#tacop.outdisk = outdisk
			#tacop.inext = 'FG'
			#tacop.invers = 1
			#tacop.ncount = 0   # these two copy *every* FG table
			#tacop.go()

			# Remove unaveraged file, then rename the averaged file
			uvdata.zap()
			uvdata = AIPSUVData(fitsfil.name,'UVAVG',indisk,fitsfil.seq)
			uvdata.rename(fitsfil.name,'UVFIX',fitsfil.seq)

		i = i + 1
else :
	print "You requested no time averaging, skipping."


i = 0
# AVSPC files only if they are of type UVFIX
if nchan != -1 :
	print "Averaging channels with AVSPC."
	for fitsfil in pca[indisk]:
		if fitsfil.klass == 'UVFIX' :
			filelist.append(fitsfil.name)
			uvdata = AIPSUVData(fitsfil.name, fitsfil.klass, indisk, fitsfil.seq)
			uvdata.header['naxis'][2] # pulls out the number of channels
			print "Averaging " + fitsfil.name + '.' + fitsfil.klass + '.' + format(fitsfil.seq)
			avspc.indata = uvdata
			avspc.outdisk = outdisk
			avspc.channel = uvdata.header['naxis'][2] / nchan
			avspc.avoption = 'SUBS'
			avspc.go()
		i = i + 1
else :
	print "You requested no spectral averaging, skipping."



print "Averaging complete, DBCONing."


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


# DBCON files only if they are of type 'SUB SP'
for fitsfil in pca[indisk]:
	if ((nchan != -1) and (fitsfil.klass == 'SUB SP')) :
		queueNAME.append(fitsfil.name)
		queueKLAS.append(fitsfil.klass)
		queueSEQ.append(fitsfil.seq)
	if ((nchan == -1) and (fitsfil.klass == 'UVFIX')) :
		queueNAME.append(fitsfil.name)
		queueKLAS.append(fitsfil.klass)
		queueSEQ.append(fitsfil.seq)
	i = i + 1



i = 0
while len(queueNAME)>1 :
	print "Combining " + queueNAME[0] + '.' + queueKLAS[0] + '.' + format(queueSEQ[0]) + " and " + queueNAME[1] + '.' + queueKLAS[1] + '.' + format(queueSEQ[1])
	uvdata1 = AIPSUVData(queueNAME.popleft(), queueKLAS.popleft(), indisk, queueSEQ.popleft())
	uvdata2 = AIPSUVData(queueNAME.popleft(), queueKLAS.popleft(), indisk, queueSEQ.popleft())
	dbcon.indata = uvdata1
	dbcon.in2data = uvdata2
	dbcon.outname = format(i)
	dbcon.outclass = 'DBCON'
	dbcon.outdisk = outdisk
	dbcon.go()

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

	queueNAME.append(format(i))
	queueKLAS.append('DBCON')
	queueSEQ.append(1)

	i = i + 1


# Run MSORT on the final DBCONd file
print "Final sort and index...."
data = AIPSUVData(queueNAME[0], queueKLAS[0], indisk, queueSEQ[0])
msort.indata = data
msort.outdata = data
msort.go()
indxr.indata = data
indxr.go()
data.rename('ALL','DBCON',0)


# Tidy up: delete AVSPC and most of DBCON files
print "Tidying up your AIPS catalogue."
pca = AIPSCat(indisk)
j = 0
for fitsfil in pca[indisk]:
	if fitsfil.klass == 'SUB SP' :
		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 "Your AIPS catalog now looks like this:"
print AIPSCat(indisk)

print "Your data are now DBCONd and reduced to " + format(nchan) + " channels on disk " + format(outdisk) + "."
print "You should probably edit them before calibrating."