import numpy as np
import matplotlib.pyplot as plt
import RMT
import matplotlib
matplotlib.rcParams.update({'font.size': 18})
matplotlib.rcParams.update({'lines.linewidth': 2})
matplotlib.rcParams.update({'axes.linewidth': 1.5})
matplotlib.rcParams.update({'xtick.major.width': 1.5})
matplotlib.rcParams.update({'ytick.major.width': 1.5})

for ii in range(20):
    for N_inel in [2,4]:
	N_hf=9
    #    N_inel=2
	N_closed=N_hf-N_inel-1

	d=1.0
	k_univ=1.0
	z=0.15
	T_sa=1.0

	range_min=-1
	range_max=4
	Gamma_sa=T_sa*d/(2*np.pi)
	Gamma_sse=z*Gamma_sa
	Gamma_laser=0.000

	Gamma_inel_1=N_inel*Gamma_sse+Gamma_laser
	Peak_1=k_univ*Gamma_sa/Gamma_inel_1
	Gamma_inel_2=Gamma_sa+Gamma_laser
	Peak_2=k_univ*Gamma_sse/Gamma_inel_2
	Gamma_inel_3=N_inel*Gamma_sse+Gamma_laser
	Peak_3=k_univ*Gamma_sse/Gamma_inel_3
#	print Peak_1,Peak_2,Peak_3
#	print Gamma_inel_1,Gamma_inel_2,Gamma_inel_3
#	print Peak_1,Gamma_inel_1
#	print Peak_2,Gamma_inel_2
#	print Peak_3,Gamma_inel_3

        RMT.init("Run{}".format(ii))
        pos_1=RMT.GOE_unfolded()
        pos_2=[]
        for j in xrange(N_inel):
            pos_2.extend(RMT.GOE_unfolded())
        pos_3=[]
        for j in xrange(N_closed):
            pos_3.extend(RMT.GOE_unfolded())
	pos_1=[x for x in pos_1 if (x>=range_min and x<=range_max)]
	pos_2=[x for x in pos_2 if (x>=range_min and x<=range_max)]
	pos_3=[x for x in pos_3 if (x>=range_min and x<=range_max)]


	plt.clf()
	E_range=np.linspace(0,3*d,3001)
	profile_tot=np.zeros_like(E_range)
	profile_type_tot=np.zeros_like(E_range)
	for pos in pos_1:
	    profile=Peak_1/((2*(E_range-pos)/Gamma_inel_1)**2+1)
	    profile_tot+=profile
	    profile_type_tot+=profile
	    plt.plot(E_range,profile,"r--")
	plt.plot(E_range,profile_type_tot,"r-")
	profile_type_tot=np.zeros_like(E_range)
	for pos in pos_3:
	    profile=Peak_3/((2*(E_range-pos)/Gamma_inel_3)**2+1)
	    profile_tot+=profile
	    profile_type_tot+=profile
	    plt.plot(E_range,profile,"b--")
	plt.plot(E_range,profile_type_tot,"b-")
	profile_type_tot=np.zeros_like(E_range)
	for pos in pos_2:
	    profile=Peak_2/((2*(E_range-pos)/Gamma_inel_2)**2+1)
	    profile_tot+=profile
	    profile_type_tot+=profile
	    plt.plot(E_range,profile,"g--")
	plt.plot(E_range,profile_type_tot,"g-")
	plt.plot(E_range,profile_tot,"k-")
	plt.plot([0.0,1.0],[-1.0,-1.0],"r--",label="incoming")
	plt.plot([0.0,1.0],[-1.0,-1.0],"g--",label="outgoing")
	plt.plot([0.0,1.0],[-1.0,-1.0],"b--",label="closed")
	plt.plot([0.0,1.0],[-1.0,-1.0],"k-",label="total")
	plt.legend(loc='upper left')
	plt.xlabel(r"$E/d$")
	plt.ylabel(r"$k_2/k_2^{\rm univ}$")
#	plt.yscale('log')
#	plt.ylim(1e-2,1e1)
	plt.ylim(0.0,2.5)
	plt.xlim(0.0,3.0)
	name="plot_{}_{}.pdf".format(ii,N_inel)
	plt.tight_layout()
	plt.savefig(name)
	print name

