#!/usr/bin/env python3
"""
Fresnel transmittivity and reflectivity of a multilayer sample.

Uses DepthprobeSimulation.transmitted() and .reflected().
"""

import bornagain as ba
from bornagain import deg, nm

depth = 100 * nm


def get_sample():
    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 = 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():
    alpha = 0.64 * deg
    scan = ba.AlphaScan(1, alpha, alpha)
    scan.setWavelength(0.3 * nm)

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

    T = simulation.transmitted().flat().valAt(0)
    R = simulation.reflected().flat().valAt(0)
    return T, R


if __name__ == '__main__':
    ba.showSample3D(get_sample(), sample_size=180*nm, seed=0)
    T, R = simulate()
    print(f'transmittivity: {T}')
    print(f'reflectivity:   {R}')
