# 2/17/2011: readf1f2a.py
# read atomic scattering factors
# data file: f1f2chantler.py
# f1, f2, mu_photoelectric, mu_total
# mu_total should be used for filter transmission
# get_delta returns la for total crosssection as well

from math import pi

from larch import use_plugin_path
use_plugin_path('xray')

from chemparser import chemparse
from xraydb_plugin import atomic_mass, atomic_number, chantler_data

from physical_constants import R_ELECTRON_CM, AVOGADRO, BARN, PLANCK_HC

KNOWN_DENSITIES = {'Ag': 10.5 ,
                   'Al': 2.72 ,
                   'Al2O3': 3.97,  # Sapphire
                   'Au': 19.37 ,
                   'Be': 1.848 ,
                   'BN': 2.29, # Boron Nitride
                   'C': 3.51 , # diamond
                   'C_graphite': 2.25, # graphite
                   'C2F4': 2.2 ,	# Teflon
                   'C3H4O2': 1.65 ,  # PAA
                   'C22H10O4N2': 1.42,	 # Kapton
                   'CaCO3': 2.71 ,     # Calcite
                   'CaF2': 3.18,	 # Fluorite
                   'Cr': 7.19 ,
                   'Co': 8.9 ,
                   'Cu': 8.94 ,
                   'Fe': 7.86,
                   'Fe2O3': 5.26  , # hematite
                   'Ge': 5.323 ,
                   'H2O': 1.0 ,
                   'Ir': 22.42 ,
                   'KAl3Si3O12H2': 2.83 ,  # Mica
                   'Mn': 7.42 ,
                   'Ni': 8.9 ,
                   'Pb': 11.34 ,
                   'PbTiO3': 7.9 ,
                   'Pd': 12.16 ,
                   'Pt': 21.37 ,
                   'Rh': 12.44 ,
                   'Si': 2.33 ,
                   'SiO2': 2.2 , # silica
                   'Sr': 2.54 ,
                   'SrTiO3': 5.12 ,
                   'Ti': 4.54 ,
                   'TiO2': 4.26,   # Rutile
                   'V': 6.11 ,
                   'Zn': 7.14 ,
                   'ZnO': 5.7,
                   }

def nominal_density(composition):
    "returns nominal density in g/cc for known compounds"
    return KNOWN_COMPOUNDS.get(composition, 1.0)

class Scatterer:
    """Scattering Element

    lamb=PLANCK_HC /(eV0/1000.)*1e-11 # in cm, 1e-8cm = 1 Angstrom
    Xsection=2* R_ELECTRON *lamb*f2/BARN    # in Barns/atom
    """
    def __init__(self, symbol, energy=10000, _larch=None):
        # atomic symbol and incident x-ray energy (eV)
        self.symbol  = symbol
        self.number  = atomic_number(symbol, _larch=_larch)
        self.mass     = atomic_mass(symbol, _larch=_larch)
        self.f1       = chantler_data(symbol, energy, 'f1', _larch=_larch)
        self.f1       = self.f1 + self.number
        self.f2       = chantler_data(symbol, energy, 'f2', _larch=_larch)
        self.mu_photo = chantler_data(symbol, energy, 'mu_photo', _larch=_larch)
        self.mu_total = chantler_data(symbol, energy, 'mu_total', _larch=_larch)


def xray_delta_beta(material, density, energy, photo_only=False, _larch=None):
    """
    calculate scattering factors for a material (chemical formula)

    arguments:
    ----------
       material:   chemical formula  ('Fe2O3', 'CaMg(CO3)2', 'La1.9Sr0.1CuO4')
       density:    material density in g/cm^3
       energy:     x-ray energy in eV
       photo_only: boolean for returning photo cross-section component only
                   if False (default), the total cross-section is returned
    outputs:
    ---------
      (delta, beta, atlen)

    where
      delta :  real part of index of refraction
      beta  :  imag part of index of refraction
      atlen :  attenuation length in cm

    n = 1 - delta - i*beta = 1 - lamdba**2 * r0/(2*pi) Sum_j (n_j * fj)

    Yong Choi
    """
    lamb_cm = 1.e-8 * PLANCK_HC / energy # lambda in cm
    elements = []
    for symbol, number in chemparse(material).items():
        elements.append((number, Scatterer(symbol, energy, _larch=_larch)))

    total_mass, delta, beta_photo, beta_total = 0, 0, 0, 0
    for (number, scat) in elements:
        weight      = density*number*AVOGARDO
        delta      += weight * scat.f1
        beta_photo += weight * scat.f2
        beta_total += weight * scat.f2*(scat.mu_total/scat.mu_phot)
        total_mass += number * scat.mass

    scale = lamb_cm * lamb_cm * R_ELECTRON_CM / (2*pi * total_mass)
    delta = delta * scale
    beta  = beta_total * scale
    if photo_only:
        beta  = beta_photo * scale
    return delta, beta, lamb_cm/(4*pi*beta)

if __name__=='__main__':
    option=0
    if option==0:
        print AtNum2f1f2Xsect(26, 7000)
        print AtNum2f1f2Xsect(26, 8000)
    if option==1:
        import time
        list_eV=range(500, 30000, 100)
        # below takes 2.72 seconds
        a=time.clock()
        for eV0 in list_eV:
            AtNum2f1f2Xsect(26, eV0)
        print AtNum2f1f2Xsect(13, 35000.)
        b=time.clock()-a
        print b
    #-------------------------------------
    if option==2:
        eV0=30014.05
        temp=AtNum2f1f2Xsect(13, eV0)
        print temp[0], temp[1], temp[2], temp[3], temp[4]
        temp=get_delta('Al2O3', 3.97, 35000)
        print temp
    #-------------------------------------
    if option==3:
        import numpy
        energies=[]
        atoms=['Si', 'Fe', 'Co', 'Ni', 'Gd', 'Ta', 'Au']
        for i in numpy.arange(500., 1500., 0.2):
            energies.append(i)
        get_ScatFact(atoms, energies, 'scatfact')
    #-------------------------------------
    if option==4:
        import numpy
        energies=[]
        atoms=['Fe','N', 'O', 'Ga', 'As', 'Au']
        for i in numpy.arange(7000., 8100., 1):
            energies.append(i)
        get_ScatFact(atoms, energies, 'scatfact')
    #-------------------------------------
    if option==5:
        eV0=8800.0
        out=AtNum2f1f2Xsect(25, eV0)
        print out[2]
        out=AtNum2f1f2Xsect(26, eV0)
        print out[2]
