"""
This module use  the xraylib library
and provide high level functions for
absorption factors for chemical compounds

August 2012 by OT
based on work by AT
"""
from string import *
import numpy
import os
from pySAXS import xraylib
from .absorption_utils import *

KEV2ANGST = xraylib.KEV2ANGST #12.39841875

def getZName(name):
    '''
    return the Atomic Number Z from compound name
    >>> getZName("Copper")
    29
    '''
    Z=ATOMS_NAME.index(name)+1
    return Z

def SymbolToAtomicNumber(symbol):
    '''
    get the atomic number of a symbol
    copy of xraylib.SymbolToAtomicNumber(symbol)
    >>> SymbolToAtomicNumber('Cu')
    29
    '''
    return xraylib.SymbolToAtomicNumber(symbol)

def AtomicNumberToSymbol(Z):
    '''
    get the symbol from an atomic number
    copy of xraylib.AtomicNumberToSymbol(Z)
    >>> AtomicNumberToSymbol(29)
    'Cu'
    '''
    return xraylib.AtomicNumberToSymbol(Z)

def getNameZ(Z):
    '''
    return the name from a Z
    >>> getNameZ(29)
    'Copper'
    '''
    return ATOMS_NAME[Z-1]
    

def getAtomsFormula(S):
    '''
    Transform a string containing a chemical formula ('C 1 O 2') in two array
    - list of atoms
    - numeric array of atoms number
    7-13-2012 by OT
    >>> getAtomsFormula('C 1 O 2')
    ['C','O'] [1,2]
    '''
    S=S.strip()
    L=S.split(' ') #cut at all spaces
    rS=L[::2] #Symbols to be return
    rN=numpy.array(L[1::2],dtype='float') #Numbers to be return
    return rS,rN

def getFormulaAtoms(atomS,atomN):
    '''
    transform a compound list and number in string
    >>> AtomsToCompound(['H','O'],[2,1])
    'H 2 O 1' 
    '''
    S=''
    if len(atomS)==len(atomN):
        for i in range(len(atomS)):
            S=S+' '+str(atomS[i])+' '+str(atomN[i])
    return S.strip()

def getMuZ(Z,energy=8.028):
    '''
    return Mu from Atomic number Z
    >>> getMuZ(4)
    0.8719998598098755
    '''
    return xraylib.CS_Photo(Z,energy)

def getMuSymbol(S,energy=8.028):
    """
    This function returns the xray absorption coefficient of the atom with
    symbole S (ie 'C' fo carbon, 'Au' for gold) by default energy is 8.03 keV
    and mass-energy absorption coefficient are return. To get the mass
    absorption coefficient ISEN has to be 0
    """
    Z=xraylib.AtomicNumberToSymbol(S)
    return getMuZ(Z ,energy)
    
def getMuName(name,energy=8.028):
    '''
    This function returns the xray absorption coefficient of the atom with
    name S (ie 'Carbon' , 'Gold') by default energy is 8.028 keV and mass-energy
    absorption coefficient are return. To get the mass absorption coefficient
    '''
    Z=getZName(name)
    return getMuZ(Z,energy)
    


def getMasseZ(Z):
    '''
    returns the molar mass (atomic weight) of the atomic number Z
    '''
    return xraylib.AtomicWeight(Z)

def getMasseSymbol(S):
    '''
    returns the molar mass (atomic weight) of the atom with atomic number Z
    (ie 'C' fo carbon, 'Au' for gold)
    '''
    Z=xraylib.SymbolToAtomicNumber(S)
    return getMasseZ(Z)



def getMuFormula(S,energy=8.028):
    '''
    This function returns the mass attenuation coefficient for a chemical
    formula in the form 'C 6 H 6 O 2 N 1'
    '''
    #29-8-2012 by OT use of xraylib
    #7-13-2012 by OT improved for formula with fraction 'Si 1.2 Al 0.8'
    
    if len(S)!=0:
        S,N=getAtomsFormula(S)
        Mu=numpy.zeros(len(S),float)
        Ma=numpy.zeros(len(S),float)
        CMa=0.0
        CMu=0.0
        for i in numpy.arange(len(S)):
            Z=xraylib.SymbolToAtomicNumber(S[i])
            Mu[i]=xraylib.CS_Photo(Z,energy)
            Ma[i]=xraylib.AtomicWeight(Z)
            CMa=CMa+Ma[i]*N[i]
            CMu=CMu+Mu[i]*Ma[i]*N[i]
        #print CMu
        return CMu/CMa
    else:
        return 0.0

def getMasseFormula(S,energy=8.028):
    """
    This function returns the molar mass a chemical formula
    in the form  'C 6 H 6 O 2 N 1'
    """
    if len(S)!=0:
        S,N=getAtomsFormula(S)
        Ma=numpy.zeros(len(S),float)
        CMa=0.0
        for i in numpy.arange(len(S)):
            Ma[i]=getMasseSymbol(S[i])
            CMa=CMa+Ma[i]*N[i]
        return CMa
    else:
        return 0.0

def getElectronNumber(S):
    S,N=getAtomsFormula(S)
    Z=0
    for i in numpy.arange(len(S)):
        Z=Z+(xraylib.SymbolToAtomicNumber(S[i]))*N[i]
    return Z

def getElectronDensity(S,rho):
    '''
    return the electron density and the scattering length density
    '''
    #Thomson scattering length
    Belec=0.282e-12
    Navo=xraylib.AVOGNUM*1e24
    M=getMasseFormula(S)
    Z=getElectronNumber(S)
    ED=(rho*Z*Navo)/M
    #print "rho=",rho,"\tZ=",Z,"\tNavo=",Navo,"\tM=",M,"\tED=",ED
    return ED,ED*Belec
    
    
def getEnergyFromSource(source='Cu'):
    '''
    return the KA LINE (most used) energy from the x-ray source
    >>> getEnergyFromSource('Cu')
    8.04105057076251
    >>> getEnergyFromSource('Mo')
    17.443217030477935
    '''
    Z=xraylib.SymbolToAtomicNumber(source)
    return xraylib.LineEnergy(Z,xraylib.KA_LINE)

def getAngstFromSource(source='Cu'):
    '''
    return the KA LINE (most used) energy in ANGSTROM from the x-ray source
    >>> getAngstFromSource('Cu')
    1.5418904085842968
    '''
    Z=xraylib.SymbolToAtomicNumber(source)
    return xraylib.KEV2ANGST/xraylib.LineEnergy(Z,xraylib.KA_LINE)

def getTransmission(formula,thickness=1.0,density=1.0,energy=8.03):
    '''
    return the transmission of a compound
    thickness in cm
    >>> getTransmission('H 2 O 1',0.1,density=1.0,energy=8.03)
    0.376802369048
    '''
    #calculate mu for compound
    mu=getMuFormula(formula,energy)
    #calculate Transmission
    return numpy.exp(-density*mu*thickness)

def getThickness(formula,transmission=0.5,density=1.0,energy=8.03):
    '''
    return the estimated thickness (cm) from the Transmission Value
    '''
    #print formula, transmission
    mu=getMuFormula(formula,energy)
    #print mu
    eb=-(1/(mu*density))*numpy.log(transmission)
    #print eb
    return eb
                            
if __name__== '__main__':
    formula="H 2 O 1"
    density=1.0
    thickness=0.1 #1 mm
    source='Cu'
    energy=getEnergyFromSource(source)
    print("energy for ",source," : ",energy," keV ->",getAngstFromSource(source),' A')
    print("Calculate Mu from formula for water : ",formula)
    mu=getMuFormula(formula,energy)
    print(mu)
    print("X-ray transmission = ",numpy.exp(-density*mu*thickness))
    print(getTransmission(formula,thickness,density,energy))
    print("--------------------")
    source='Mo'
    energy=getEnergyFromSource(source)
    print("energy for ",source," : ",energy," keV ->",getAngstFromSource(source),' A')
    print("Calculate Mu from formula for water : ",formula)
    mu=getMuFormula(formula,energy)
    print(mu)
    print("X-ray transmission = ",numpy.exp(-density*mu*thickness))
    print(getTransmission(formula,thickness,density,energy))
    print("--------------------")
    
