#-*- coding: utf-8 -*-
# Comparação do espectros de espalhamento baseados na função modificada de Bessel do primeiro tipo de ordem 1 
# com o modelo Gaussiano. 
# Referência: R. P. Pezzi, et al. Applied Physics Letters, 92 164102 (2008)

from pylab import *
from numpy import nan_to_num, zeros
from scipy.special import i1
from Ewindow import Edict
from layers import Ldict
PI=math.pi

####################################################################################################################################################
# Seção de choque Rutherford
def CSRf(Z1, Z2, Ein, M1, M2, Theta):
    if M1 < M2:
        return ((Z1*Z2*4.8e-20*1e8)/(4*Ein))**2 * sin(Theta)**-4 * (sqrt(1-((M1/M2)*sin(Theta))**2)+cos(Theta))**2 / sqrt(1-((M1/M2)*sin(Theta))**2) 
    else:
        return 0

####################################################################################################################################################
# Função de convolução com forma de linha
def convoluiF0(espectro, alpha,passoE):

    areaantes = espectro.sum()

    temp = espectro[::-1]
    eixo = arange(0,12.*alpha,passoE)
    lineshape = alpha*exp(-eixo/alpha)

#    ls_plt=plt.plot(lineshape/lineshape.sum(),lw=2)
#    sig_plt=plt.plot(temp/temp.sum(),lw=2)

    temp = convolve(temp/temp.sum(), lineshape/lineshape.sum())
    temp.resize(len(espectro))

#    result_plt=plt.plot(temp/temp.sum(),'o')
#    legend([ls_plt,sig_plt,result_plt],["Lineshape","Sinal","Resultado"])

    return temp[::-1]*areaantes

####################################################################################################################################################
# Função de convolução com Gaussiana
def convoluiGauss(espectro, Sigma,passoE):

    areaantes=espectro.sum()
    temp=zeros(len(espectro),float)
    wg4_step=4.0*Sigma/passoE;
    wg2=Sigma**2
    
    #eixo=arange(-6.*Sigma,6.*Sigma,passoE)
    eixo=arange(0,len(temp)*passoE,passoE)

    for i in range(0,len(eixo)):
        for j in range(0,i):
            temp[i]+=espectro[j]*exp(-1.0*pow(passoE*(i-j-wg4_step),2.0)/(wg2*2.0))
            
    return temp*areaantes/temp.sum(), wg4_step*passoE


# Função de convolução com Gaussiana
def convoluiGaussdev(espectro, Sigma,passoE):

    areaantes=espectro.sum()
    temp=zeros(len(espectro),float)
    wg4_step=4.0*Sigma/passoE;
    wg2=Sigma**2
    
    #eixo=arange(-6.*Sigma,6.*Sigma,passoE)
    eixo=arange(0,len(temp)*passoE,passoE)
    gaussi = exp(-1.0*pow(passoE*(eixo-wg4_step),2.0)/(wg2*2.0))

    for i in range(0,len(eixo)):
        for j in range(0,i):
            temp[i]+=espectro[j]*gaussi[i-j]
            
    return temp*areaantes/temp.sum(), wg4_step*passoE

####################################################################################################################################################
####################################################################################################################################################
# Funcoes em comum

def spect_file(x, y, shift, dose):
    data2 = open('temp.sim', 'w')
    for i in range(len(y)):
        print>>data2, x[i]-shift,y[i]*dose
    data2.close()

def set_lim(x, y, shift, dose):  
    ylim([0,max(y*dose)*1.1])          
    xlim([min(x-shift)*0.95,max(x-shift)*1.05])

####################################################################################################################################################
####################################################################################################################################################

def spectro(p, modelo, ion, OFlabentries, LScontrol, CGausscontrol, tempor):

    EPasso = float(OFlabentries[16].get())
    passox = float(OFlabentries[17].get())
    e = arange(float(OFlabentries[14].get()), float(OFlabentries[15].get()), EPasso)
    resol = float(OFlabentries[13].get())
    dose = float(OFlabentries[24].get())
    gshift=0
    X = e*0.

    Theta_s = PI - (p['Theta_in'] + p['Theta_out'])*PI/180
    mi = ion['mass']

    ## Calcula o espectro de cada elemento da amostra
    for componente in Edict:
   
        mt = Edict[componente]['mass']
        c = Edict[componente]['dist']
        alpha_lineshape = Edict[componente]['LineShape']

        k = ((sqrt (mt**2-mi**2*(sin(Theta_s)**2)) + mi*cos(Theta_s) ) / (mi+mt) )**2	

        #k = mi*cos(Theta_s) + sqrt (mt**2 - mi**2*sqrt(sin(Theta_s)))  
        #k = k*k/(mi+mt)/(mi+mt)

        dedxef=p['dedx']*(k/cos(p['Theta_in']*PI/180.) + 1/cos(p['Theta_out']*PI/180.))
        dw2dxef=p['dW2dx']*(k**2./cos(p['Theta_in']*PI/180.)+1./cos(p['Theta_out']*PI/180.))

        sinalbessel = e*0
        sinalgaussiana = e*0

        # Calcula a contribuição de cada camada
        for x in arange(passox, Edict[componente]['profundidademax'], passox): 

            # Seção de choque - CSRf(Z1, Z2, Ein, M1, M2, Theta):
            CS = c[int(x/passox)]*k*CSRf(ion['Z'], Edict[componente]['Z'], k*p['E0']-dedxef*x, mi, mt, Theta_s)*passox 
            # É multiplicada pelo passo para manter a escala do gráfico independente do mesmo.

            # Bessel:  
            # Ainda falta incluir a delta de Dirac na origem
            # Existe um valor máximo em energia no qual a função de Bessel pode ser computada em função
            # da precisão numérica. Uma possível solução pode ser trabalhar em outras unidades que não
            # resultem em argumentos tão grandes para a função exponencial.         
            # Também temos um problema com a primeira camada, que tem profundidade zero.               
            # Perceba que os cálculos de Bessel correspondentes às camadas mais fundas estão truncados
            #  para uma perda de energia maior que 25 keV. (Observado com Theta_in=70, para x > 30)
            if modelo == 'bessel':
                alpha=dedxef*(2./dw2dxef)
                m=alpha*dedxef
                lbd=m*x*alpha
                besselcamada = lbd*exp(-m*x-alpha*(k*p['E0']-e))*i1(2.*sqrt(lbd*(k*p['E0']-e)))/(sqrt(lbd*(k*p['E0']-e)))
                sinalbessel = besselcamada*CS+sinalbessel

            # Gaussiana:
            elif modelo == 'gaussiana':
                Em=k*p['E0']-x*dedxef
                sigma=x*dw2dxef
                gaussianacamada = exp((-(e-Em)**2)/(2.*sigma))/sqrt(2.*PI*sigma) 
                sinalgaussiana = gaussianacamada*CS+sinalgaussiana

        if modelo == 'bessel':
            besself = nan_to_num(sinalbessel) 
            if LScontrol == 1:
                besself = convoluiF0(besself, alpha_lineshape, EPasso)
            X = besself + X
            linecolor='r'

        elif modelo == 'gaussiana':
            if LScontrol == 1:     
                sinalgaussiana = convoluiF0(sinalgaussiana, alpha_lineshape, EPasso)           
            X = sinalgaussiana + X
            linecolor='b'

    if CGausscontrol == 1:
        X, gshift = convoluiGauss(X,resol/2.35482,EPasso)    

    #tempor.set_xdata(e-gshift)           
    #tempor.set_ydata(X*dose)
    plot(e-gshift, X*dose)

    set_lim(e, X, gshift, dose)
    spect_file(e, X, gshift, dose)
    
    plt.show()

####################################################################################################################################################
####################################################################################################################################################

def spectrolayer(p, modelo, ion, OFlabentries, LScontrol, CGausscontrol, tempor):

    EPasso = float(OFlabentries[16].get())
    passox = float(OFlabentries[17].get())
    e = arange(float(OFlabentries[14].get()), float(OFlabentries[15].get()), EPasso)
    resol = float(OFlabentries[13].get())
    dose = float(OFlabentries[24].get())
    X = e*0.
    gshift=0.
    curdepht=passox    #Avoid zero division

    Theta_s = PI - (p['Theta_in'] + p['Theta_out'])*PI/180
    mi = ion['mass']

    for camada in Ldict:

        xfinal = curdepht + float(camada.LandE[5].get())

        ## Calcula o espectro de cada elemento da amostra
        for componente in range(len(camada.Elements)):

            mt = camada.Elements[componente]['mass']
            c = float(camada.Spinboxes[componente].get())
            alpha_lineshape = camada.Elements[componente]['LineShape']
            dedx = float(camada.LandE[3].get())
            dW2dx = float(camada.LandE[4].get())

            k = ((sqrt (mt**2-mi**2*(sin(Theta_s)**2)) + mi*cos(Theta_s) ) / (mi+mt) )**2	
            dedxef=dedx*(k/cos(p['Theta_in']*PI/180.) + 1/cos(p['Theta_out']*PI/180.))
            dw2dxef=dW2dx*(k**2./cos(p['Theta_in']*PI/180.)+1./cos(p['Theta_out']*PI/180.))

            sinalbessel = e*0
            sinalgaussiana = e*0

            # Calcula a contribuição de cada passo
            for x in arange(curdepht, xfinal, passox): 

                # Seção de choque - CSRf(Z1, Z2, Ein, M1, M2, Theta):
                CS = c*k*CSRf(ion['Z'], camada.Elements[componente]['Z'], k*p['E0']-dedxef*x, mi, mt, Theta_s)*passox 
                # É multiplicada pelo passo para manter a escala do gráfico independente do mesmo.

                # Bessel:  
                if modelo == 'bessel':
                    alpha=dedxef*(2./dw2dxef)
                    m=alpha*dedxef
                    lbd=m*x*alpha
                    besselcamada = lbd*exp(-m*x-alpha*(k*p['E0']-e))*i1(2.*sqrt(lbd*(k*p['E0']-e)))/(sqrt(lbd*(k*p['E0']-e)))
                    sinalbessel = besselcamada*CS+sinalbessel

                # Gaussiana:
                elif modelo == 'gaussiana':
                    Em=k*p['E0']-x*dedxef
                    sigma=x*dw2dxef
                    gaussianacamada = exp((-(e-Em)**2)/(2.*sigma))/sqrt(2.*PI*sigma) 
                    sinalgaussiana = gaussianacamada*CS+sinalgaussiana

            if modelo == 'bessel':
                besself = nan_to_num(sinalbessel) 
                if LScontrol == 1:
                    besself = convoluiF0(besself, alpha_lineshape, EPasso)
                X = besself + X
                linecolor='r'

            elif modelo == 'gaussiana':
                if LScontrol == 1:     
                    sinalgaussiana = convoluiF0(sinalgaussiana, alpha_lineshape, EPasso)           
                X = sinalgaussiana + X
                linecolor='b'

        curdepht = xfinal

    if CGausscontrol == 1:
        X, gshift = convoluiGauss(X,resol/2.35482,EPasso)  
  
    tempor.set_xdata(e-gshift)           
    tempor.set_ydata(X*dose)

    set_lim(e, X, gshift, dose)
    spect_file(e, X, gshift, dose)
  
    plt.show()

####################################################################################################################################################
####################################################################################################################################################

def spectroRNA(p, modelo, ion, OFlabentries, LScontrol, CGausscontrol, tempor):

    EPasso = float(OFlabentries[16].get())
    passox = float(OFlabentries[17].get())
    e = arange(float(OFlabentries[14].get()), float(OFlabentries[15].get()), EPasso)
    resol = float(OFlabentries[13].get())
    dose = float(OFlabentries[24].get())
    X = e*0.
    gshift=0

    Theta_s = PI - (p['Theta_in'] + p['Theta_out'])*PI/180
    mi = ion['mass']

    ## Calcula o espectro o elemento da amostra
   
    mt = Edict[0]['mass']
    c = Edict[0]['dist']
    alpha_lineshape = Edict[0]['LineShape']

    k = ((sqrt (mt**2-mi**2*(sin(Theta_s)**2)) + mi*cos(Theta_s) ) / (mi+mt) )**2	
    dedxef=p['dedx']*(k/cos(p['Theta_in']*PI/180.) + 1/cos(p['Theta_out']*PI/180.))
    dw2dxef=p['dW2dx']*(k**2./cos(p['Theta_in']*PI/180.)+1./cos(p['Theta_out']*PI/180.))

    sinalbessel = e*0
    sinalgaussiana = e*0

    # Calcula a contribuição de cada camada
    for x in arange(passox, Edict[0]['profundidademax'], passox): 

        # Seção de choque - CSRf(Z1, Z2, Ein, M1, M2, Theta):
        CS = c[int(x/passox)]*k*CSRf(ion['Z'], Edict[0]['Z'], k*p['E0']-dedxef*x, mi, mt, Theta_s)*passox 
        # É multiplicada pelo passo para manter a escala do gráfico independente do mesmo.

        # Bessel:  
        if modelo == 'bessel':
            alpha=dedxef*(2./dw2dxef)
            m=alpha*dedxef
            lbd=m*x*alpha
            besselcamada = lbd*exp(-m*x-alpha*(k*p['E0']-e))*i1(2.*sqrt(lbd*(k*p['E0']-e)))/(sqrt(lbd*(k*p['E0']-e)))
            sinalbessel = besselcamada*CS+sinalbessel

        # Gaussiana:
        elif modelo == 'gaussiana':
            Em=k*p['E0']-x*dedxef
            sigma=x*dw2dxef
            gaussianacamada = exp((-(e-Em)**2)/(2.*sigma))/sqrt(2.*PI*sigma) 
            sinalgaussiana = gaussianacamada*CS+sinalgaussiana

    if modelo == 'bessel':
        besself = nan_to_num(sinalbessel) 
        if LScontrol == 1:
            besself = convoluiF0(besself, alpha_lineshape, EPasso)
        X = besself + X
        linecolor='r'

    elif modelo == 'gaussiana':
        if LScontrol == 1:     
            sinalgaussiana = convoluiF0(sinalgaussiana, alpha_lineshape, EPasso)           
        X = sinalgaussiana + X
        linecolor='b'

    if CGausscontrol == 1:
        X, gshift = convoluiGauss(X,resol/2.35482,EPasso)    

    tempor.set_xdata(e-gshift)           
    tempor.set_ydata(X*dose)

    set_lim(e, X, gshift, dose)
    spect_file(e, X, gshift, dose)
     
    plt.show()

####################################################################################################################################################
####################################################################################################################################################