#!/usr/bin/env python

#**************************************************************************
#   Copyright (C) 2008-2011 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, exit
from os import chdir, popen, system, getcwd
from string import strip, split, find, lower, upper

program = 'difxsniff'
version = '2.0'
verdate = '20100902'
author  = 'Walter Brisken'

def usage():
	print '\n%s ver. %s   %s %s' % (program, version, verdate, author)
	print '\nA program to "sniff" FITS files made by DiFX\n'
	print 'Usage: ./difxsniff [options] <refants> <filebase1> [<filebase2> [...] ]'
	print '\nOptions can include:'
	print '  -h or --help : print this info and quit.'
	print 'each "file" is a FITS file made by difx2fits\n'

def reverseapd(dataline):
	s = split(strip(dataline))
	str = '%5d %8.5f %2d %-10s %2d %2d %-3s %-3s %2d' % \
		(int(s[0]), float(s[1]), int(s[2]), s[3], int(s[4]), int(s[5]), s[7], s[6], int(s[8]))
	n = int(s[8])
	for i in range(n):
		s2 = ' %10.4f %7.5f %10.4f %10.6f' % \
			(-float(s[9+4*i]), float(s[10+4*i]), -float(s[11+4*i]), -float(s[12+4*i]))
		str += s2
	str += '\n'
	return str

# Executes a program wanting command line input.  The args here are those
# inputs in order.
def run(pgm, args):
	print pgm
	p = popen(pgm, 'w')
	for a in args:
		p.write('%s\n' % a)
	p.close()

def sniff(filebases, refants, dir, newNames):

	if len(dir) > 0:
		system('mkdir -p ' + dir)

	# make apdfile
	apd = open(dir+'datafile.lis', 'w')
	first = 1
	for f in filebases:
		f = open(f+'.apd', 'r')
		l = f.readline()
		if first:
			if l[0:8] != 'obscode:':
				apd.write('obscode:  ')
			apd.write(l)
			first = 0
		if len(refants) > 1:
			while 1:
				l = f.readline()
				if l == '':
					break
				for r in refants:
					if find(l, r) > 0:
						apd.write(d)
						break
		else:
			while 1:
				l = f.readline()
				if l == '':
					break
				s = split(l)
				if s[6] == refants[0]:
					apd.write(l);
				elif s[7] == refants[0]:
					apd.write(reverseapd(l))
		f.close()
	apd.close()

	# make wtsfile
	wts = open(dir+'wtsfile.lis.sum', 'w')
	first = 1
	for f in filebases:
		f = open(f+'.wts', 'r')
		l = f.readline()
		# only write the first line for the first file
		if first:
			first = 0
			wts.write(l)
		while 1:
			l = f.readline()
			if l == '':
				break
			wts.write(l)
		f.close()
	wts.close()

	# make acbandfile
	acb = open(dir+'acbandfile.lis', 'w')
	for f in filebases:
		f = open(f+'.acb', 'r')
		for l in f:
			acb.write(l)
		f.close()
	acb.close()

	# make xcbandfile
	xcb = open(dir+'xcbandfile.lis', 'w')
	for f in filebases:
		f = open(f+'.xcb', 'r')
		while 1:
			data = []
			l = f.readline()
			if l == '':
				break
			data.append(l)
			v = split(l)
			nif = int(v[-1])
			nch = int(v[-3])
			n = 2 + nif + nif*nch
			for i in range(1,n):
				l = f.readline()
				if l == '':
					print 'premature end of file: ', (f+'.xcb')
					break
				data.append(l)
			if len(data) != n:
				break;
			hasref = 0
			for r in refants:
				if find(data[2+nif], r) > 0:
					hasref += 1
			if hasref > 0:
				for l in data:
					xcb.write(l)
		f.close()
	xcb.close()
	
	if len(dir) > 0:
		chdir(dir)

	if newNames:
		base = split(filebases[0], '/')[-1]
		apdfilename = '%s.apd.ps' % base
		wtsfilename = '%s.wts.ps' % base
		acbfilename = '%s.acb.ps' % base
		xcbfilename = '%s.xcb.ps' % base
	else:
		apdfilename = 'apdfile.ps'
		wtsfilename = 'wtsfile.ps'
		acbfilename = 'acbandfile.ps'
		xcbfilename = 'xcbandfile.ps'

	run('plotapd', ['datafile.lis',    '%s/vps' % apdfilename, 'ALL', 'ALL'])
	run('plotwt',  ['wtsfile.lis.sum', '%s/vps' % wtsfilename])
	run('plotbp',  ['acbandfile.lis',  '%s/vps' % acbfilename, ''])
	run('plotbp',  ['xcbandfile.lis',  '%s/vps' % xcbfilename, ''])
	
files = []
refants = []
if getcwd()[-7:] == 'sniffer':
	dir = getcwd() + '/'
else:
	dir = 'sniffer/'

if len(argv) < 2:
	usage()
	exit(0)

for a in argv[1:]:
	p = find(a, '.FITS')
	if p < 1:
		p = find(a, '.fits')
	if p > 0:
		files.append(a[:p])
	elif find(a, '.') >= 0 and len(a) > 2:
		files.append(a)
	elif find(a, '/') > 0:
		dir = a
	elif a[0] == '-':
		if a == '-h' or a == '--help':
			usage()
			exit(0)
		else:
			print 'Error: unknown option : ' + a
	else:
		refants.append(upper(a))

if len(files) == 0:
	print 'Error: no input files provided.'
	exit(0)

if len(refants) == 0:
	print 'Error: no refants provided.'
	exit(0)

if dir[-1] != '/':
	dir += '/'

if len(refants) == 1:
	dir += refants[0] + '/'

nFITS = 0
for f in files:
	if f[-5:] == '.FITS':
		nFITS += 1

if nFITS == 0:
	for f in files:
		sniff([f], refants, dir, True)
elif nFITS == len(files):
	sniff(files, refants, dir, False)
else:
	print 'Error -- mixed FITS files and file prefixes'