#!/usr/bin/env python3
"""
An example of taking into account beam angular divergence
and beam footprint correction in reflectometry calculations.
"""
import numpy as np
import bornagain as ba
from bornagain import angstrom, deg
from os import path
from matplotlib import pyplot as plt

# input parameters
wavelength = 1.54*angstrom
alpha_i_min = 0  # min incident angle, deg
alpha_i_max = 2*deg  # max incident angle, rad
beam_sample_ratio = 0.01  # beam-to-sample size ratio

# convolution parameters
d_ang = 0.01*deg  # spread width for incident angle
n_sig = 3  # number of sigmas to convolve over
n_points = 25  # number of points to convolve over

# substrate (Si)
si_sld_real = 2.0704e-06  # \AA^{-2}
# layer parameters
n_repetitions = 10
# Ni
ni_sld_real = 9.4245e-06  # \AA^{-2}
d_ni = 70*angstrom
# Ti
ti_sld_real = -1.9493e-06  # \AA^{-2}
d_ti = 30*angstrom


def get_sample():
    # defining materials
    m_vacuum = ba.MaterialBySLD("Vacuum", 0, 0)
    m_ni = ba.MaterialBySLD("Ni", ni_sld_real, 0)
    m_ti = ba.MaterialBySLD("Ti", ti_sld_real, 0)
    m_substrate = ba.MaterialBySLD("SiSubstrate", si_sld_real, 0)

    vacuum_layer = ba.Layer(m_vacuum)
    ni_layer = ba.Layer(m_ni, d_ni)
    ti_layer = ba.Layer(m_ti, d_ti)
    substrate_layer = ba.Layer(m_substrate)
    multi_layer = ba.MultiLayer()
    multi_layer.addLayer(vacuum_layer)
    for i in range(n_repetitions):
        multi_layer.addLayer(ti_layer)
        multi_layer.addLayer(ni_layer)
    multi_layer.addLayer(substrate_layer)
    return multi_layer


def create_real_data():
    """
    Loading data from genx_angular_divergence.dat
    """
    filepath = path.join(path.dirname(path.realpath(__file__)),
                         "genx_angular_divergence.dat.gz")
    ax_values, real_data = np.loadtxt(filepath,
                                      usecols=(0, 1),
                                      skiprows=3,
                                      unpack=True)

    # translating axis values from double incident angle # to incident angle
    ax_values *= 0.5

    return ax_values, real_data


def get_simulation(sample, scan_size=500):
    """
    Returns a specular simulation with beam and detector defined.
    """
    footprint = ba.FootprintSquare(beam_sample_ratio)
    alpha_distr = ba.RangedDistributionGaussian(n_points, n_sig)

    scan = ba.AngularSpecScan(wavelength, scan_size, alpha_i_min,
                              alpha_i_max)
    scan.setFootprintFactor(footprint)
    scan.setAbsoluteAngularResolution(alpha_distr, d_ang)

    simulation = ba.SpecularSimulation()
    simulation.setScan(scan)
    simulation.setSample(sample)

    return simulation


if __name__ == '__main__':
    import ba_plot
    sample = get_sample()
    simulation = get_simulation(sample)
    ba_plot.run_and_plot(simulation)
    genx_axis, genx_values = create_real_data()
    plt.semilogy(genx_axis, genx_values, 'ko', markevery=300)
    plt.legend(['BornAgain', 'GenX'], loc='upper right')
    plt.show()
