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, kE0):
    # Funcao que retorna a perda de energia de acordo com a funcao 
    #   modificada de bessel do primeiro tipo de ordem 1
    #   N FUNCIONA!
    lbd = m*alpha*x
    #print  lbd*exp((-m*x-alpha*(e-kE0))/1000)*i1(2*sqrt(lbd*(e-kE0)/1000))/(sqrt(lbd*(e-kE0)/1000 +.0000001))
    return lbd*exp((-m*x-alpha*(e-kE0))/1000)*i1(2*sqrt(lbd*(e-kE0)/1000))/(sqrt(lbd*(e-kE0)/1000 +.0000001))

## 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, kE0):
        # Adiciona uma funcao modificada de bessel do primeiro tipo:
        #  energia media Em, com variancia sigma com o peso dado
        #self.Y=self.Y +peso*modbesselp(self.e, Em, sigma, dedx, dw2dx, x) 
        self.Y=self.Y +peso*modbessel(self.e, alpha, m, x, kE0)
    
    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: 
                passo = 0.01
                c = SLOT[ATOM]['dist']
                x = arange(0,20,passo)
                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 [0.1]: #x:
                    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)
                    if modelo == 'gaussiana':
                        self.addgauss(Em, sigma, c[int(i/passo)]*k*CS)
                    elif modelo == 'bessel':
			print Em
                        alpha = (2.* param['dedx']/10. * Em )  / ((param['dW2dx']/10.) *(sigma/i))
                        m = alpha * param['dedx']/10. * Em
                        self.addbessel(c[int(i/passo)]*k*CS, alpha, m, i,k*self.E0)