from utils import*
from math import factorial as fact
path='../multislice/docs_fig/multi2D/Pcoh_kin_dual_dyn'
plt.close('all')
topt = 1        #number of multiple scattering to show
rho   = 0.106   #atoms/A^3
sig_e = np.array([0.001,0.0025,0.005]) #A^2
le_s  = 1/(sig_e*rho) #A

z = np.linspace(0,10000,1000) #A

cs=dsp.getCs('Reds',topt+2)
l_e  = le_s[1]
y = [1/fact(m)*(z/l_e)**m*np.exp(-z/l_e) for m in range(topt+1)]
plts = [ [z/10,y[m],cs[m],'$%d$' %m] for m in range(topt+1)]
plts+= [[z/10,1-np.array(y).sum(axis=0),'k--','']]
dsp.stddisp(plts,lw=2,labs=[r'$z(nm)$','P(z)'],#title=tle,
    name=path+'%d.svg' %0, opt='p',fonts={'leg':20,'title':20})


# for i,le in zip(range(le_s.size),le_s):
#     Pcoh = np.exp(-z/le)
#     Pkin = z/le*np.exp(-z/le)
#     Pdual = 0.5*(z/le)**2*np.exp(-z/le)
#     if topt:
#         Ptri = 1/6*(z/le)**3*np.exp(-z/le)
#         Pdyn = 1-np.exp(-z/le)*(1+z/le+0.5*(z/le)**2+(z/le)**3/6)
#         S = Pcoh+Pkin+Pdual+Ptri+Pdyn
#     else:
#         Pdyn = 1-np.exp(-z/le)*(1+z/le+0.5*(z/le)**2)
#         S = Pcoh+Pkin+Pdual+Pdyn
#
#     tle = r'$\sigma_e=%.3f\AA^2$, $l_e$=%.d nm' %(sig_e[i],int(le/10))
#     cs=dsp.getCs('jet',4)
#     plts = [[z/10,Pcoh,cs[0],'$coh$'],[z/10,Pkin,cs[1],'$kin$']]
#     plts+= [[z/10,Pdual,cs[2],'$double$']]
#     if topt: plts+= [[z/10,Ptri ,cs[3],'$tri$']]
#     plts+= [[z/10,Pdyn,'k','$dyn$']]#,[z/10,S,'k--']]
#     dsp.stddisp(plts,lw=2,labs=[r'$z(nm)$','P(z)'],#title=tle,
#         name=path+'%d.svg' %i, opt='p',fonts={'leg':20,'title':20})
