from string import *
from math import *
import atomic_sub
import mucal


##	Correction uses formulae from Booth & Bridges, 
##	http://xxx.lanl.gov/cond-mat/0306252
##	YOU MUST cite this article, or its in-print version from
##	XAFS XII (in press) if you use this correction.


## FUNCTION form2sym
## USAGE:
##  [symbol,num_sym,tot_syms]=form2sym(formula)
##
##  convert formula into a pair of arrays
##  given:
##    formula - a character string such as
##              Ni2O4
##              YBa2Fe0.31Cu2.69O7
##  returns:
##    symbol - sequence with the atoms in the formula
##    num_sym - sequence of reals with number of symbols
##    tot_syms - number of symbols in formula (an integer)
def form2sym(formula):
    #initialize
    form_len=len(formula)
    symbol=[]
    num_sym=[]
    i=0
    isym=0
    numbers='0123456789.'
    #parse formula and build symbol
    for c in formula:
        if c in uppercase:
            isym=isym+1
            elen=1
            if (i+1<form_len) and (formula[i+1] in lowercase):
                symbol.append(c+formula[i+1])
                elen=2
            else:
                symbol.append(c)
            #find subscript
            if (i+elen<form_len) and (formula[i+elen] in numbers):
                num=''
                j=0
                for n in formula[i+elen:]:
                    if (j==0) and (n in numbers):
                        num=num+n
                    else:
                        j=1
                num_sym.append(float(num))
            else:
                num_sym.append(1.0)
        i=i+1
    #print symbol,num_sym
    return [symbol,num_sym,isym]

## FUNCTION atomic_weight
## USAGE:
##  weight=atomic_weight(Z)
##  returns atomic weight of element Z in g/mol
def atomic_weight(Z):
    weight=[1.0079,4.00260,6.941,9.0128,10.81,12.011,14.0067,
            15.9994,18.998403,20.179,22.98977,24.305,26.98154,
            28.0855,30.97376,32.06,35.453,39.948,39.0983,40.08,
            44.9559,47.90,50.9415,51.996,54.9380,55.847,58.9332,
            58.70,63.546,65.38,69.72,72.59,74.9216,78.96,79.904,
            83.80,85.4678,87.62,88.9059,91.22,92.9064,95.94,98,
            101.07,102.9055,106.4,107.868,112.41,114.82,118.69,
            121.75,127.60,126.9045,131.30,132.9054,137.33,0,0,0,
            0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,
            0,0,0,0,0,0,0 ]
    return weight[int(Z-1)]


## FUNCTION density
## USAGE:
##  dens=density(formula,formvol)
##  returns the density of atoms per barn-cm
##  which is the same as atoms per cubic angstrom
def density(formula,formvol):
    [symbol,num_sym,tot_syms]=form2sym(formula)
    dens=0.0
    for i in num_sym:
        dens=dens+i/formvol
    return dens


## FUNCTION tot_weight
## USAGE:
##  twt=tot_weight(formula)
##  returns total weight of chemical formula
def tot_weight(formula):
    [symbol,num_sym,tot_syms]=form2sym(formula)
    twt=0.0
    for i in range(tot_syms):
        Z=atomic_sub.sym2Z(symbol[i])
        twt=twt+atomic_weight(Z)*num_sym[i]
    return twt


## FUNCTION get_tot_abs
## USAGE:
##  abs=get_tot_abs(formula,energy,formvol)
##  returns absorption in inverse cm
def get_tot_abs(formula,energy,formvol):
    #initialize
    temp=0
    lam_e=12389.2/energy
    [symbol,num_sym,tot_syms]=form2sym(formula)
    #have symbol list...
    #need to figure out which edge coefficients we
    #   need to use from the Victoreen formula
    cross_section=0.0
    for i in range(tot_syms):
        Z=atomic_sub.sym2Z(symbol[i])
        #have edge, get coeffs and add up absorption
        #using McMaster tables...
        [err,energyret,xsec,fl_yield,errmsg] = mucal.mucal(Z=Z,ephot=energy/1000.0,unit='B')
        temp=xsec[3]
        conversion=num_sym[i]
        conversion=conversion*density(symbol[i],formvol)
        #put mu-Z into inverse cm
        mu_Z=conversion*temp
        cross_section=cross_section+mu_Z
    return cross_section


## FUNCTION get_noedge_abs
## USAGE:
##  abs=get_noedge_abs(formula,energy,formvol)
##  returns absorption in inverse cm, MINUS the current absorption edge 
def get_noedge_abs(formula,energy,formvol):
    #initialize
    temp=0
    lam_e=12389.2/energy
    [symbol,num_sym,tot_syms]=form2sym(formula)
    #have symbol list...
    #need to figure out which edge coefficients we
    #   need to use from the Victoreen formula
    cross_section=0.0
    for i in range(tot_syms):
        Z=atomic_sub.sym2Z(symbol[i])
        #have edge, get coeffs and add up absorption
        #using McMaster tables...
        [err,energyret,xsec,fl_yield,errmsg] = mucal.mucal(Z=Z,ephot=energy/1000.0,unit='B')
        temp=xsec[11]
        conversion=num_sym[i]
        conversion=conversion*density(symbol[i],formvol)
        #put mu-Z into inverse cm
        mu_Z=conversion*temp
        cross_section=cross_section+mu_Z
    return cross_section


## FUNCTION thick_chi
## USAGE:
##  chi=thick_chi(chi_exp,mu_t,mu_a,mu_f,phi,theta)
##  in the thick approximation ala Troger and Baberschke et al
##  PRB 46, 3283 (1992)
def thick_chi(chi_exp,mu_t,mu_a,mu_f,phi,theta):
    g=sin(phi)/sin(theta)
    alpha=mu_t+mu_f*g
    chi=chi_exp/(1.0-mu_a/alpha)
    return chi


## FUNCTION info_depth
## USAGE:
##  depth=info_depth(tot_abs,fluor_abs,phi,g)
##  from Troger and Baberschke
def info_depth(tot_abs,fluor_abs,phi,g):
    tcrazy=sin(phi)/(tot_abs+fluor_abs*g)
    if tcrazy<0:
        print phi,g,tot_abs,fluor_abs
    return tcrazy


## FUNCTION corr_chi
## USAGE:
##  chi=corr_chi(k,chi_exp,mu_t,mu_a,mu_f,d_given,phi,g,option)
##   k:       k vector
##   chi_exp: experimental measurement of XAFS chi
##   mu_t:    total abs coeff of materia
##   mu_a:    absorbing atoms total absoption contribution to mu_t
##   mu_f:    total absorption at the fluo energy
##   d_given: thickness of material in cm
##   phi:     angle of incident beam wrt sample surface in rads
##   g:       sin(phi)/sin(theta) where theta is the angle of the
##            exit beam wrt the sample surface
##   option:  if 'u', undo correction.  Else make correction.
##            if correction is undone, chi_exp is really chi...
##
##   Correction should work in the thin as well as thick limits
def corr_chi(k,chi_exp,mu_t,mu_a,mu_f,d_given,phi,g,option):
    #set a maximum limit on thickness
    if d_given>0.01:
        d=0.01
    else:
        d=d_given
    alpha=mu_t+mu_f*g
    beta=alpha*mu_a*d*exp(-alpha*d/sin(phi))/sin(phi)
    gamma=1.0-exp(-alpha*d/sin(phi))
    #thick limit is more accurate as it includes the XAFS oscillations
    if lower(option)=='u':
        thick=chi_exp*(alpha-mu_a)/(alpha+chi_exp*mu_a)
    else:
        thick=chi_exp/(1.0-mu_a*(1.0+chi_exp)/alpha)
    #criteria for checking if roundoff error is starting to be a problem
    crit1=(4.0*beta*alpha*gamma*chi_exp)/(gamma*(alpha-mu_a*(chi_exp+1.0))+beta)
    crit2=beta

    if (abs(crit1)>1e-7 and abs(crit2)>1e-7 and d<=0.01):
        #otherwise use the "nearly exact" expression
        if lower(option)=='u':
            #this is exact...
            #chi=(1.0-dexp(-(alpha+chi_exp*mu_a)*d/sin(phi)))
            #chi=chi*(chi_exp+1)*alpha/(gamma*(alpha+chi_exp*mu_a))
            #chi=chi-1.0
            #this is from the same approx used in the else
            F=gamma*mu_a/(2*beta)
            E=-(gamma*(alpha-mu_a)+beta)/(2*beta)
            H=alpha*gamma/beta
            chi=(chi_exp*(chi_exp-2*E))/(2*F*chi_exp+H)
        else:
            chi=-(gamma*(alpha-mu_a*(chi_exp+1))+beta)
            chi=chi+sqrt(chi*chi+4*beta*alpha*gamma*chi_exp)
            chi=chi/(2*beta)
    else:
        #if roundoff potentially bad, do the thick limit
        chi=thick
        print """WARNING: Using 'thick limit' equation at k="""+str(k)
    #check main approximation
    reality=1.0-exp(-(alpha+chi*mu_a)*d/sin(phi))
    approx=1.0-exp(-alpha*d/sin(phi))*(1.0-chi*mu_a*d/sin(phi))
    error=abs((approx-reality)/reality)
    if (error>0.02 and k>1.5):
        if error>1.0: error=1.0
        print 'WARNING: integrand approx. only good to '+str(error*100)+' percent at k='+str(k)
    return chi

## GET MUS HERE
## USAGE:
##  [tot_abs,fluor_abs,atomic_abs]=get_mus(formula,edge,energy,
##                           fluor_energy,concentration,formvol)
## Returns mus for corrections
def get_mus(formula,edge,energy,fluor_energy,concentration,formvol):
    numbers='0123456789.'
    #ignoring CB "first" structure...
    k=0
    if concentration<0:
        atomic_sym='   '
        #isolate the atom from the kind of edge
        i=find(edge,' ')
        #find position of edge in formula
        j=find(formula,edge[0:i])
        #make sure number of atoms stay in atomic_sym
        x=0
        for m in formula[j+i:]:
            if m in numbers and x==0:
                k=k+1
            else:
                x=1
        atomic_sym=formula[j:j+i+k]
    else:
        atomic_sym=edge[0:2]
    #clean up end of atomic_sym with blanks... not needed?
    #atomic_sym=rstrip(atomic_sym)
    #atomic_sym=ljust(atomic_sym,2)
    fluor_abs=get_tot_abs(formula,fluor_energy,formvol)
    #correction needs the total absorption
    tot_abs=get_tot_abs(formula,energy,formvol)
    r1=get_tot_abs(atomic_sym,energy,formvol)
    r2=get_noedge_abs(atomic_sym,energy,formvol)
    if concentration<0:
        atomic_abs=r1-r2
    else:
        atomic_abs=concentration*(r1-r2)
    #diagnostic check
    r3=get_tot_abs(atomic_sym,fluor_energy,formvol)
    if ((r3-r1)/r1>0.03):
        print 'BIG TROUBLE'
        print 'DIAG: energy='+str(energy)
        print 'DIAG: r1='+str(r1)
        print 'DIAG: r2='+str(r2)
        print 'DIAG: r3='+str(r3)
    return [tot_abs,fluor_abs,atomic_abs]


## FLUOR CORRECTION
## USAGE:
##  corr_chi=fluor_corr(k,chi,d,phi,formvol,formula,edge,concentration,io,option)
##  returns fluo corrected chi array given k,chi and other params
##   k:             Sequence of k wavevectors
##   chi:           Sequence of fluo impacted chi meaurements
##   d:             Thickness of sample, in ANGSTROMS
##   phi:           Angle of incident beam wrt sample in degrees
##   formvol:       Volume of formula cell
##   formula:       Formula of compound of interest
##   edge:          Element and edge, eg. 'Mn K' or 'U LIII'
##   fluor_energy:  energy of fluo line -- taken from edge!
##   concentration: concentration (-1 by default)
##   io:            output option-
##                    'p' or 'P' for printout to stdoutput
##                    other for maximum verbosity
##   option:        if 'U' or 'u', will undo a correction, otherwise it'll do it
def fluor_corr(k,chi,d,phi,formvol,formula,edge,concentration,io,option):
    #initialize new array
    sabcor_chi=[]
    EV2INVANG=0.512393
    #need E0 for calculating edge energy
    [Z,edge_code,e_not]=atomic_sub.get_atomic(edge)
    fluor_energy=atomic_sub.get_fluo_energy(edge)    
    #convert angles to radians
    phi_rad=2*pi*phi/360
    theta_rad=pi/2-phi_rad
    g=sin(phi_rad)/sin(theta_rad)
    #begin correction loop
    for idat in range(len(chi)):
        #k is the wavevector
        #need it in terms of eV, since that is what the Vict. Coeffs are in terms of...
        if k[idat]>=0.0:
            energy=e_not+(k[idat]/EV2INVANG)**2
        else:
            print 'Trouble in fluo corr, k='+str(k[idat])
        #correction needs the total absorption
        [tot_abs,fluor_abs,atomic_abs]=get_mus(formula,edge,energy,fluor_energy,concentration,formvol)
        #separate out chi
        #  difference from CB version -- input is chi, NOT k-weighted chi
        new_chi=corr_chi(k[idat],chi[idat],tot_abs,atomic_abs,fluor_abs,d*1e-8,phi_rad,g,option)
        try:
            correction=new_chi/chi[idat]
        except:
            correction=1
        sabcor_chi.append(new_chi)
        temp=1e8*info_depth(tot_abs,fluor_abs,phi_rad,g)
        if k[idat]<=5.0:
            corr_at_5=correction
            info_at_5=temp
    #make sure new corr_chi is same length as old one        
    if len(sabcor_chi)!=len(chi):
        print 'SERIOUS ERROR -- CORRECTED CHI NOT SAME LENGTH'
    #stat report?
    if lower(io)=='p':
        print
        print 'Self-Absorption correction statistics:'
        print
        print 'info depth at 5 inv ang: '+str(info_at_5)
        print 'correction chi_true/chi_exp at 5 inv ang: '+str(corr_at_5)
        print        
    #return final corrected chi
    return sabcor_chi

#tests

#get_mus('ZnS','Zn K',10000,8500,-1,10)
#get_mus('ZnS','Zn K',10000,8500,1,10)
#get_mus('ZnS','S K',2500,2308,-1,10)
#get_mus('YBa2Fe0.31Cu2.69O7','Fe K',7800,6500,-1,10)

#print get_tot_abs('ZnS',10000,1)
#print tot_weight('NaCl')
#print atomic_sub.get_atomic('Mn K')
#form2sym('ZnS')
#form2sym('Ni2O4')
#print form2sym('YBa2Fe0.31Cu2.69O7')
#print atomic_weight(1)
#print atomic_sub.sym2Z('U')
#print atomic_weight(atomic_sub.sym2Z('U '))
#print fluor_corr([1,2,3,4,5],[1,1,1,1,1],45000,45,50,'MnO2','Mn K',-1,'v','x')
