#!/usr/bin/env python3
"""
Realistic (imperfect) example for polarized reflectometry.
Sample contains a magnetic layer,
similar to Devishvili et al., Rev. Sci. Instrum. 84, 025112 (2013).
"""
import numpy
import bornagain as ba
from bornagain import nm, ba_plot as bp, deg, R3
from math import sin, cos

def get_sample():
    # Materials
    Bmag = 1.6e6
    Bangle = 0 * deg
    B = R3(Bmag*sin(Bangle), Bmag*cos(Bangle), 0)
    vacuum = ba.Vacuum()
    pd_color = (0.93, 0.72, 0.25)
    pd_mat = ba.SLDMaterial("Pd", pd_color, 4.0099e-6, 1.3019e-09)
    fe_color = (0.48, 0.32, 0.80)
    fe_mat = ba.SLDMaterial("Fe", fe_color, 8.0241e-06, 6.0448e-10, B)
    substrate_color = (0.28, 0.57, 0.82)
    substrate_mat = ba.SLDMaterial("MgO", substrate_color, 5.9803e-06, 9.3996e-12)

    autocorr = ba.SelfAffineFractalModel(2*nm, 0.7, 25*nm)
    transient = ba.TanhTransient()
    roughness = ba.Roughness(autocorr, transient)

    # Layers
    layer_vacuum = ba.Layer(vacuum)
    layer_Pd = ba.Layer(pd_mat, 12*nm, roughness)
    layer_Fe = ba.Layer(fe_mat, 100*nm, roughness)
    layer_substrate = ba.Layer(substrate_mat, roughness)

    # Multilayer
    sample = ba.Sample()
    sample.addLayer(layer_vacuum)
    sample.addLayer(layer_Pd)
    sample.addLayer(layer_Fe)
    sample.addLayer(layer_substrate)

    return sample


def simulate(p_dir, a_dir, p_eff, a_trans, a_supp, title):
    sample = get_sample()

    qzs = numpy.linspace(0.1, 1.5, 1500)
    distr = ba.DistributionGaussian(0., 1., 25, 4.)

    scan = ba.QzScan(qzs)
    scan.setAbsoluteQResolution(distr, 0.008)

    scan.setPolarization(p_dir*p_eff)
    scan.setAnalyzer(a_dir, a_trans, a_supp)

    simulation = ba.SpecularSimulation(scan, sample)
    simulation.setBackground(ba.ConstantBackground(1e-7))

    result = simulation.simulate()
    result.setTitle(title)

    return result


if __name__ == '__main__':
    ba.showSample3D(get_sample(), sample_size=80*nm, seed=0)
    # polarizer and analyzer efficiencies
    p_eff = 0.986
    a_trans = 0.970
    a_supp = 0.015

    results_pp = simulate(R3(0, +1, 0), R3(0, +1, 0), p_eff, a_trans, a_supp, "$++$")
    results_pm = simulate(R3(0, +1, 0), R3(0, -1, 0), p_eff, a_trans, a_supp, "$+-$")
    results_mp = simulate(R3(0, -1, 0), R3(0, +1, 0), p_eff, a_trans, a_supp, "$-+$")
    results_mm = simulate(R3(0, -1, 0), R3(0, -1, 0), p_eff, a_trans, a_supp, "$--$")

    results = [results_pp, results_pm, results_mp, results_mm]
    bp.plot_multicurve(results)
    bp.plt.show()
