#!/usr/bin/python

import re, pylab, numpy, math

#Properties of simulation data
injOutName = 'rad_ccuw1.png'
injOutName2 = 'rad_ccuw1_2.png'
filename = 'combined.dat'
injbottom = 0.0
injcutoff = 50.0
injTitle = 'Particle lost at CCU_W'

Qx0 = 95
Qx1 = 135
dQx = 0.5
Qy0 = 75
Qy1 = 115
dQy = 0.5

#Properties of tune plane
windXmin = Qx0
windXmax = Qx1
windYmin = Qy0
windYmax = Qy1
pMax = 4  # [pQ_x+qQ_y+rQ_s = n]
qMax = 2
rMax = 2
nMax = 2
pqrMax = 8
pqMax = 6
s = 0.0464

#Properties of plot
plotXmin = Qx0
plotXmax = Qx1
plotYmin = Qy0
plotYmax = Qy1

#operation point:
posQx=116.1 #positron
posQy=299.0
posQx2=99.5 #electron
posQy2=315.0

frev=390.136

def additionalItemsToPlot():
	pylab.plot(posQx, posQy,'go',markersize=10)
	pylab.plot(posQx2, posQy2,'rd',markersize=10)
	pylab.plot([windXmin,windXmax],[windYmin,windYmin], 'k-')
	pylab.plot([windXmin,windXmax],[windYmax,windYmax], 'k-')
	pylab.plot([windXmin,windXmin],[windYmin,windYmax], 'k-')
	pylab.plot([windXmax,windXmax],[windYmin,windYmax], 'k-')

def additionalItemsToPlot2():
	pylab.plot(posQx/frev, posQy/frev,'go',markersize=10)
	pylab.plot(posQx2/frev, posQy2/frev,'rd',markersize=10)
	pylab.plot([windXmin,windXmax],[windYmin,windYmin], 'k-')
	pylab.plot([windXmin,windXmax],[windYmax,windYmax], 'k-')
	pylab.plot([windXmin,windXmin],[windYmin,windYmax], 'k-')
	pylab.plot([windXmax,windXmax],[windYmin,windYmax], 'k-')

#-------------------------------
#- Begin Main Program
#-------------------------------

#-------------------------
#- Read in simulation data
#-------------------------
xsteps = int(round( (Qx1-Qx0)/dQx, 0 )) + 1
ysteps = int(round( (Qy1-Qy0)/dQy, 0 )) + 1

xgridmin = Qx0 - dQx/2.0
xgridmax = Qx0 - dQx/2.0 + xsteps*dQx + 0.00001
lqx = numpy.arange(xgridmin, xgridmax, dQx)

ygridmin = Qy0 - dQy/2.0
ygridmax = Qy0 - dQy/2.0 + ysteps*dQy + 0.00001
lqy = numpy.arange(ygridmin, ygridmax, dQy)

Qx, Qy = numpy.meshgrid(lqx,lqy)

amph = numpy.zeros(xsteps*ysteps).reshape(ysteps,xsteps)

# autofill: any missing data points will be taken at cutoff
for ix in range(xsteps):
	for iy in range(ysteps):
		amph[iy,ix] = injcutoff

file = open(filename, 'r')
lines = file.readlines()[1:]
for line in lines:
	fileLine = line.split()
	if len(fileLine) == 0 or fileLine[0][0] == '!': continue
	Datax = float(fileLine[4])
	Datay = float(fileLine[5])
	ah = float(fileLine[7]) #ccu_w1 (south one, close to 7w)
#	ah = float(fileLine[8]) #ccu_w2
	steps = float(fileLine[0])
	init_orb = float(fileLine[1])

	if ( (ah > injcutoff) ): ah = injcutoff

	Dataix = int(round((Datax - Qx0) / dQx, 0))
	Dataiy = int(round((Datay - Qy0) / dQy, 0))
	amph[Dataiy,Dataix] = ah

file.close()

#--------------------------------------------
#Plot the data
#--------------------------------------------

#Plot in kHz 
pylab.clf()
pylab.gca().set_autoscale_on(False)
pylab.axes().set_aspect('equal')

#pylab.pcolor(Qx,Qy,amph,rasterized=True)
pylab.pcolor(Qx,Qy,amph,rasterized=False)

pylab.xlim(plotXmin,plotXmax)
pylab.ylim(plotYmin,plotYmax)
pylab.clim(injbottom,injcutoff)

pylab.colorbar()
pylab.xlabel('Qx (kHz)')
pylab.ylabel('Qy (kHz)')

#additionalItemsToPlot()  #Plot additional items specified at top of this file

pylab.title(injTitle)
pylab.savefig(injOutName, dpi=150)
#pylab.savefig(injOutName, format='eps')


#------------------------
#Plot in fractional tunes
pylab.clf()
pylab.gca().set_autoscale_on(False)
pylab.axes().set_aspect('equal')

#pylab.pcolor(Qx,Qy,amph,rasterized=True)
pylab.pcolor(Qx/frev,Qy/frev,amph,rasterized=False)


pylab.xlim(plotXmin/frev,plotXmax/frev)
pylab.ylim (plotYmin/frev,plotYmax/frev)
pylab.clim(injbottom,injcutoff)

pylab.colorbar()
pylab.xlabel('Qx')
pylab.ylabel('Qy')
#pylab.grid(True)

#additionalItemsToPlot2()  #Plot additional items specified at top of this file

pylab.title(injTitle)
pylab.savefig(injOutName2, dpi=150)
#pylab.savefig(injOutName2, format='eps')