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

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


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

    sample = ba.Sample()
    sample.addLayer(ba.Layer(vac_mat))
    sample.addLayer(ba.Layer(a_mat, depth))
    sample.addLayer(ba.Layer(sub_mat))

    return sample

def simulate(depth):
    sample = get_sample(depth)
    n = <%= test_mode ? 20 : 500 %>
    alpha_max = 0.8 * deg
    scan = ba.AlphaScan(n, alpha_max / n, alpha_max)
    scan.setWavelength(0.3*nm)

    z_axis = ba.ListScan("z (nm)", [-depth])
    simulation = ba.DepthprobeSimulation(scan, sample, z_axis, 0)

    result = simulation.simulate().flat()
    result.setTitle(f'{depth} nm')
    return result

if __name__ == '__main__':
    depths = [20, 40, 60]
    results = [simulate(d) for d in depths]
    <%- if test_mode -%>
    from bornagain import ba_check
    ba_check.persistence_test(result)
    <%- elsif figure_mode -%>
    plotargs = bp.parse_commandline()
    bp.plot_multicurve(results, legendloc='lower right', **plotargs)
    bp.export(**plotargs)
    <%- else -%>
    ba.showSample3D(get_sample(depths[0]), sample_size=120*nm, seed=0)
    bp.plot_multicurve(results, legendloc='lower right')
    bp.plt.show()
    <%- end -%>
