#!/usr/bin/env python
# this script requires Python3, numpy, scipy, matplotlib, and xraydb modules. Use:
#        pip install xraydb scipy matplotlib
import numpy as np
import matplotlib.pyplot as plt
import xraydb

# X-ray atomic scattering factors
# inputs from web form
elem   = 'I'
energy = np.arange(1000, 50000+50, 50)

mu_total = xraydb.mu_elam(elem, energy, kind='total')
mu_photo = xraydb.mu_elam(elem, energy, kind='photo')
mu_incoh = xraydb.mu_elam(elem, energy, kind='incoh')
mu_coher = xraydb.mu_elam(elem, energy, kind='coh')

plt.plot(energy, mu_total, label='Total')
plt.plot(energy, mu_photo, label='Photo-electric')
plt.plot(energy, mu_incoh, label='Incoherent')
plt.plot(energy, mu_coher, label='Coherent')
plt.xlabel('Energy (eV)')
plt.ylabel(r'$\mu/\rho \rm\,(cm^2/gr)$')
plt.legend()
plt.yscale('log')
plt.title('Mass Attenuation for I')
plt.show()

atz = xraydb.atomic_number(elem)
if atz < 93:
    f1 = xraydb.f1_chantler(elem, energy)
    f2 = xraydb.f2_chantler(elem, energy)
    plt.plot(energy, f1, label='f1')
    plt.plot(energy, f2, label='f2')
    plt.xlabel('Energy (eV)')
    plt.ylabel('f1, f2 (electrons/atom)')
    plt.title('Resonant Scattering factors for I')
    plt.legend(True)
    plt.show()
else:
    print('f1 and f2 are only available for Z<93')