#!/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
from matplotlib import pyplot as plt
import bornagain as ba
from bornagain import angstrom, ba_plot as bp, deg, std_samples

# 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.setAngleDistribution(alpha_distr)

    return ba.SpecularSimulation(scan, sample)


if __name__ == '__main__':
    datadir = os.getenv('BA_DATA_DIR', '')
    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)

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

    sample = get_sample()
    simulation = get_simulation(sample)
    result = simulation.simulate()
    bp.plot_simulation_result(result)
    plt.legend(['GenX', 'BornAgain'], loc='upper right')
    bp.show_or_export()
