#!/usr/bin/env python

#**************************************************************************
#   Copyright (C) 2014 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.             *
#**************************************************************************


from sys import argv, exit
from string import split, strip
from math import *
from pylab import np, plot, show, scatter, xlabel, ylabel, title, axes, savefig

program = 'plotpcal'
author = 'Walter Brisken <wbrisken@nrao.edu>'
version = '0.1'

legalOps = ['amp', 'phase', 'delay', 'xy']

def usage(prog):
	print '\n%s ver. %s  %s\n' % (program, version, author)
	print 'Usage: %s [options] operation toneSelections dataFiles\n' % prog

# returns: op, toneSelect, files, options
def parseCommandLine():
	op = None
	toneSelect = []
	data = []
	options = {}
	options['verbose'] = 1
	for a in argv[1:]:
		if a[0] == '-':
			# this is an option
			if a in ['-h', '--help']:
				usage(argv[0])
				exit(0)
			elif a in ['-v', '--verbose']:
				options['verbose'] += 1
			elif a in ['-q', '--quiet']:
				options['verbose'] -= 1
			else:
				if len(a) > 2 and a[1] == '-':
					options[a[2:]] = 1
				else:
					options[a[1:]] = 1
		elif '=' in a:	
			# this is a tone selector
			toneSelect.append(a)
		elif op == None:
			op = a
		else:
			data.append(a)
	
	die = False
	if len(data) == 0:
		print 'Error: no data files provided'
		die = True
	if len(toneSelect) == 0:
		print 'Error: no tone selection'
		die = True
	if op == None:
		print 'Error: operation not specified'
		die = True
	elif not op in legalOps:
		print 'Error: operation must be one of ', legalOps
		die = True

	if die:
		print '\nRun with --help for more info\n'
		exit(0)

	if options['verbose'] > 1:
		print 'op = %s' % op
		print 'toneSection =', toneSelect
		print 'data files =', data
		print 'options =', options

	return op, toneSelect, data, options

def fileSummary(filename):
	start = None
	stop = None
	pcalVersion = 0
	data = open(filename, 'r').readlines()
	for d in data:
		if d[0] == '#':
			if d[:17] == '# File version = ':
				pcalVersion = int(d[17:])
		else:
			s = split(d)
			ant = s[0]
			start = float(s[1])
			firstLine = strip(d)
			break

	for d in reversed(data):
		if d[0] != '#':
			s = split(d)
			stop = float(s[1])
			break
	
	if start == None or stop == None:
		print 'Error: %s seems ill-formatted.  Is this a DiFX PCAL file?' % filename
		exit(0)

	return ant, start, stop, firstLine, pcalVersion

def generateDataDictionary(data):
	dataDict = {}
	for filename in data:
		ant, start, stop, firstLine, pcalVersion = fileSummary(filename)
		if not ant in dataDict:
			dataDict[ant] = []
		dataDict[ant].append([start, stop, filename, firstLine, pcalVersion])

	return dataDict

def summarizeData(dataDict):
	ants = dataDict.keys()
	ants.sort()
	for a in ants:
		print 'Antenna %s:' % a
		data = dataDict[a]
		data.sort()
		for d in data:
			print '  %s from %f to %f ver %d' % (d[2], d[0], d[1], d[4])

# each tone has 4 columns.  This function returns the index to the first of them (0-based)
# the columns are: recBand (0-based), toneFreq (MHz), re, im
def getOneToneColumns(toneSelect, dataLine, pcalVersion):

	toneOffsets = [9, 6]	# indexed by pcalVersion

	columns = []
	s = split(strip(dataLine))
	toneOffset = toneOffsets[pcalVersion]	# index to first tone

	if pcalVersion == 0:
		nPol = int(s[4])
		nFreq = int(s[5])
		nTone = int(s[6])
		nRecBand = int(s[8])
		totalTones = nPol*nFreq*nTone
		for ts in toneSelect:
			tss = split(ts, '=')
			if tss[0] == 'rt':
				# look for specific tone number.  Should be unambiguous
				cd = split(tss[1], ',')	# column identifying data
				rb = int(cd[0])		# rec band (0 based)
				if rb < 0:
					rb += nRecBand
				tn = int(cd[1])		# tone num
				if tn < 0:
					# negative number means count from back, python style
					tn += nTone
				if tn < nTone:
					for t in range(0, totalTones, nTone):
						if int(s[toneOffset + 4*t]) == rb:
							columns.append(toneOffset+4*(t+tn))
							break
			elif tss[0] == 'rf':
				# look for frequency within a specific record band.  Should be unambiguous
				cd = split(tss[1], ',')
				rb = int(cd[0])		# rec band (0 based)
				if rb < 0:
					rb += nRecBand
				if cd[1][0] == '%':
					fq = int(cd[1][1:])	# freq (MHz)
					for t in range(totalTones):
						if int(s[toneOffset + 4*t]) == rb and int(s[toneOffset + 4*t + 1]) % fq == 0:
							columns.append(toneOffset+4*t)
				else:
					fq = int(cd[1])	# freq (MHz)
					for t in range(totalTones):
						if int(s[toneOffset + 4*t]) == rb and int(s[toneOffset + 4*t + 1]) == fq:
							columns.append(toneOffset+4*t)
							break
			elif tss[0] == 'f':
				# look for specific frequency.  There could be many
				cd = split(tss[1], ',')
				if cd[0][0] == '%':
					fq = float(cd[0][1:])	# freq (MHz)
					for t in range(totalTones):
						if int(s[toneOffset + 4*t]) >= 0:
							if int(s[toneOffset + 4*t + 1]) % fq == 0:
								columns.append(toneOffset+4*t)
				else:
					fq = float(cd[0])	# freq (MHz)
					for t in range(totalTones):
						if int(s[toneOffset + 4*t]) >= 0:
							if int(s[toneOffset + 4*t + 1]) == fq:
								columns.append(toneOffset+4*t)
			elif tss[0] == 'r':
				# get all tones for one record band.
				cd = split(tss[1], ',')
				rb = int(cd[0])		# rec band (0 based)
				if rb < 0:
					rb += nRecBand
				for t in range(totalTones):
					if int(s[toneOffset + 4*t]) == rb:
						columns.append(toneOffset+4*t)
			elif tss[0] == 't':
				# get specific tone number for all recorded bands, or "all", meaning all tones
				cd = split(tss[1], ',')
				if cd[0] == 'all':
					for t in range(totalTones):
						if int(s[toneOffset + 4*t]) >= 0:
							columns.append(toneOffset+4*t)
				else:
					tn = int(cd[0])		# tone num
					if tn < 0:
						# negative number means count from back, python style
						tn += nTone
					if tn < nTone:
						for t in range(tn, totalTones, nTone):
							if int(s[toneOffset + 4*t]) >= 0:
								columns.append(toneOffset+4*t)
	elif pcalVersion == 1:
		nTone = int(s[5])
		nRecBand = int(s[4])
		totalTones = nRecBand*nTone
		for ts in toneSelect:
			tss = split(ts, '=')
			if tss[0] == 'rt':
				# look for specific tone number.  Should be unambiguous
				cd = split(tss[1], ',')	# column identifying data
				rb = int(cd[0])		# rec band (0 based)
				tn = int(cd[1])		# tone num
				if tn < 0:
					# negative number means count from back, python style
					tn += nTone
				if rb < 0:
					rb += nRecBand
				if tn >= 0 and tn < nTone and rb >= 0 and rb < nRecBand:
					columns.append(toneOffset + 4*(tn + nTone*rb))
			elif tss[0] == 'rf':
				# look for frequency within a specific record band.  Should be unambiguous
				cd = split(tss[1], ',')
				rb = int(cd[0])		# rec band (0 based)
				if rb < 0:
					rb += nRecBand
				if rb >= 0 and rb < nRecBand:
					if cd[1][0] == '%':
						fq = int(cd[1][1:])	# freq (MHz)
						for t in range(nTone):
							if int(s[toneOffset + 4*(rb*nTone + t)]) % fq == 0:
								columns.append(toneOffset + 4*(rb*nTone + t))
					else:
						fq = int(cd[1])	# freq (MHz)
						for t in range(nTone):
							if int(s[toneOffset + 4*(rb*nTone + t)]) == fq:
								columns.append(toneOffset + 4*(rb*nTone + t))
								break
			elif tss[0] == 'f':
				# look for specific frequency.  There could be many
				cd = split(tss[1], ',')
				if cd[0][0] == '%':
					fq = float(cd[0][1:])	# freq (MHz)
					for t in range(totalTones):
						if int(s[toneOffset + 4*t]) >= 0:
							if int(s[toneOffset + 4*t]) % fq == 0:
								columns.append(toneOffset+4*t)
				else:
					fq = float(cd[0])	# freq (MHz)
					for t in range(totalTones):
						if int(s[toneOffset + 4*t]) >= 0:
							if int(s[toneOffset + 4*t + 1]) == fq:
								columns.append(toneOffset+4*t)
			elif tss[0] == 'r':
				# get all tones for one record band.
				cd = split(tss[1], ',')
				rb = int(cd[0])		# rec band (0 based)
				if rb < 0:
					rb += nRecBand
				if rb >= 0 and rb < nRecBand:
					for t in range(nTone):
						columns.append(toneOffset+4*(rb*nTone + t))
			elif tss[0] == 't':
				# get specific tone number for all recorded bands, or "all", meaning all tones
				cd = split(tss[1], ',')
				if cd[0] == 'all':
					for t in range(totalTones):
						if int(s[toneOffset + 4*t]) >= 0:
							columns.append(toneOffset+4*t)
				else:
					tn = int(cd[0])		# tone num
					if tn < 0:
						# negative number means count from back, python style
						tn += nTone
					if tn < nTone:
						for t in range(tn, totalTones, nTone):
							if int(s[toneOffset + 4*t]) >= 0:
								columns.append(toneOffset+4*t)
	else:
		print 'Error: pcal file with version %d is not supported' % pcalVersion
		exit(0)
	
	return columns

def getTwoToneColumns(toneSelect, dataLine, pcalVersion):

	toneOffsets = [9, 6]	# indexed by pcalVersion
	
	columns = []
	s = split(strip(dataLine))
	toneOffset = toneOffsets[pcalVersion]	# index to first tone

	if pcalVersion == 0:
		nPol = int(s[4])
		nFreq = int(s[5])
		nTone = int(s[6])
		nRecBand = int(s[8])
		totalTones = nPol*nFreq*nTone
		for ts in toneSelect:
			tss = split(ts, '=')
			if tss[0] == 'rt':
				# look for specific tone number.  Should be unambiguous
				cd = split(tss[1], ',')	# column identifying data
				rb = int(cd[0])		# rec band (0 based)
				if rb < 0:
					rb += nRecBand
				if cd[1] == 'pairs':
					for tn1 in range(0, nTone, 2):
						tn2 = tn1+1
						for t in range(0, totalTones, nTone):
							if int(s[toneOffset + 4*t]) == rb:
								columns.append([toneOffset+4*(t+tn1), toneOffset+4*(t+tn2)])
								break

				else:
					if cd[1] == 'ends':
						tn1 = 0
						tn2 = -1
					elif cd[1] == 'good':
						if nTone == 2:
							tn1 = 0
							tn2 = 1
						elif nTone == 4:
							tn1 = 1
							tn2 = 2
						else:
							tn1 = nTone / 8
							tn2 = nTone - 1 - tn1
					else:
						tn1 = int(cd[1])	# first tone num
						tn2 = int(cd[2])	# last tone num
					if tn1 < 0:
						# negative number means count from back, python style
						tn1 += nTone
					if tn2 < 0:
						# negative number means count from back, python style
						tn2 += nTone
					if tn1 < nTone and tn2 < nTone:
						for t in range(0, totalTones, nTone):
							if int(s[toneOffset + 4*t]) == rb:
								columns.append([toneOffset+4*(t+tn1), toneOffset+4*(t+tn2)])
								break
			elif tss[0] == 't':
				# look for specific tone number.  Should be unambiguous
				cd = split(tss[1], ',')	# column identifying data
				if cd[0] == 'ends':
					tn1 = 0
					tn2 = -1
				elif cd[0] == 'good':
					if nTone == 2:
						tn1 = 0
						tn2 = 1
					elif nTone == 4:
						tn1 = 1
						tn2 = 2
					else:
						tn1 = nTone / 8
						tn2 = nTone - 1 - tn1
				else:
					tn1 = int(cd[0])	# first tone num
					tn2 = int(cd[1])	# last tone num
				if tn1 < 0:
					# negative number means count from back, python style
					tn1 += nTone
				if tn2 < 0:
					# negative number means count from back, python style
					tn2 += nTone
				if tn1 < nTone and tn2 < nTone:
					for t in range(0, totalTones, nTone):
						if int(s[toneOffset + 4*(t+tn1)]) != -1 and int(s[toneOffset + 4*(t+tn2)]) != -1:
							columns.append([toneOffset+4*(t+tn1), toneOffset+4*(t+tn2)])
	elif pcalVersion == 1:
		nTone = int(s[5])
		nRecBand = int(s[4])
		totalTones = nRecBand*nTone
		for ts in toneSelect:
			tss = split(ts, '=')
			if tss[0] == 'rt':
				# look for specific tone number.  Should be unambiguous
				cd = split(tss[1], ',')	# column identifying data
				rb = int(cd[0])		# rec band (0 based)
				if rb < 0:
					rb += nRecBand
				if rb >= 0 and rb < nRecBand:
					if cd[1] == 'pairs':
						for tn1 in range(0, nTone, 2):
							tn2 = tn1+1
							if int(s[toneOffset+4*(rb*nTone+tn1)]) != -1 and int(s[toneOffset+4*(rb*nTone+tn2)]) != -1:
								columns.append([toneOffset+4*(rb*nTone+tn1), toneOffset+4*(rb*nTone+tn2)])

					else:
						if cd[1] == 'ends':
							tn1 = 0
							tn2 = -1
						elif cd[1] == 'good':
							if nTone == 2:
								tn1 = 0
								tn2 = 1
							elif nTone == 4:
								tn1 = 1
								tn2 = 2
							else:
								tn1 = nTone / 8
								tn2 = nTone - 1 - tn1
						else:
							tn1 = int(cd[1])	# first tone num
							tn2 = int(cd[2])	# last tone num
						if tn1 < 0:
							# negative number means count from back, python style
							tn1 += nTone
						if tn2 < 0:
							# negative number means count from back, python style
							tn2 += nTone
						if tn1 >= 0 and tn1 < nTone and tn2 >= 0 and tn2 < nTone:
							if int(s[toneOffset+4*(rb*nTone+tn1)]) != -1 and int(s[toneOffset+4*(rb*nTone+tn2)]) != -1:
								columns.append([toneOffset+4*(rb*nTone+tn1), toneOffset+4*(rb*nTone+tn2)])
			elif tss[0] == 't':
				# look for specific tone number.  Should be unambiguous
				cd = split(tss[1], ',')	# column identifying data
				if cd[0] == 'ends':
					tn1 = 0
					tn2 = -1
				elif cd[0] == 'good':
					if nTone == 2:
						tn1 = 0
						tn2 = 1
					elif nTone == 4:
						tn1 = 1
						tn2 = 2
					else:
						tn1 = nTone / 8
						tn2 = nTone - 1 - tn1
				else:
					tn1 = int(cd[0])	# first tone num
					tn2 = int(cd[1])	# last tone num
				if tn1 < 0:
					# negative number means count from back, python style
					tn1 += nTone
				if tn2 < 0:
					# negative number means count from back, python style
					tn2 += nTone
				if tn1 >= 0 and tn1 < nTone and tn2 >= 0 and tn2 < nTone:
					for rb in range(nRecBand):
						if int(s[toneOffset + 4*(rb*nTone+tn1)]) != -1 and int(s[toneOffset + 4*(rb*nTone+tn2)]) != -1:
							columns.append([toneOffset+4*(rb*nTone+tn1), toneOffset+4*(rb*nTone+tn2)])
	else:
		print 'Error: pcal file with version %d is not supported' % pcalVersion
		exit(0)

	return columns
	

def calculateComplex(plotDataReal, plotDataImag, row, difxPcalLine, toneColumns):
	s = split(strip(difxPcalLine))
	for i in range(len(toneColumns)):
		x = float(s[toneColumns[i]+2])
		y = float(s[toneColumns[i]+3])
		plotDataReal[i][row] = x
		plotDataImag[i][row] = y

def calculateAmplitudes(plotData, row, difxPcalLine, toneColumns):
	s = split(strip(difxPcalLine))
	plotData[0][row] = float(s[1])
	for i in range(len(toneColumns)):
		x = float(s[toneColumns[i]+2])
		y = float(s[toneColumns[i]+3])
		plotData[i+1][row] = sqrt(x*x+y*y)

def calculatePhases(plotData, row, difxPcalLine, toneColumns):
	s = split(strip(difxPcalLine))
	plotData[0][row] = float(s[1])
	for i in range(len(toneColumns)):
		x = float(s[toneColumns[i]+2])
		y = float(s[toneColumns[i]+3])
		plotData[i+1][row] = atan2(y, x)

def unwrapPhases(phaseData):
	lastPhase = phaseData[0]
	for i in range(1, len(phaseData)):
		while phaseData[i] > lastPhase+pi:
			phaseData[i] -= 2*pi
		while phaseData[i] < lastPhase-pi:
			phaseData[i] += 2*pi
		lastPhase = phaseData[i]

def unwrapDelays(delayData, ambig):
	lastDelay = 0
	for i in range(len(delayData)):
		while delayData[i] > lastDelay + ambig/2:
			delayData[i] -= ambig
		while delayData[i] < lastDelay - ambig/2:
			delayData[i] += ambig
		lastDelay = delayData[i]


# in units of picoseconds
def calculateDelays(plotData, row, difxPcalLine, toneColumns, pcalVersion):
	
	freqColumns = [1, 0]	# indexed by pcalVersion
	
	s = split(strip(difxPcalLine))
	plotData[0][row] = float(s[1])
	
	freqColumn = freqColumns[pcalVersion]

	for i in range(len(toneColumns)):
		x1 = float(s[toneColumns[i][0]+2])
		y1 = float(s[toneColumns[i][0]+3])
		x2 = float(s[toneColumns[i][1]+2])
		y2 = float(s[toneColumns[i][1]+3])
		df = float(s[toneColumns[i][1]+freqColumn])-float(s[toneColumns[i][0]+freqColumn])
		plotData[i+1][row] = 1.0e3*(atan2(y2, x2) - atan2(y1, x1))/(2*pi*df)

def generatePlots(op, toneSelect, dataDict, options):
	ants = dataDict.keys()
	ants.sort()
	colors = ['red', 'orange', 'green', 'blue', 'purple', 'brown', 'cyan', 'magenta']
	for a in ants:
		data = dataDict[a]
		data.sort()

		start = data[0][0]
		stop = data[-1][1]

		title('Pulse cal %s for antenna %s' % (op, a))

		for d in data:
			filename = d[2]
			firstLine = d[3]
			pcalVersion = d[4]
			if op == 'amp':
				toneColumns = getOneToneColumns(toneSelect, firstLine, pcalVersion)
				if options['verbose'] > 0:
					print 'From %s, the following tone columns will be used:' % filename, toneColumns
				nColumn = 1+len(toneColumns)
				fileData = open(filename).readlines()
				nRow = 0
				for d in fileData:
					if d[0] != '#':
						nRow += 1
				plotData = np.zeros((nColumn, nRow))
				row = 0
				for d in fileData:
					if d[0] != '#':
						calculateAmplitudes(plotData, row, d, toneColumns)
						row += 1

				xlabel('Time (Day of Year)')
				ylabel('Amplitude')
				for i in range(1, nColumn):
					if 'lines' in options:
						plot(plotData[0], plotData[i], color=colors[(i-1)%len(colors)], linewidth=1)
					else:
						scatter(plotData[0], plotData[i], color=colors[(i-1)%len(colors)], s=3)
						
			elif op == 'phase':
				toneColumns = getOneToneColumns(toneSelect, firstLine, pcalVersion)
				nColumn = 1+len(toneColumns)
				fileData = open(filename).readlines()
				nRow = 0
				for d in fileData:
					if d[0] != '#':
						nRow += 1
				plotData = np.zeros((nColumn, nRow))
				row = 0
				for d in fileData:
					if d[0] != '#':
						calculatePhases(plotData, row, d, toneColumns)
						row += 1
				for i in range(1, nColumn):
					unwrapPhases(plotData[i])

				xlabel('Time (Day of Year)')
				ylabel('Phase (radians)')
				for i in range(1, nColumn):
					if 'lines' in options:
						plot(plotData[0], plotData[i], color=colors[(i-1)%len(colors)], linewidth=1)
					else:
						scatter(plotData[0], plotData[i], color=colors[(i-1)%len(colors)], s=3)

			elif op == 'delay':
				toneColumns = getTwoToneColumns(toneSelect, firstLine, pcalVersion)
				s = split(firstLine)
				if options['verbose'] > 0:
					print 'From %s, the following tone columns will be used:' % filename, toneColumns
				nColumn = 1+len(toneColumns)
				fileData = open(filename).readlines()
				nRow = 0
				for d in fileData:
					if d[0] != '#':
						nRow += 1
				plotData = np.zeros((nColumn, nRow))
				row = 0
				for d in fileData:
					if d[0] != '#':
						calculateDelays(plotData, row, d, toneColumns, pcalVersion)
						row += 1

				for i in range(nColumn-1):
					freqOffset = [1, 0][pcalVersion]
					ambig = fabs(1.0/(float(s[toneColumns[i][0]+freqOffset])-float(s[toneColumns[i][1]+freqOffset])))
					unwrapDelays(plotData[i+1], ambig*1.0e3)

				xlabel('Time (Day of Year)')
				ylabel('Delay (ns)')
				for i in range(1, nColumn):
					if 'lines' in options:
						plot(plotData[0], plotData[i], color=colors[(i-1)%len(colors)], linewidth=1)
					else:
						scatter(plotData[0], plotData[i], color=colors[(i-1)%len(colors)], s=3)

			elif op == 'xy':
				toneColumns = getOneToneColumns(toneSelect, firstLine, pcalVersion)
				nColumn = len(toneColumns)
				fileData = open(filename).readlines()
				nRow = 0
				for d in fileData:
					if d[0] != '#':
						nRow += 1
				plotDataX = np.zeros((nColumn, nRow))
				plotDataY = np.zeros((nColumn, nRow))
				row = 0
				for d in fileData:
					if d[0] != '#':
						calculateComplex(plotDataX, plotDataY, row, d, toneColumns)
						row += 1

				xlabel('Real')
				ylabel('Imag')
				axes().set_aspect('equal')
				for i in range(nColumn):
					if 'lines' in options:
						plot(plotDataX[i], plotDataY[i], color=colors[i%len(colors)], linewidth=1)
					else:
						scatter(plotDataX[i], plotDataY[i], color=colors[i%len(colors)], s=3)
				

		savefig('%s-pcal-%s.pdf' % (a, op), bbox_inches='tight')
		show()


# main()

op, toneSelect, data, options = parseCommandLine()
dataDict = generateDataDictionary(data)

if options['verbose'] > 0:
	summarizeData(dataDict)

generatePlots(op, toneSelect, dataDict, options)

# vim: tabstop=8:softtabstop=8:shiftwidth=8:noexpandtab