#!/usr/bin/env python3
"""
Polarized neutron reflectometry of a single magnetic layer on a
substrate. TanhTransient interface roughness.

Four-channel (PP, PM, MP, MM) angle-scan. The layer magnetization
is tilted 45 deg from the Y axis (R3(1, 1, 0) direction), enabling
spin-flip scattering (PM, MP non-zero).
"""
import bornagain as ba
from bornagain import angstrom, ba_plot as bp, deg, nm, R3


def get_sample():
    layer_field = R3(1e6, 1e6, 0)
    substr_field = R3(0, 1e6, 0)
    layer_color = (0.86, 0.24, 0.18)
    layer_mat = ba.RefractiveMaterial(
        "MagLayer", layer_color, 6e-4, 2e-8, layer_field)
    substrate_color = (0.28, 0.57, 0.82)
    substrate_mat = ba.RefractiveMaterial(
        "Substrate", substrate_color, 7e-6, 2e-8, substr_field)

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

    sample = ba.Sample()
    sample.addLayer(ba.Layer(ba.Vacuum()))
    sample.addLayer(ba.Layer(layer_mat, 20*nm, roughness))
    sample.addLayer(ba.Layer(substrate_mat, roughness))
    return sample


def simulate(sample, pol, an, title):
    n = 50
    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"),
    ]
    from bornagain import ba_check
    ba_check.persistence_test(results)
