#!/usr/bin/env python
# v0.2 13 October 2008
import time,sys,re
from pylab import *  #pylab is the same as matplotlib.pylab
from mycolors import *
#############################################
from optparse import OptionParser
usage="usage: %prog DATAFILE   [options]"
parser = OptionParser(usage=usage, version="%prog 0.1")
parser.add_option("-d", dest="deltas",type="string",help="contour increments",default='.1 .1 .1 .1')
parser.add_option("-B", dest="baren",action='store_true',help="don't plot headers",default=False)
parser.add_option("-x", dest="xlim",type="float",help="plot window x limit",default=.5)
parser.add_option("-y", dest="ylim",type="float",help="plot window y limit",default=1.)
parser.add_option("-o", dest="outputfile",type="string",help="output filename",default=None)
parser.add_option("-c", dest="colorlegend",action='store_true',help="add color legend",default=False)
parser.add_option("-l", dest="drawl",action='store_true',help="plot contour lines",default=False)
parser.add_option("-e", dest="extendl",action='store_true',help="use extended color legend",default=False)
parser.add_option("-v", dest="varstring",type="string",help="variables to plot",default='u v w p')
parser.add_option("--dpi", dest="dpi",type="int",help="dotsperinch",default=72)
parser.add_option("-n", dest="mlev",type="int",help="number of contour levels",default=20)
parser.add_option("-b", dest="black",action='store_true',help="reverse foreground & background colors",default=False)
parser.add_option("-g", dest="drawg",action='store_true',help="draw grid",default=False)
parser.add_option("-C", dest="colorname",type="string",help="color scheme",default="default")
parser.add_option("-V", dest="verbose",action='store_true',help="debugging info",default=False)
(opt, args) = parser.parse_args()
print(opt)
print(args)
xlim=opt.xlim
ylim=opt.ylim
dotsperinch=opt.dpi
extendl=opt.extendl
mlev=opt.mlev
deltas=[float(x) for x in opt.deltas.split()]
tcol='k'
fcol='w'
if opt.black:
	tcol='w'
	fcol='k'
try:
	filename=args[0]
except:
	filename="eza.dat"
#################################
def drawgrid(theaxes):
	im,jm= xa.shape
	segs=[]
	regs=[]
	grcol=[0,0,0]
	for j in arange(0,jm):
		xl=xa[:,j]
		yl=ya[:,j]
		myline= Line2D(xl,yl,color=grcol,linewidth=1)
		myline.set_aa(False)
		theaxes.add_line(myline)
	for i in arange(0,im):
		xl=xa[i,:]
		yl=ya[i,:]
		myline= Line2D(xl,yl,color=grcol,linewidth=1)
		myline.set_aa(False)
		theaxes.add_line(myline)
#################################
def drawframe(theaxes):
	grcol=[0,0,0]
	for x in [0,xlim]:
		xl=[x,x]
		yl=[0,ylim]
		myline= Line2D(xl,yl,color=grcol,linewidth=1)
		myline.set_aa(False)
		theaxes.add_line(myline)
	for y in [0,ylim]:
		xl=[0,xlim]
		yl=[y,y]
		myline= Line2D(xl,yl,color=grcol,linewidth=1)
		myline.set_aa(False)
		theaxes.add_line(myline)
#################################
def contourit(theaxes,za,xlim,ylim,lab):
	global mylevs
	mylevs=makelevels()
	if extendl: mylevs=extendlevels(mylevs)
	print("start contourf and contour",time.process_time())
	contourf(xa,ya,za,mylevs+[1.e30],colors=mycol,antialiased=False,zorder=3)
	if opt.drawl: contour(xa,ya,za,mylevs,colors=('k',))
	print ("end contourf and contour",time.process_time())
	print ("LAB",lab)
	texttop(za,lab,theaxes)
	textbottom(delta,lab,theaxes)
	if not opt.baren:
		textbottom(delta,lab,theaxes)
	theaxes.set_xlim([0,xlim])
	theaxes.set_ylim([0,ylim])
#################################
def extendcolors(mycol):
	morecols=[(0,1.,0),(0.5,1.,0.5),(0.9,1,0.9)]
	newcol=mycol+morecols
	morecols.reverse()
	newcol=morecols+newcol
	return newcol
#################################
def sliceit(i,r,sz):
	j=i+sz[0]*sz[1]
	slice=r[i:j].copy()
	slice.resize(sz)
	return slice,j
#################################
from  matplotlib import  rc
import matplotlib.lines as mpllines
#rc("tick", labelsize=18) # older versions use just tick
rc("xtick", labelsize=18)
rc("ytick", labelsize=18)
#######################################
############# START MAIN PROGRAM ##################################
inf=open(filename,'rb') # 'b' is needed only for Windows
a=inf.read()
namelistfile=re.sub(r'\/.*dat','/eza.input',filename)
print(namelistfile)
nmlf=open(namelistfile,'r')
#r=fromstring(a[4:-4],typecode='d') # for older Numeric package
#Fortran has 4 extra bytes at beginning and end of files, so strip off:
r=fromstring(a[4:-4],dtype='d')
a=''
i=0
j=100
ap=r[0:j].copy()
namelistlines=nmlf.readlines()
nml={}
n=0
for line in namelistlines:
	v=line.split('=')
	if len(v)<2: continue
	name=v[0].strip()
	print(name)
	nml[name]=ap[n]
	n+=1
print(nml)
#n=1

print("the ap array is:\n",ap)
timeap=ap[0]
useprim=int(nml['useprim'])
fortran_in=int(ap[98])
fortran_jn=int(ap[99])
print('model was run with in=',fortran_in,' jn=',fortran_jn)
sz=(fortran_jn,fortran_in) #note python Numpy follows C (not FORTRAN) convention for array storage
i=j
b,i=sliceit(i,r,sz)
bdf,i=sliceit(i,r,sz)
#bdf=bdf+.4*b #temporary
stm,i=sliceit(i,r,sz)
vor,i=sliceit(i,r,sz)
u,i=sliceit(i,r,sz)
v,i=sliceit(i,r,sz)
w,i=sliceit(i,r,sz)
viscu,i=sliceit(i,r,sz)
j=i+sz[1]
xa=resize(r[i:j],sz)
i=j
j=i+sz[0]
ytemp=resize(r[i:j].copy(),(sz[1],sz[0]))
ya=transpose(ytemp)
########################################
if nml['useprim']>0:
	p=vor
	#p=vor-vor[-1,0] #temporary
	modeltype='prim'
else:
	p=stm
	modeltype='stvr'
########################################
if opt.verbose: #print out debugging info
#examples of debugging info:
	print(dir(b))
	print(shape(b),b.iscontiguous())
	print(max(b.flat))
	print(w.shape)
	print(w.max)
	print(w.min) 
####################################
###################################
guiheight=8.
nlev=mlev+1
mycol=[niceColors(n/(nlev-1.0),colorname=opt.colorname) for n in range(0,nlev)]
if extendl: mycol=extendcolors(mycol)
def minmax(f,xa,ya,xlim,ylim):
	im,jm=f.shape
	maxf=f[0,0]
	minf=f[0,0]
	for i in range(0,im):
		for j in range(0,jm):
			if xa[i,j]>xlim: continue
			if ya[i,j]>ylim: continue
			if f[i,j]>maxf: maxf=f[i,j]
			if f[i,j]<minf: minf=f[i,j]
	return [minf,maxf]
def extendlevels(mylevs):
	delta=mylevs[2]-mylevs[1]
	cmax=mylevs[-1]+delta/2.
	morelevs=[1.5*cmax,2.5*cmax,4*cmax]
	newlevs=mylevs+morelevs
	morelevs=[-1.5*cmax,-2.5*cmax,-4*cmax]
	morelevs.reverse()
	newlevs=newlevs[0:1]+morelevs+newlevs[1:]
	return newlevs
def makelevels():
	global delta,deltas
	delta=deltas.pop(0)
	deltas.append(delta)
	mylevs=[-(mlev/2-.5)*delta+delta*n for n in range(0,mlev)]
	mylevs.insert(0,-1.e30)
	return mylevs
def texttop(f,lab,theaxes):
	fmin,fmax=minmax(f,xa,ya,xlim,ylim)
	fex=max(abs(fmin),abs(fmax))
	format=r'%6.2f'
	if fex<.1: format=r'%6.3f'
	leftstring= format+' <' 
	leftstring= leftstring % (fmin)
	rightstring= '< '+format % (fmax)
	topstring= r'$%s$' % (lab) 
	#topstring= r'$\psi$' # LaTeX demo 
	theaxes.text(.0*xlim,1.02*ylim  ,leftstring ,fontsize=18,
			verticalalignment='bottom',horizontalalignment='left',color=tcol)
	theaxes.text(xlim,1.02*ylim  ,rightstring ,fontsize=18,
			verticalalignment='bottom',horizontalalignment='right',color=tcol)
	theaxes.text(.5*xlim,1.01*ylim  ,topstring ,fontsize=36,
			verticalalignment='bottom',horizontalalignment='center',color=tcol)
def textbottom(d,lab,theaxes):
	if not opt.baren:
		thestring= '$\Delta %s=%5.2f$' % (lab,d)
		theaxes.text(.1*xlim,-.06*ylim  ,thestring ,fontsize=24,verticalalignment='top',color=tcol)
	theaxes.text(.5*xlim,-.01*ylim  ,'$x$' ,fontsize=24,verticalalignment='top',color=tcol)
	theaxes.text(-.01*xlim,.5*ylim  ,'$z$' ,fontsize=24,horizontalalignment='right',color=tcol)
def setsubplot(n):
	global sa
	sa=axes([(dw+n*(dpw+din))/guiasp,dw,dpw/guiasp,pt],zorder=1)
	#sa=axes([(dw+n*(dpw+din)),dw,dpw,pt],zorder=1)
	sa.set_xticks([0,xlim])
	sa.set_yticks([0,ylim])
	lines =sa.get_xticklines()
	for line in lines:
		line.set_marker(mpllines.TICKDOWN)
		line.set_color(tcol)
	lines =sa.get_yticklines()
	for line in lines:
		line.set_marker(mpllines.TICKLEFT)
		line.set_color(tcol)
	labs=sa.get_xticklabels()
	for lab in labs:
		lab.set_color(tcol)
	labs=sa.get_yticklabels()
	for lab in labs:
		lab.set_color(tcol)
vv=opt.varstring.split()
command="varstoplot=["+','.join(vv)+"]"
print(command)
exec(command)
#######################
if True:  
	pt=.7 #sub-plot height
	plotasp=1. #means isotropic distances 
	dpw=pt*xlim/ylim
	din=.07
	pw=plotasp*(dpw+din)*len(vv)
	dw=.1
	cw=0.
	if opt.colorlegend: cw=.1
	guiwidth=(2*dw+pw+2*cw)*guiheight
	guiasp=guiwidth/guiheight
	print(guiwidth,guiheight)
	figure(figsize=(guiwidth,guiheight),facecolor=fcol)
	nframe=0
	for f in varstoplot:
		print("will call setsubplot")
		setsubplot(nframe)
		thelab=vv[nframe]
		if useprim==0 and thelab=='p': thelab="\psi"
		contourit(sa,f,xlim,ylim,thelab)
		drawframe(sa)
		if nframe==0:
			params1="time=%5.2f   rei=%9.2e   reit=%9.2e   swirl=%9.2e" %   \
					(nml['time'], nml['rei'], nml['reit'], nml['swirl'])
			params2="rwid=%3.1f   hite=%3.1f   order=%d   in=%d   jn=%d   %s"  % \
					(nml['rwid'], nml['hite'], int(nml['order']), fortran_in, fortran_jn, modeltype)
			if not opt.baren:
				sa.text(0,1.17*ylim,params1,fontsize=18,color=tcol)
				sa.text(0,1.10*ylim,params2,fontsize=18,color=tcol)
			if opt.drawg: drawgrid(sa)
		nframe+=1
else:
	print("not sure what this is for")	
#######################
if opt.colorlegend:
	sleg=axes([(2*dw+pw)/guiasp,dw,cw/guiasp,pt])   
	sleg.set_axis_off()
	labels=['%5.2f' % x for x in mylevs[1:]]
	nlab=len(labels)
	dy=1./(nlab+0)
	for i in range(0,len(labels)):
		regcolor=mycol[i]
		xrect=[0,1,1,0]
		ybot=max(0,(i-.5)*dy)
		ytop=min(1,(i+.5)*dy)
		yrect=[ybot,ybot,ytop,ytop]
		sleg.fill(xrect,yrect,facecolor=regcolor,edgecolor=regcolor,linewidth=0) 
		sleg.text(1.1, (i+.5)*dy   , labels[i],verticalalignment='center',color=tcol)
	regcolor=mycol[-1]
	yrect=[ytop,ytop,ytop+.5*dy,ytop+.5*dy]
	sleg.fill(xrect,yrect,facecolor=regcolor,edgecolor=regcolor,linewidth=0)
	for i in range(0,len(labels)):
		segcolor='k'
		ytop=min(1,(i+.5)*dy)
		myline= Line2D([0,1],[ytop,ytop],color=segcolor,linewidth=2)
		if mylevs[i+1]<0: myline.set_dashes([5,5])
		sleg.add_line(myline)
if opt.outputfile!=None:
	savefig(opt.outputfile, dpi=dotsperinch, facecolor='w', edgecolor='w', orientation='portrait')
else:
	show()
