# Automated loading and housekeeping of e-MERLIN data.
# OFF-SITE VERSION
# Original version: mkargo 20110717


####################################################################################
# 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 ".fits" 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'.
####################################################################################

# Changelog
# 20110807 - removed the need for any specific filename structure, as long as the file names 
# 	are shorter than 46 characters and end ".fits" they will be loaded.  Also now handles 
#	different directories with inseq, making catalog entries unique without risking names 
#	which are too long for AIPS.


# Import useful things

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



print "------------------------------------------------------------"
print "Load e-MERLIN 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, UVFIXd, MSORTd and INDXRd."



# 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')
uvfix = AIPSTask('UVFIX')
zap = AIPSTask('ZAP')
msort = AIPSTask('MSORT')
msort.sort = 'TB'
indxr = AIPSTask('INDXR')

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


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

datadir = os.getcwd()					# search directories below where we are now
os.listdir(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))]
	for fitsfile in mylist:				# for each file...
		if fnmatch.fnmatch(fitsfile, '*.fits'):
			#index = fitsfile.find('_')
			descriptor = "./" + dir + "/" + fitsfile
			print "Loading " + descriptor
			fitld.datain = descriptor
			outn = 'TMP'
			fitld.digicor = -1
			fitld.douvcomp = -1
			fitld.outdata = AIPSUVData(outn,'UVDATA',outdisk,thisdir)
			fitld.go()				# Load the data

			# reset the filename based on the source name in the SU table
			datafile = AIPSUVData('TMP','UVDATA',outdisk,thisdir)
			sourcelist = datafile.table('SU',1)
			outn = datafile.sources[0]
			if len(outn)>12 :
				outn = outn[0:12]		# catch long source names
			inname = outn

			# uvfix each file
			uvfix.indata =  AIPSUVData('TMP','UVDATA',indisk,thisdir)
			uvfix.outdata = AIPSUVData(inname,'UVFIX',indisk,thisdir)
			uvfix.go()

			# msort each file, writing back to same file
			msort.indata = AIPSUVData(inname,'UVFIX',indisk,thisdir)
			msort.outdata = AIPSUVData(inname,'UVFIX',indisk,thisdir)
			msort.go()

			# index each file
			indxr.indata = AIPSUVData(inname,'UVFIX',indisk,thisdir)
			indxr.go()

			# remove now-superfluous UVDATA files
			mydatain = AIPSUVData('TMP','UVDATA',indisk,thisdir)
			mydatain.zap()



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

print "Your data are now loaded in AIPS under user number " + format(AIPS.userno) + " on disk " + format(outdisk) + "."
print "You should carefully edit them before proceeding with calibration."