#!/usr/bin/env python3
"""
Basic example of depth-probe simulation with BornAgain.

Sample layers are Si | Ti | Pt | Ti | TiO2 | D2O.
Beam comes from Si side.
Therefore we model the stack with Si on top.
The z axis points from D2O to Si; z=0 is at the Si/Ti interface.
"""
import bornagain as ba
from bornagain import angstrom, ba_plot as bp, deg, nm


# layer thicknesses in angstroms
t_Ti = 130*angstrom
t_Pt = 320*angstrom
t_Ti_top = 100*angstrom
t_TiO2 = 30*angstrom

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

# convolution parameters
d_ang = 0.01*ba.deg  # spread width for incident angle

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


def get_sample():
    """
    Constructs a sample with one resonating Ti/Pt layer
    """

    # Materials
    d2o_color = (0.90, 0.93, 0.97)
    d2o_mat = ba.RefractiveMaterial("D2O", d2o_color, 0.00010116, 1.809e-12)
    pt_color = (0.93, 0.72, 0.25)
    pt_mat = ba.RefractiveMaterial("Pt", pt_color, 0.00010117, 3.01822e-08)
    si_color = (0.30, 0.62, 0.86)
    si_mat = ba.RefractiveMaterial("Si", si_color, 3.3009e-05, 0)
    ti_color = (0.05, 0.62, 0.55)
    ti_mat = ba.RefractiveMaterial("Ti", ti_color, -3.0637e-05, 1.5278e-08)
    tio2_color = (0.48, 0.32, 0.80)
    tio2_mat = ba.RefractiveMaterial("TiO2", tio2_color, 4.1921e-05, 8.1293e-09)

    # Layers
    layer_1 = ba.Layer(si_mat)
    layer_2 = ba.Layer(ti_mat, 13*nm)
    layer_3 = ba.Layer(pt_mat, 32*nm)
    layer_4 = ba.Layer(ti_mat, 10*nm)
    layer_5 = ba.Layer(tio2_mat, 3*nm)
    layer_6 = ba.Layer(d2o_mat)

    # Sample
    sample = ba.Sample()
    sample.addLayer(layer_1)
    sample.addLayer(layer_2)
    sample.addLayer(layer_3)
    sample.addLayer(layer_4)
    sample.addLayer(layer_5)
    sample.addLayer(layer_6)

    return sample


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

    scan = ba.AlphaScan(na, ai_min, ai_max)
    scan.setWavelength(wl)

    alpha_distr = ba.DistributionGaussian(0, d_ang, 25, 3.)
    scan.setGrazingAngleDistribution(alpha_distr)

    z_axis = ba.EquiDivision("z (nm)", 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()
    plotargs = bp.parse_commandline()
    bp.plot_datafield(result, frame_aspect=1.618, **plotargs)
    bp.export(**plotargs)
