#!/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

# mirror reflectivity calculations
# inputs from web form
formula = 'Pt'  # mirror chemical formula
density = 21.4500  # mirror density in gr/cm^3
angle   = 2.5000  # mirror angle in mrad
rough   = 10.0000  # mirror roughness in Angstroms
polar   = 's'  # mirror polarization ('s' for vert deflecting with horiz-polarized source)
energy = np.arange(1000, 51000+50, 50)


reflectivity = xraydb.mirror_reflectivity(formula, 0.001*angle, energy, density,
                                          roughness=rough, polarization=polar)

delta, beta, _ = xraydb.xray_delta_beta(formula, density, energy)
ang_crit = 1000*np.arccos(1 - delta - 1j*beta).real

plt.plot(energy, reflectivity, label='reflectivity')
plt.xlabel('Energy (eV)')
plt.ylabel('Reflectivity')
plt.title('X-ray Reflectivity for %s' % formula)
plt.show()

plt.plot(energy, ang_crit, label='Critical Angle')
plt.xlabel('Energy (eV)')
plt.ylabel('Critical Angle (mrad)')
plt.title('Critical Angle for %s' % formula)
plt.show()