#mdot


def plot_infall_vel(xaxis,ax, path_exp,path_wind,label,shape,marker,markevery,color):
	'''
	Plots the infall velocity vs free-fall time or mass coordinate

	Parameter
	xaxis: 'ff': free-fall time
	        'mass' : mass coordinate
	

	'''

	plt.figure(fig)
	import nugridse as mp
	sefiles=mp.se(path_exp)
	cyc = int(sefiles.se.cycles[0])
	cyc = cyc - 1
	sefiles=mp.se(path_wind)



	G =6.67259e-8 #cm^3 g^-1 s^-2 
	msun = 1.9891e33 #Msun in g

	r = sefiles.get(cyc,'radius') #cm
	rho = sefiles.get(cyc,'rho') #g/cm3
	m = sefiles.get(cyc,'mass')*msun # cm
	print 'cycle lookup',cyc

	r = np.array(r)

	t_ff = np.pi/2 * r**(3./2.)/np.sqrt(2*G*m) #s

	v_ff = np.sqrt(2*G*m/r)

	mdot = 4.*np.pi*r**2 * rho * v_ff /msun
	print 'plot'
	if xaxis=='ff':
		ax.plot(t_ff,mdot,label=label,linestyle=shape,markevery=markevery,marker=marker,color=color)
		plt.xlabel(r't$_{ff}$ [s]')
		plt.xscale('log')

	if xaxis=='mass':
		ax.plot(m/msun,mdot,label=label,linestyle=shape,markevery=markevery,marker=marker,color=color)
		plt.xlabel("M [M$_{\odot}$]")

	plt.ylabel(r"$\dot{M}$ [M$_{\odot}$/s]")
	plt.yscale('log')



fign = 1
fig=fign
markevery=1000
props = dict(boxstyle='square', facecolor='w', alpha=1)
f, axarr = plt.subplots(2, 2)
metals=['Z=0.02','Z=0.01','Z=0.001','Z=0.0001']
marker=['o','^','s','D']
shape=['-','--',':','-.']
color=['r','b','g','k']

xaxis='ff'
#xaxis='mass'



props=dict(boxstyle='square',facecolor='w',alpha=1)

path='/apod2/NuGrid/data/set1ext'
textr=0.25
log=True
for i in range(len(metals)):

	if i==0:
		ax=plt.sca(axarr[0,0])
		axarr[0,0].text(textr, 0.94, metals[i], transform=axarr[0,0].transAxes, fontsize=16,verticalalignment='top', bbox=props)	
		path_s=path+'/set1.2'
		k=0
		for m in [12.,15.,20.,25.]:
			try:
				path_exp1=path_s+'/see_exp/M'+str(m)+'Z2.0e-02.delay'
				path_wind1=path_s+'/see_wind/M'+str(m)+'Z2.0e-02/M'+str(m)+'Z2.0e-02'
				label=str(m)+'$M_{\odot}$' #, Z='+str(metals[i])
				plot_infall_vel(xaxis,axarr[0,0], path_exp1,path_wind1,label=label,shape=shape[k],marker=marker[k],markevery=markevery,color=color[k])
				k=k+1
			except:
				print '#######################################################################'
				print 'faulty data ',label
				print '#######################################################################'

				pass
		plt.legend(loc=1,fontsize=14)
		#axarr[0,0].locator_params(nbins=4, axis='x')
		plt.ylim(1e-8,1e2)		
                axarr[0,0].set_xlabel('')
                axarr[0,0].set_xticklabels([])
		plt.xlim(1e-2,7)	
		plt.ylim(0.06,1e2)		
	
		
	if i==1:
		ax=plt.sca(axarr[0,1])
		axarr[0,1].text(textr, 0.94, metals[i], transform=axarr[0,1].transAxes, fontsize=16,verticalalignment='top', bbox=props)	
		path_s=path+'/set1.1'
		k=0
		for m in [12.,15.,20.,25.]:
			try:
				path_exp1=path_s+'/see_exp/M'+str(m)+'Z1.0e-02.delay'
				path_wind1=path_s+'/see_wind/M'+str(m)+'Z1.0e-02/M'+str(m)+'Z1.0e-02'
				label=str(m)+'$M_{\odot}$' #, Z='+str(metals[i])
				plot_infall_vel(xaxis,axarr[0,1], path_exp1,path_wind1,label=label,shape=shape[k],marker=marker[k],markevery=markevery,color=color[k])
				k=k+1
			except:
				print 'faulty data ',label
				pass
		plt.legend(loc=1,fontsize=14)
		#axarr[0,0].locator_params(nbins=4, axis='x')
		plt.ylim(1e-8,1e2)		
                axarr[0,1].set_xlabel('')
                axarr[0,1].set_xticklabels([])
                plt.legend().set_visible(False)
		plt.xlim(1e-2,7)	
		plt.ylim(0.06,1e2)		
	
		
	if i==2:
		ax=plt.sca(axarr[1,0])
		axarr[1,0].text(textr, 0.94, metals[i], transform=axarr[1,0].transAxes, fontsize=16,verticalalignment='top', bbox=props)	
		path_s=path+'/set1.4a'
		k=0
		for m in [12.,15.,20.,25.]:
			try:
				path_exp1=path_s+'/see_exp/M'+str(m)+'Z1.0e-03.delay'
				path_wind1=path_s+'/see_wind/M'+str(m)+'Z1.0e-03/M'+str(m)+'Z1.0e-03'
				label=str(m)+'$M_{\odot}$' #, Z='+str(metals[i])
				plot_infall_vel(xaxis,axarr[1,0], path_exp1,path_wind1,label=label,shape=shape[k],marker=marker[k],markevery=markevery,color=color[k])
				k=k+1
			except:
				print 'faulty data ',label
				pass
		plt.legend(loc=1,fontsize=14)
		#axarr[0,0].locator_params(nbins=4, axis='x')
		plt.ylim(1e-8,1e2)		
                plt.legend().set_visible(False)
		plt.xlim(1e-2,7)	
		plt.ylim(0.06,1e2)		
	
	if i==3:
		ax=plt.sca(axarr[1,1])
		axarr[1,1].text(textr, 0.94, metals[i], transform=axarr[1,1].transAxes, fontsize=16,verticalalignment='top', bbox=props)	
		path_s=path+'/set1.5a'
		k=0
		for m in [12.,15.,20.,25.]:
			try:
				path_exp1=path_s+'/see_exp/M'+str(m)+'Z1.0e-04.delay'
				path_wind1=path_s+'/see_wind/M'+str(m)+'Z1.0e-04/M'+str(m)+'Z1.0e-04'
				label=str(m)+'$M_{\odot}$' #, Z='+str(metals[i])
				plot_infall_vel(xaxis,axarr[1,1], path_exp1,path_wind1,label=label,shape=shape[k],marker=marker[k],markevery=markevery,color=color[k])
				k=k+1
			except:
				print 'faulty data ',label
				pass
		plt.legend(loc=1,fontsize=14)
		#axarr[0,0].locator_params(nbins=4, axis='x')
		plt.ylim(1e-8,1e2)		
                plt.legend().set_visible(False)
		plt.xlim(1e-2,7)	
		plt.ylim(0.06,1e2)		
	

f.set_size_inches(12,10,forward=True)
f.subplots_adjust(hspace=0.1,wspace=0.4)

f.set_size_inches(8,6,forward=True)
axarr[0,0].texts[0].set_fontsize(12)
axarr[0,1].texts[0].set_fontsize(12)
axarr[1,1].texts[0].set_fontsize(12)
axarr[1,0].texts[0].set_fontsize(12)
axarr[0,0].legend(fontsize=10)



plt.savefig('infall_velocities.pdf')