#!/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
    material_D2O = ba.RefractiveMaterial("D2O", 0.00010116, 1.809e-12)
    material_Pt = ba.RefractiveMaterial("Pt", 0.00010117, 3.01822e-08)
    material_Si = ba.RefractiveMaterial("Si", 3.3009e-05, 0)
    material_Ti = ba.RefractiveMaterial("Ti", -3.0637e-05, 1.5278e-08)
    material_TiO2 = ba.RefractiveMaterial("TiO2", 4.1921e-05, 8.1293e-09)

    # Layers
    layer_1 = ba.Layer(material_Si)
    layer_2 = ba.Layer(material_Ti, 13*nm)
    layer_3 = ba.Layer(material_Pt, 32*nm)
    layer_4 = ba.Layer(material_Ti, 10*nm)
    layer_5 = ba.Layer(material_TiO2, 3*nm)
    layer_6 = ba.Layer(material_D2O)

    # 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()
    plotargs['aspect'] = 'auto'
    bp.plot_simulation_result(result, **plotargs)
    bp.export(**plotargs)
