#!/usr/bin/python

#from katcp import BlockingClient, Message
'''PLOTS CROSS AND AUTO CORRELATION FUNCTION OF WIDEBAND POCO n-CH DESIGN ON ROACH BOARD
This program is a generalized program with few modifications to plot and dump the data in
analysis program-tax native format into the given specified format.
'''
import matplotlib
from matplotlib import rcParams
matplotlib.use('GTKAgg')
import katcp, numpy, pylab
import gtk, gobject, time, struct, sys, math

#device_host 	= "192.168.100.72"	# For Lab Testing
#device_host 	= "192.168.4.91"	# For Receiver Room Testing
device_port 	= 7147
acq_time	= 2**27/(150.0*(10**6))		# 0.89 is 1 sync period time
n_chans		= 1024
ants		= 2
max_ants	= 2
baselines	= (ants*(ants-1))/2
bram_offset	= ants*2
e_o_offset	= 2
n_brams		= 4 #( even + odd + re + im )

#print "baselines = %i\tbaselines+ants = %i"%(baselines,baselines+ants)

#print 'Connecting to ROACH through katcp'
#client =  katcp.BlockingClient(device_host, device_port)
#client.start()

if __name__ == '__main__':
        from optparse import OptionParser
        p = OptionParser()
        p.set_usage('tut4_poco_plot.py <ROACH_HOSTNAME_or_IP> [options] ')
        p.set_description(__doc__)
        p.add_option('-l', '--log', dest='log', action='store_true', default=False,
            help='''Plot the power in logarithmic scale (requires some non-zero value signal).''')
        p.add_option('--hold', dest='hold', action='store_true', default=False,
            help='''Turn on hold. This will plot subsequent specra on top of each other.''')
        p.add_option('-f', '--file', action='store', type='string', dest='filename',
            help='''file dump the data in native tax format''', metavar='FILE')
        opts, args = p.parse_args(sys.argv[1:])

        if args==[]:
            print 'Please specify a ROACH board. \nExiting.'
            exit()
        else:
            roach = args[0]

client =  katcp.BlockingClient(roach, device_port)
client.start()

print "baselines = %i\tbaselines+ants = %i"%(baselines,baselines+ants)

bram_map=['dir_x1_aa_real','dir_x1_bb_real','dir_x1_cc_real','dir_x1_dd_real','dir_x1_ac_real','dir_x1_ac_imag','dir_x1_bd_real','dir_x1_bd_imag']	#BRAM for 2-Ant Correlator System.
#bram_map=['dir_x1_aa_real','dir_x1_bb_real','dir_x1_cc_real','dir_x1_dd_real',]

x=range(0,n_chans/2)
xlab=('FFT channel number')

fig = pylab.figure(1)
manager = pylab.get_current_fig_manager()
cnt = 0


#        (opts, args) = p.parse_args()

if (opts.filename != None):
        f_name = opts.filename #+ opts.filename
        print "Tax File is %s" %f_name
        fptr = open (f_name,'wb')


# CONTINUOUS LOOP
def updatefig(*args):
    	global cnt, ants
    	cnt += 1
    	print 'Grabbing spectrum number %i'%cnt

	rply		= [[],[],[],[],[],[],[],[],[]]
	self		= [[],[]]
	self_tax	= [[],[]]
	cross           = [[],[]]
	cross_tax	= [[],[],[],[]]
	cx_tax		= [[],[],[],[]]
        cx_intrleav     = []
	
	main_loop_t	= time.time()
# DATA GRABBING UTILITY
	for input in range(2*ants+baselines*4):
		reply, inform 	= client.blocking_request(katcp.Message.request("wordread", bram_map[input], "0", "256"))
		rply[input]     = map(eval, reply.arguments[1:n_chans/4+1])
		print"rply[%i] = %i"%(input,numpy.size(rply[input]))	
#		print 'size of rply[%i] = %i'%(input, numpy.size(rply[input]))
		if input > (bram_offset-1):
			print "bram_map[%i]=%s"%(input,bram_map[input])
			for i in range(256):                            # For 2's Complement interpretation :)
        	                if (rply[input][i] > 2147483647):
 	              	                rply[input][i] = rply[input][i] - 4294967296

	k=0 #self[0] =>input[0] & input[1] and self[1] => input[2] & input[3]
	#--------------- SELF DATA MANIPULATION--------------------#
	for input in range(ants):
		rply_1		= list(numpy.copy(rply[k+input]))
		rply_2		= list(numpy.copy(rply[1+k+input]))
        	j		= 0
		print 'size of rply_1[%i] = %i, rply_2[%i] = %i'%((k+input),numpy.size(rply_1),(k+input+1),numpy.size(rply_2))
	        for i in range(1,(n_chans/2+1),2):                       		# This is for interleaving Logic.
       		        rply_1.insert(i,rply_2[j])
               		j=j+1
	        k=k+1
#		print 'After size of rply_1 = %i, rply_2 = %i'%(numpy.size(rply_1), numpy.size(rply_2))
       	 	self[input] 	= list(numpy.array(rply_1) + 0.1)
#			print input
#			print 'size of self[%i] = %i'%(input, numpy.size(self[input]))
	#----------- for tax program compatibility --------------#		
#			self_tmp 	= list(numpy.array(rply_1)) 
		self_tmp 	= list(numpy.copy(self[input]))
#		print 'size of self[%i] = %i, self_tmp = %i'%(input, numpy.size(self[input]), numpy.size(self_tmp))
		for i in range(1,((n_chans)+1),2):                         # This is for interleaving Logic Depends upon your n_chans value.
                       	self_tmp.insert(i,0)
		self_tax[input]	= list(numpy.copy(self_tmp))
#			print 'After size of self[%i] = %i, self_tmp = %i, self_tax[%i] = %i'%(input, numpy.size(self[input]), numpy.size(self_tmp), input, numpy.size(self_tax[input]))

	#------------FOR CROSS DATA MANIPULATION-------------#
	for input in range(baselines):										# This Number depends upon number of Cross Spectrums to be observed
		#-------------- Bin Interleaving-----------------#
		for m in range(n_brams/2):								# Real Interleaving and then Imaginary Interleaving 
			print "input = %i\tm = %i\t cx_tmp[%i]\tcx_tmp1[%i]"%(input,m,(bram_offset+input*4+m),(bram_offset+input*4+m+e_o_offset))    # Is logic me gadbad hui thi....... rply number me. :( {Added for Generalized prog}
			cx_tmp  = list(numpy.copy(numpy.array(rply[bram_offset+input*4+m])))          	# Odd Samples    	{for Generalized} 
                        cx_tmp1 = list(numpy.copy(numpy.array(rply[bram_offset+input*4+m+e_o_offset])))      # Even Samples		{for Generalized}
			j 	= 0
			for i in range(1,((n_chans/2)+1),2):                         	# This is for interleaving Logic Depends upon your n_chans value.
				cx_tmp.insert(i,cx_tmp1[j])			     	# insert before ith position
			        j = j+1
			cx_tax[m] = list(numpy.copy(cx_tmp))				# Even and Odd Samps Interleaved and cx_tax[0] = Real cx_tax[1] = Imaginary
			print 'size of cx_tax[%i] = %i'%(m, numpy.size(cx_tax[m]))
		#----------------- For Python Plotting Program --------------------#
		cross[input] = numpy.copy(numpy.array(cx_tax[input*2]) + numpy.array(cx_tax[input*2+1])*1j +0.1)
		#-------------- Real & Imaginary Interleaving-----------------#
		for n in range(baselines):		
                        cross_1 = list(numpy.copy(cx_tax[n*2]))                       # Real Part 
                        cross_2 = list(numpy.copy(cx_tax[n*2+1]))                     # Imaginary Part 
			print 'size of cx_tax[%i] = %i\tSize of cx_tax[%i] = %i\n'%(n*2,(numpy.size(cross_1)),(n*2+1),(numpy.size(cross_2)))
			j = 0
			for i in range(0,n_chans,2):                         	# This is for interleaving Logic Depends upon your n_chans value.
				cross_2.insert(i,cross_1[j])
				j = j+1
			print 'size of cross_2 = %i\n'%(numpy.size(cross_2))
			cross_tax[n] = list(numpy.copy(cross_2))

#----------------- DATA DUMPING UTILITY ------------------------#
	#--------------File Write-----------------#	
	if (opts.filename != None):
		a 	 = 0
		cx_index = 0
		ants_var = ants
		while(ants_var >= 0):
			for j in range(ants_var):
				if j == 0:
					fptr.write(numpy.array(self_tax[a], dtype = 'f'))
					print 'size of self_tax[%i] = %i'%(a, numpy.size(self_tax[a]))
				else:
					fptr.write(numpy.array(cross_tax[(cx_index+j-1)], dtype = 'f'))	
					print 'size of cross_tax[%i] = %i'%((cx_index+j-1), numpy.size(cross_tax[(cx_index+j-1)]))
			a = a+1	
			ants_var = ants_var-1		
			cx_index = (cx_index + ants_var)
	#------------PLOTTING UTILITY-------------#
	k = 0
	for i in range(ants):			# For 2-Ants : Selfs n = 2. Will be n=4 for 4-Ants
		num = (baselines*2+ants)*100 + 11 + k
		print 'plot number (num)= %i'%num
		pylab.subplot(11+(baselines*2+ants)*100+k)	# For 2-Ants : n*100 = n-self plots	
		pylab.ioff()
		pylab.hold(False)
#		print 'size of self[%i] = %i'%(i,numpy.size(self[i]))
		if opts.log:
			pylab.semilogy(x, self[i])
       			pylab.ylabel('Power (log, arb units)')
		else:
			pylab.plot(x, self[i])
       			pylab.ylabel('Power (linear, arb units)')
		pylab.xlim(x[2],x[n_chans/2-1])
		pylab.grid()
	        pylab.title('Power: Self-%i Spectrum number %i'%(i,cnt))
		k = k+1
	for i in range(baselines):
		#------------Cross-Amplitude--------------#	
		num = (baselines*2+ants)*100 + 11 + k
		print 'plot number (num)= %i'%num
		pylab.subplot((11+(baselines*2+ants)*100)+k)   # For 2-Ants : n*100 = n-self plots     
		pylab.ioff()
		pylab.hold(opts.hold)
		if opts.log:
	                pylab.semilogy(x, (numpy.abs(cross[i])))
                	pylab.ylabel('Power (log, arb units)')
		else:
			pylab.plot(x, (numpy.abs(cross[i])))
                	pylab.ylabel('Power (linear, arb units)')
#                pylab.semilogy(x, (numpy.abs(cross_tax[i])))
#                pylab.plot(cross_1)
                pylab.xlim(x[2],x[n_chans/2-1])
                pylab.grid()
                pylab.title('Power: Cross-Amplitude %i. Spectrum number %i'%(i,cnt))
		k = k+1
		#------------Cross-Phase--------------#	
		num = (baselines*2+ants)*100 + 11 + k
		print 'plot number (num)= %i'%num
		pylab.subplot((11+(baselines*2+ants)*100)+k)   # For 2-Ants : n*100 = n-self plots     
		pylab.ioff()
		pylab.hold(opts.hold)
                pylab.plot(numpy.angle(cross[i],deg=True))
#                pylab.plot(cx_tax[1])
                pylab.xlim(x[2],x[n_chans/2-1])
#                pylab.ylim(-180,+180)
                pylab.ylabel('Phase in Deg')
                pylab.grid()
                pylab.title('Power: Cross-Phase %i. Spectrum number %i'%(i,cnt))
		k = k+1
	pylab.xlabel(xlab)
#	pylab.semilogy(x, self[1])
#       	pylab.ylabel('Power (log scale, arbitrary units)')
#	pylab.grid()
#	k = k+1
#	pylab.subplot(413)
#	pylab.semilogy(x, numpy.abs(cross[0]))
#       	pylab.ylabel('Cross Power (arbitrary units)')
#	pylab.grid()
  #      pylab.title('Power: Input %i. Spectrum number %i'%(input,cnt))
#	k = k+1
#	pylab.subplot(414)
#	pylab.semilogy(x, numpy.angle(cross[0],deg=True))
#	k = k+1
#	pylab.grid()
#	pylab.setp(ax[input].get_xticklabels(),visible=False)
#       	pylab.ylabel('Cross Angle (in Deg.)')
#        pylab.title('Power: Input %i. Spectrum number %i'%(input,cnt))
#	for i in range(2):
#        	if opts.log:
       	#ax[i].set_ylim(3.5,ymax*1.001)
#	        else:
#        	        ax[i].set_ylim(0,ymax*1.1)

	#pylab.setp(ax[input].get_xticklabels(),visible=True)
	loop_1_t	= time.time()
	time.sleep(acq_time-(loop_1_t-main_loop_t))
	print "Loop time = %f \nSleep = %f\n"%((loop_1_t-main_loop_t), acq_time-(loop_1_t-main_loop_t))
	pylab.xlabel(xlab)
	pylab.draw()
#	if (cnt == (8*60*60+3168)):
#		return False
#	else:
#		return True
	return True
#	return False	# For single dump
gobject.idle_add(updatefig)
pylab.show()
client.stop()
client.join()
if(opts.filename != None):
        fptr.close()
print 'Done with all'