#!/usr/bin/env python3
"""
Reflectivity of a multilayer, taking into account beam angular divergence
and beam footprint correction, simulated with BornAgain and GenX.
"""
import numpy as np, os, sys
import bornagain as ba
from bornagain import angstrom, ba_plot as bp, deg, std_samples, nm

# input parameters
wavelength = 1.54*angstrom
beam_sample_ratio = 0.01  # beam-to-sample size ratio

def reference_data(filename):
    """
    Loads and returns reference data from GenX simulation
    """
    ax_values, data = np.loadtxt(filename,
                                      usecols=(0, 1),
                                      skiprows=3,
                                      unpack=True)

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

    return ax_values, data


def get_sample():
    return std_samples.alternating_layers()


def get_simulation(sample, **kwargs):
    """
    A specular simulation with beam and detector defined.
    """
    n = 50
    footprint = ba.FootprintSquare(beam_sample_ratio)
    alpha_distr = ba.DistributionGaussian(0, 0.01 * deg, 25, 3.)

    scan = ba.AlphaScan(n, 2*deg/n, 2*deg)
    scan.setWavelength(1.54*angstrom)
    scan.setFootprint(footprint)
    scan.setGrazingAngleDistribution(alpha_distr)

    return ba.SpecularSimulation(scan, sample)


if __name__ == '__main__':
    datadir = os.getenv('BA_DATA_DIR', '')
    if not datadir:
        raise Exception("Environment variable BA_DATA_DIR not set")
    data_fname = os.path.join(datadir, "specular/genx_angular_divergence.dat.gz")
    print(f"Loading GenX reference data from {data_fname}")
    genx_axis, genx_values = reference_data(data_fname)

    bp.plt.yscale('log')
    bp.plt.plot(genx_axis, genx_values, 'ko', markevery=300)

    sample = get_sample()
    simulation = get_simulation(sample)
    result = simulation.simulate()
    from bornagain import ba_check
    ba_check.persistence_test(result)
