from pylab import *

PI=math.pi

## Primeiro defino algumas funcoes que sao 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 modbesselp(e, param,x):
    # Funcao que retorna a perda de energia de acordo com a funcao 
    #   modificada de bessel do primeiro tipo de ordem 1
    # **** NAO FUNCIONA!! ****
    alpha=2.*param['dedx']/param['dW2dx']
    m = alpha * param['dedx']
    lbd=m*x*alpha
    return (lbd*exp(-m*x-alpha*e)/(sqrt(lbd*e)))*i1(2.*sqrt(lbd*e))


def profile(pares, passo):
    # Funcao que retorna os vetores espessura e perfil de concentracao
    #  a partir dos pares de pontos fornecidos.
    # pares deve ser uma matriz de duas colunas: [espessura, concentracao] 
    #  Modos: degraus - distribuicao do tipo escada. Ex. [10 , 1]
    #         representa uma amostra com 10 angstroms de espessura
    #         com concentracao constante    
    depth = sum(pares[:,0]) # A espessura total eh a soma da primeira coluna
    x = arange(0,depth, passo) # Monta o eixo X
    c = [] # Inicia o vetor concentracao vazio
    for i in range(len(pares)): # popula o vetor concentracao
       temp = ones(pares[i][0]/passo)*pares[i][1]
       c.extend(temp.tolist())
    return x,c



## 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,param,peso,x):
        # 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, param, x) 

    
    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
        #  Ainda falta colocar os nomes nos eixos e outras descricoes
        # 
        plot(self.e,self.Y)
        show()

    def setparam (self, parametros):
        #Define os parametros fixos do spectro
        # parametros eh um dicionario definido com algo como
        # param = dict(dedx=287. , dW2dx=35000 , E0= 100000, Theta_in=0, Theta_out=70, fwhm0=300 )
        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.)
       # self. = parametros['']

    def addelements(self,arquivos, param,k,modelo):
        # Adiciona ao spectro os elementos definidos nos arquivos com 
        #  os parametros dados 
        #  *** Eh preciso incluir o fator cinematico apropriadamente ***
        #  Argumentos para addelements sao:
        #       Lista de arquivos do tipo ['file1.dat','file2.dat']
        #       Parametros
        #       Modelo de perda de energia: 'gaussiana' ou 'bessel'

        # ** Falta incluir informacoes de massa do elemento no arquivo
        # ** quem sabe na primeira linha incluir numero atomico e massa
        # ** e ao carregar o arquivo com np.loadtxt(), retirar estes valores
        # ** dos vetores x e c

        for arquivo in arquivos:
            amostra = np.loadtxt(arquivo)
            x=amostra[:,0]*0.01
            c=amostra[:,1]*1.
            self.setparam(param)
            for i in x:
               sigma=i*self.dW2dx*(k*k/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.))
               if modelo == 'gaussiana':
                   self.addgauss(Em, sigma, c[i*100]*k)

               if modelo == 'bessel':
                   self.addbessel(param, c[i*100]*k,x)
                   
            k=k-0.21


    def addelements2(self,arquivos, param,k,modelo,passo):
        # Adiciona ao spectro os elementos definidos nos arquivos com 
        #  os parametros dados 
        #  *** Eh preciso incluir o fator cinematico apropriadamente ***
        #  Argumentos para addelements sao:
        #       Lista de arquivos do tipo ['file1.dat','file2.dat']
        #       Parametros
        #       Modelo de perda de energia: 'gaussiana' ou 'bessel'

        # ** Falta incluir informacoes de massa do elemento no arquivo
        # ** quem sabe na primeira linha incluir numero atomico e massa
        # ** e ao carregar o arquivo com np.loadtxt(), retirar estes valores
        # ** dos vetores x e c

        for arquivo in arquivos:
            amostra = np.loadtxt(arquivo)
            x,c = profile(amostra,passo)
            self.setparam(param)
            for i in x:
               sigma=i*self.dW2dx*(k*k/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.))
               if modelo == 'gaussiana':
                   self.addgauss(Em, sigma, c[int(i/passo)]*k)

               if modelo == 'bessel':
                   self.addbessel(param, c[i*100]*k,x)
                   
            k=k-0.21