#!/usr/bin/env python3
"""
Print transmission of a multilayer,
using a flattened depth-probe simulation.
"""

import bornagain as ba
from bornagain import deg, nm

depth = 100 * 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.MultiLayer()
    sample.addLayer(ba.Layer(material_vac))
    sample.addLayer(ba.Layer(material_A, depth))
    sample.addLayer(ba.Layer(material_sub))

    return sample

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

    z_axis = ba.EquiDivision("z (nm)", 1, -depth, -depth)
    simulation = ba.DepthprobeSimulation(scan, get_sample(), z_axis, 0)

    result = simulation.simulate().flat()
    assert result.size() == 1
    return result.valAt(0)

if __name__ == '__main__':
    print(simulate())
