#!/bin/env python

from sys import argv, exit
from string import split, upper, strip
from os import system, popen, stat
from os.path import isfile
from pwd import getpwuid
import datetime
from pyx import *

program = 'plottsys'
version = '0.2'
verdate = '20130111'
author = 'Walter Brisken <wbrisken@nrao.edu>'

mjd0 = datetime.datetime(1858, 11, 17, 0, 0)


def usage():
	print '\n%s ver. %s  %s  %s\n' % (program, version, verdate, author)
	print 'Usage: %s [options] <tsysfile> <vexfile> [<plotfile>]\n' % argv[0]
	print 'options can be:\n'
	print '  -h or --help       print help info and quit.'
	print '  -v or --verbose    make more verbose.'
	print '  -q or --quiet      make less verbose.\n'
	print 'At least one of <tsysfile> or <vexfile> must be provided.'
	print 'If vexfile is not provided, some information will be left off the plots.'
	print ''

def mjd2yd(mjd):
	mjdi = int(mjd)
	dt = datetime.timedelta(mjdi, 0)
	t = mjd0 + dt
	tt = t.timetuple()
	return [tt.tm_year, tt.tm_yday + (mjd-mjdi)]

def readTsysFile(tsysFile, verbose):
	
	data = {}	# indexed by antenna name
	nBad = 0
	nComment = 0
	nData = 0
	nTsys = 0
	nOverflow = 0
	nMask = 0

	if verbose > 0:
		print 'Opening: %s' % tsysFile
	lines = open(tsysFile, 'r').readlines()
	lastAnt = ''
	for line in lines:
		s = split(split(line, '#')[0])
		l = len(s)
		if l < 4:
			nComment += 1
			continue
		nChan = int(s[3])
		if l != nChan*2 + 4:
			if verbose > 1:
				print 'Malformed line: %s' % line
			nBad += 1
		else:
			nData += 1
		if s[0] != lastAnt:
			lastAnt = s[0]
			if verbose > 0:
				print 'Reading antenna: %s' % lastAnt
			if data.has_key(lastAnt):
				antArray = data[lastAnt]
			else:
				antArray = []
				data[lastAnt] = antArray
		while len(antArray) < nChan:
			antArray.append([[],[]])
		doy = float(s[1])
		for c in range(nChan):
			tsys = float(s[4 + 2*c])
			if tsys == 999:
				nMask += 1
			elif tsys < 999.0 and tsys >= 0.0:
				antArray[c][0].append(doy)
				antArray[c][1].append(tsys)
				nTsys += 1
			else:
				nOverflow += 1
	
	if verbose > 0:
		print '%d lines were malformed' % nBad
		print '%d lines were comments' % nComment
		print '%d lines were valid' % nData
		print '%d Tsys values were read' % nTsys
		print '%d overflow values were read' % nOverflow
		print '%d masked values were read' % nMask

	return data


def plotAntennaTsys(pages, data, antName, obsCode, timeRange, verbose):
	rv = 0
	lastPage = -1

	tmax = 0
	tmin = 999
	xmax = 0
	xmin = 1e9

	if len(data[0][0]) == 0:
		print '  Warning: no Tsys data for station %s' % antName
		return 0

	for chan in range(len(data)):
		chanData = data[chan]
		if len(chanData[0]) == 0 or len(chanData[1]) == 0:
			print 'Warning: no data for %s 0-based channel %d' % (antName, chan)
		else:
			m = max(chanData[0])
			if m > xmax:
				xmax = m
			m = min(chanData[0])
			if m < xmin:
				xmin = m
			m = max(chanData[1])
			if m > tmax:
				tmax = m
			m = min(chanData[1])
			if m < tmin:
				tmin = m
	
	if tmin >= tmax:
		print 'Error: No data to plot!'

		return -1

	if len(timeRange) == 2:
		startTime = mjd2yd(timeRange[0])
		endTime = mjd2yd(timeRange[1])

		if verbose > 1:
			print '  Expt time range: %11.5f - %11.5f  =  %d/%f - %d/%f' % (timeRange[0], timeRange[1], startTime[0], startTime[1], endTime[0], endTime[1])
		if xmin > startTime[1]+0.05:
			dt = int((xmin-startTime[1])*1440.0 + 0.5)
			print '  Warning: Tsys data appears incomplete at beginning of experiment by %d minutes' % dt
			rv = 1
		if xmax < endTime[1]-0.05:
			dt = int((endTime[1]-xmax)*1440.0 + 0.5)
			print '  Warning: Tsys data appears incomplete at end of experiment by %d minutes' % dt
			rv = 1

		xmin = startTime[1]
		xmax = endTime[1]

	# reset bottom of range to 0K
	tmin = 0

	for chan in range(len(data)):
		chanData = data[chan]
		col = chan/4
		row = 3-(chan%4)
		page = chan/8

		if page != lastPage:
			lastPage = page
			c = canvas.canvas()
			p = document.page(c, paperformat=document.paperformat.Letter)
			pages.append(p)

		if len(chanData[0]) < 2:
			continue;

		g = graph.graphxy(width=8, height=5, \
			x=graph.axis.linear(title='Day of year', min=xmin, max=xmax), \
			y=graph.axis.linear(title='$T_{\\rm sys}$ [K]', min=tmin, max=tmax))
		g.plot(graph.data.values(x=chanData[0], y=chanData[1]), [graph.style.symbol(graph.style.symbol.circle, size=0.02)])
		g.text(0.25, 4.75, '%s Chan %d' % (antName,chan), [text.halign.left, text.valign.top])
		g.text(7.75, 4.75, '%s' % obsCode, [text.halign.right, text.valign.top])
		g.dolayout()

		c.insert(g, [trafo.translate(col*10, row*7)])

	return rv

def plotTsys(tsysFile, plotFile, obsCode, stationTimes, verbose):
	data = readTsysFile(tsysFile, verbose)
	pages = []
	ants = data.keys()
	ants.sort()
	nPartial = 0
	for a in ants:
		if stationTimes.has_key(a):
			timeRange = stationTimes[a]
		else:
			timeRange = []
		if verbose > 0:
			print 'Plotting antenna: %s' % a
		rv = plotAntennaTsys(pages, data[a], a, obsCode, timeRange, verbose)
		nPartial += rv
	doc = document.document(pages)
	doc.writePDFfile(plotFile, title=tsysFile)

	if nPartial > 0:
		print ''
		print 'NOTE: %d stations only had partial Tsys for the experiment plotted.' % nPartial
		print 'Some of the data may not have been transferred to Socorro yet.'
		print 'Rerunning this plotting program several minutes after the experiment'
		print 'may yield more complete results.\n'
	
	print 'On Linux the best way to view the results is with:\n'
	print '  evince %s\n' % plotFile

def vexPeek(vexFile):
	
	cmd = 'vexpeek %s' % vexFile
	if verbose > 0:
		print 'Executing command: %s' % cmd

	p = popen(cmd)
	data = p.readlines()

	if len(data) == 0:
		return 'Error', 'Error', 'Error'
	
	obsCode = upper(strip(data[0]))
	obsSeg = ''

	if obsCode[0:5] == 'ERROR':
		return 'Error', 'Error', 'Error'

	if len(obsCode) > 3:
		if obsCode[0].isalpha() and obsCode[1].isalpha() and obsCode[2].isdigit():
			for i in range(3, len(obsCode)):
				if obsCode[i].isalpha():
					obsSeg = obsCode[i:]
					obsCode = obsCode[0:i]
					break
		if obsCode[0].isalpha() and obsCode[1].isdigit():
			for i in range(2, len(obsCode)):
				if obsCode[i].isalpha():
					obsSeg = obsCode[i:]
					obsCode = obsCode[0:i]
					break
		

	stationTimes = {}
	for d in data[1:]:
		s = split(strip(d))
		stationTimes[upper(s[0])] = [float(s[1]), float(s[2])]

	print 'This is experiment %s %s' % (obsCode, obsSeg)

	return obsCode, obsSeg, stationTimes


verbose = 1
tsysFile = ''
plotFile = ''
vexFile = ''
obsCode = ''
generateTsys = False
stationTimes = {}

for a in argv[1:]:
	if a[0] == '-':
		if a in ['-v', '--verbose']:
			verbose += 1
		elif a in ['-q', '--quiet']:
			verbose += 1
		elif a in ['-h', '--help']:
			usage()
			exit(0)
		else:
			print 'Error: unknown option: %s' % a
			exit(0)
	elif a[-4:] == '.pdf':
		if plotFile != '':
			print 'Error: plot file provided more than once'
			exit(0)
		plotFile = a
	elif a[-4:] == '.vex':
		if vexFile != '':
			print 'Error: vex file provided more than once'
			exit(0)
		vexFile = a
	elif tsysFile == '':
		tsysFile = a
	elif obsCode == '':
		obsCode = a
	else:
		print 'Error: unexpected argument: %s' % a
		exit(0)

if vexFile != '':
	if not isfile(vexFile):
		print 'Error: file %s not found\n' % vexFile
		exit(0)
	if tsysFile == '':
		tsysFile = '/tmp/' + split(vexFile[0:-4], '/')[-1] + '.tsys.avg'
		generateTsys = True
	obsCode, obsSeg, stationTimes = vexPeek(vexFile)
	obsCode = obsCode + obsSeg
	if plotFile == '':
		plotFile = obsCode + '.tsys.pdf'

if plotFile == '':
	plotFile = tsysFile + '.pdf'

try:
	open(plotFile, 'w')
except IOError:
	print '\nError:  You are unable to write to output file %s.  This likely means you don\'t have write permission in this directory.  Try running as %s or change to a directory where you have write permission.\n\nHint: At NRAO, vlbaops should probably be running this.\n' % (plotFile, getpwuid(stat('.').st_uid).pw_name)
	exit(0)

if generateTsys:
	cmd = 'rdbetsys --interval 300 %s %s' % (vexFile, tsysFile)
	if verbose > 0:
		print 'Executing: %s' % cmd
	system(cmd)

if tsysFile == '':
	print 'Error: Tsys filename or vex filename required.'
	exit(0)
else:
	if not isfile(tsysFile):
		print 'Error: file %s not found\n' % tsysFile
		exit(0)

if verbose > 1:
	print 'vexFile = %s' % vexFile
	print 'tsysFile = %s' % tsysFile
	print 'plotFile = %s' % plotFile

plotTsys(tsysFile, plotFile, obsCode, stationTimes, verbose)