#!/usr/bin/env python3
"""
Polarized neutron reflectometry of a 4x Fe/Ni bilayer on Au substrate,
in Q-space.

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

Results must agree with PolarizedFeNiBilayerSpinFlip.py (angle-scan).
Compare PolarizedFeNiBilayerSpinFlipTanhQ.py for the roughness variant.
"""
from math import cos, pi, sin
import numpy as np
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)

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

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


def simulate(sample, pol, an, title):
    n = 50
    lam = 1.54*angstrom
    qs = [4*pi*sin(a)/lam
          for a in np.linspace(5*deg/n, 5*deg, n)]
    scan = ba.QzScan(qs)
    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"),
    ]
    from bornagain import ba_check
    ba_check.persistence_test(results)
