#!/usr/bin/env python3
"""
Polarized neutron reflectometry of a 4x Fe/Ni bilayer on Au substrate
with ErfTransient (Nevot-Croce) interface roughness.

Four-channel (PP, PM, MP, MM) angle-scan. The Fe magnetization is
rotated 38 deg from the Y axis, enabling spin-flip scattering (PM, MP
non-zero).
See PolarizedFeNiBilayerNC.py for the two-channel variant
(0 deg tilt, PM=MP=0).

Compare PolarizedFeNiBilayerSpinFlipNCQ.py for the Q-space
equivalent.
"""
from math import cos, sin
import bornagain as ba
from bornagain import angstrom, ba_plot as bp, deg, nm, R3


def get_sample():
    vacuum = ba.Vacuum()
    mag_vec = R3(sin(38*deg), cos(38*deg), 0) * 1e7
    fe_color = (0.48, 0.32, 0.80)
    fe_mat = ba.SLDMaterial("Fe", fe_color, 8.02e-6, 0, mag_vec)
    ni_color = (0.93, 0.48, 0.14)
    ni_mat = ba.SLDMaterial("Ni", ni_color, 9.4245e-6, 0)
    substrate_color = (0.93, 0.72, 0.25)
    substrate_mat = ba.SLDMaterial("Au", substrate_color, 4.6665e-6, 0)

    autocorr = ba.SelfAffineFractalModel(0.2*nm, 0.7, 25*nm)
    roughness = ba.Roughness(autocorr, ba.ErfTransient())

    stack = ba.LayerStack(4)
    stack.addLayer(ba.Layer(fe_mat, 10*nm, roughness))
    stack.addLayer(ba.Layer(ni_mat, 4*nm, roughness))

    sample = ba.Sample()
    sample.addLayer(ba.Layer(vacuum))
    sample.addStack(stack)
    sample.addLayer(ba.Layer(substrate_mat, roughness))
    return sample


def simulate(sample, pol, an, title):
    n = 500
    scan = ba.AlphaScan(n, 5*deg/n, 5*deg)
    scan.setWavelength(1.54*angstrom)
    scan.setPolarization(pol)
    scan.setAnalyzer(an)
    result = ba.SpecularSimulation(scan, sample).simulate()
    result.setTitle(title)
    return result


if __name__ == '__main__':
    sample = get_sample()
    results = [
        simulate(sample, R3(0, +1, 0), R3(0, +1, 0), "PP"),
        simulate(sample, R3(0, +1, 0), R3(0, -1, 0), "PM"),
        simulate(sample, R3(0, -1, 0), R3(0, +1, 0), "MP"),
        simulate(sample, R3(0, -1, 0), R3(0, -1, 0), "MM"),
    ]
    plotargs = bp.parse_commandline()
    bp.plot_multicurve(results, **plotargs)
    bp.export(**plotargs)
