#!/usr/bin/env python3
"""
Basic example of depth-probe simulation:
Computes intensity as function of incident angle alpha and depth z.
"""
import bornagain as ba
from bornagain import ba_plot as bp, deg, nm

def get_sample():
    # Materials
    vac_mat = ba.Vacuum()
    a_color = (0.05, 0.62, 0.55)
    a_mat = ba.RefractiveMaterial("A", a_color, 6e-5, 0)
    sub_color = (0.28, 0.57, 0.82)
    sub_mat = ba.RefractiveMaterial("Substrate", sub_color, 3e-05, 0)

    # Sample
    sample = ba.Sample()
    sample.addLayer(ba.Layer(vac_mat))
    sample.addLayer(ba.Layer(a_mat, 100*nm))
    sample.addLayer(ba.Layer(sub_mat))

    return sample


def get_simulation(sample, flags):
    n = 50

    scan = ba.AlphaScan(n, 0*deg, 1*deg)
    scan.setWavelength(0.3*nm)

    z_axis = ba.EquiDivision("z (nm)", n, -130*nm, 30*nm)
    simulation = ba.DepthprobeSimulation(scan, sample, z_axis, flags)
    # BUG The following line should be unnecessary:
    simulation.options().setUseAvgMaterials(False)

    return simulation


def run_example(flags=0):
    sample = get_sample()
    simulation = get_simulation(sample, flags)
    result = simulation.simulate()
    from bornagain import ba_check
    ba_check.persistence_test(result)
    return


if __name__ == '__main__':
    run_example()
