from pylab import *
from scipy.special import i1
from Edict import *

PI=math.pi

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

## funcoes usadas pela classe spectro
def gaussian(e, Em, sigma):
    # Funcao que retorna uma gaussiana centrada em Em e largura sigma
    return exp((-(e-Em)**2)/(2.*sigma))/sqrt(2.*PI*sigma) 

#def modbessel(e, alpha, m, x, k, E0):
    # Funcao que retorna a perda de energia de acordo com a funcao 
    #   modificada de bessel do primeiro tipo de ordem 1
    #   N FUNCIONA!

## Definindo classe referente a objetos de espectro de espalhamento
class spectro_espalhamento:
    def __init__(self, Emin, Emax, EPasso):
        # Inicia o spectro no limite dado, com o passo EPasso
        self.e = arange(Emin, Emax, EPasso)
        self.Y = e*0.

    def addbessel(self, peso, alpha, m, x, k, E0):
        # Adiciona uma funcao modificada de bessel do primeiro tipo:
        #  energia media Em, com variancia sigma com o peso dado
        lbd=m*x*alpha
        bessel = lbd*exp(-m*x-alpha*(k*E0-e))*i1(2.*sqrt(lbd*(k*E0-e)))/(sqrt(lbd*(k*E0-e)))
        self.Y = self.Y + peso*bessel #modbessel(self.e, alpha, m, x, k, E0)
        print len(self.Y), len(self.e)
    
    def addgauss(self,Em,sigma,peso):
        # Adiciona uma gaussiana ao spectro:
        #  centrada em Em, com variancia sigma com o peso dado
        self.Y = self.Y+peso*gaussian(self.e, Em, sigma) 

    def plot(self):
        # Plota o spectro
        plt.plot(self.e,self.Y)
        suptitle('Signal received', fontsize=12)
        xlabel("Enegy (eV)")
        ylabel("Yield")
        show()

    def setparam (self, parametros):
        # Define os parametros fixos do spectro
        self.E0 = parametros['E0']
        self.dedx = parametros['dedx']
        self.dW2dx = parametros['dW2dx']
        self.Theta_in = parametros['Theta_in']
        self.Theta_out = parametros['Theta_out']
        self.FWHM0 = parametros['FWHM0']
        self.sigma0 = pow(self.FWHM0/2.35,2.)

    def sweepelements(self, param, modelo, ion):
        for ATOM in SLOT:
            if SLOT[ATOM]['control']==0:
                pass
            else: 
                xstep = 0.01
                c = SLOT[ATOM]['dist']
                x = arange(0.1,20,xstep)
                self.setparam(param)

                Theta_s = PI - (self.Theta_in + self.Theta_out)*PI/180
                mt = SLOT[ATOM]['mass']
                mi = ion['mass']
                k = ( (sqrt (mt**2+mi**2*(sin(Theta_s)**2)) + mi*cos(Theta_s) ) / (mi+mt) )**2	
               
                for i in x:

                    if modelo == 'gaussiana':
                        sigma=i*self.dW2dx*(k**2./cos(self.Theta_in*PI/180.)+1./cos(self.Theta_out*PI/180.))+self.sigma0
                        Em = k*self.E0 - i*self.dedx*(k/cos(self.Theta_in*PI/180.) + 1/cos(self.Theta_out*PI/180.))
                        CS = CSRf(ion['Z'], SLOT[ATOM]['Z'], Em, ion['mass'], SLOT[ATOM]['mass'], Theta_s)
                        self.addgauss(Em, sigma, c[int(i/xstep)]*k*CS)

                    elif modelo == 'bessel':
                        dedxef = self.dedx*(k/cos(self.Theta_in*PI/180.) + 1/cos(self.Theta_out*PI/180.))
                        dw2dxef = self.dW2dx*(k**2./cos(self.Theta_in*PI/180.)+1./cos(self.Theta_out*PI/180.))
                        alpha = dedxef*(2./dw2dxef)
                        m = alpha*dedxef
                        Em = k*self.E0-dedxef*x
                        CS = CSRf(ion['Z'], SLOT[ATOM]['Z'], Em , ion['mass'], SLOT[ATOM]['mass'], Theta_s)
                        self.addbessel(c[int(i/xstep)]*k*CS*xstep, alpha, m, i, k, self.E0)