#!/usr/bin/env python3
"""
Flattened depth-probe simulation:
Intensity as function of depth z for a few incident angles alpha_i.
"""

import bornagain as ba
from bornagain import ba_plot as bp, deg, nm


def get_sample():
    material_vac = ba.RefractiveMaterial("Vacuum", 0, 0)
    material_A = ba.RefractiveMaterial("A", 6e-5, 0)
    material_sub = ba.RefractiveMaterial("Substrate", 3e-05, 0)

    sample = ba.Sample()
    sample.addLayer(ba.Layer(material_vac))
    sample.addLayer(ba.Layer(material_A, 100*nm))
    sample.addLayer(ba.Layer(material_sub))

    return sample

def simulate(sample, alpha_i_deg):
    alpha = alpha_i_deg * deg
    scan = ba.AlphaScan(1, alpha, alpha)
    scan.setWavelength(0.3*nm)

    z_axis = ba.EquiDivision("z (nm)", 500, -130*nm, 30*nm)
    simulation = ba.DepthprobeSimulation(scan, sample, z_axis, 0)

    result = simulation.simulate().flat()
    result.setTitle(f'{alpha_i_deg} deg')
    return result

if __name__ == '__main__':
    sample = get_sample()
    angles = [0.4, 0.5, 0.6, 0.65, 0.7]
    results = [simulate(sample, ai) for ai in angles]
    plotargs = bp.parse_commandline()
    plotargs['legendloc'] = 'lower right'
    bp.plot_multicurve(results, **plotargs)
    bp.plt.show()
