#!/usr/bin/env python3
"""
Depth profile of an opaque sample.
"""
import bornagain as ba
from bornagain import ba_plot as bp, deg, nm

#  beam data
ai_min = 0  # minimum incident angle
ai_max = 1*deg  # maximum incident angle
wl = 0.03*nm  # wavelength

#  depth position span
z_min = -400*nm
z_max = 100*nm


def get_sample():
    # Define materials
    vac_color = (0.90, 0.93, 0.97)
    vac_mat = ba.RefractiveMaterial("Vac", vac_color, 0, 0)
    a_color = (0.05, 0.62, 0.55)
    a_mat = ba.RefractiveMaterial("A", a_color, 1e-5, 3e-6)

    # Define layers
    layer_top = ba.Layer(vac_mat)
    layer_1 = ba.Layer(a_mat, 300*nm)
    layer_bot = ba.Layer(vac_mat)

    # Define sample
    sample = ba.Sample()
    sample.addLayer(layer_top)
    sample.addLayer(layer_1)
    sample.addLayer(layer_bot)

    return sample


def get_simulation(sample):
    """
    Returns a depth-probe simulation.
    """
    nz = 500
    na = 5000

    scan = ba.AlphaScan(na, ai_min, ai_max)
    scan.setWavelength(wl)
    footprint = ba.FootprintSquare(0.01)
    scan.setFootprint(footprint)

    z_axis = ba.EquiDivision("z", nz, z_min, z_max)
    simulation = ba.DepthprobeSimulation(scan, sample, z_axis)

    return simulation


if __name__ == '__main__':
    sample = get_sample()
    simulation = get_simulation(sample)
    result = simulation.simulate()
    <%- if test_mode -%>
    from bornagain import ba_check
    ba_check.persistence_test(result)
    <%- else -%>
    <%- if figure_mode -%>
    plotargs = bp.parse_commandline()
    <%- else -%>
    plotargs = {}
    <%- end -%>
    bp.plot_datafield(result, intensity_min=2.5e-91, intensity_max=4,
                      frame_aspect=1.618, **plotargs)
    # Customize colorbar ticks
    import numpy as np
    fig = bp.plt.gcf()
    if len(fig.axes) > 1:
        cbar_ax = fig.axes[1]  # Colorbar axes
        ticks = [1, 1e-15, 1e-30, 1e-45, 1e-60, 1e-75, 1e-90]
        cbar_ax.yaxis.set_ticks(ticks)
    <%- if figure_mode -%>
    bp.export(**plotargs)
    <%- else -%>
    ba.showSample3D(sample, sample_size=500*nm, seed=0)
    bp.plt.show()
    <%- end -%>
    <%- end -%>
