#!/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 matplotlib.pyplot as plt
import bornagain as ba
from bornagain import angstrom, 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.MaterialBySLD("Vacuum", 0, 0)
    material_Pd = ba.MaterialBySLD("Pd", 4.0099e-6, 1.3019e-09)
    material_Fe = ba.MaterialBySLD("Fe", 8.0241e-06, 6.0448e-10, B)
    material_substrate = ba.MaterialBySLD("MgO", 5.9803e-06, 9.3996e-12)

    # Layers
    layer_vacuum = ba.Layer(vacuum)
    layer_Pd = ba.Layer(material_Pd, 120*angstrom)
    layer_Fe = ba.Layer(material_Fe, 1000*angstrom)
    layer_substrate = ba.Layer(material_substrate)

    roughness = ba.LayerRoughness(20*angstrom)

    # Multilayer
    sample = ba.MultiLayer()
    sample.addLayer(layer_vacuum)
    sample.addLayerWithTopRoughness(layer_Pd, roughness)
    sample.addLayerWithTopRoughness(layer_Fe, roughness)
    sample.addLayerWithTopRoughness(layer_substrate, roughness)

    return sample


def simulate(p_dir, a_dir, p_eff, a_eff, 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_eff)

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

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

    return result


if __name__ == '__main__':
    # polarizer and analyzer efficiencies
    p_eff = 0.986
    a_eff = 0.970

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

    results = [results_pp, results_pm, results_mp, results_mm]
    bp.plot_multicurve_specular(results)
