#!/usr/bin/env python3
"""
Introductory example for polarized neutron reflectivity.
Sample is a magnetic layer.
"""
import bornagain as ba
from bornagain import angstrom, ba_plot as bp, deg, nm, R3
from math import sin, cos


def get_sample():
    # Magnetic field
    Bmag = 1e8
    Bangle = 60*deg
    B = R3(Bmag*cos(Bangle), Bmag*sin(Bangle), 0)

    # Materials
    vacuum = ba.MaterialBySLD("Vacuum", 0, 0)
    material_layer = ba.MaterialBySLD("Layer", 0.0001, 1e-08, B)
    material_substrate = ba.MaterialBySLD("Substrate", 7e-05, 2e-06)

    # Layers
    layer_1 = ba.Layer(vacuum)
    layer_2 = ba.Layer(material_layer, 10*nm)
    layer_3 = ba.Layer(material_substrate)

    # Sample
    sample = ba.MultiLayer()
    sample.addLayer(layer_1)
    sample.addLayer(layer_2)
    sample.addLayer(layer_3)

    return sample


def simulate(sample, polarizer_vec, analyzer_vec, title):
    n = 50
    scan = ba.AlphaScan(n, 5*deg/n, 5*deg)
    scan.setWavelength(1.54*angstrom)
    scan.setPolarization(polarizer_vec)
    scan.setAnalyzer(analyzer_vec)

    simulation = ba.SpecularSimulation(scan, sample)

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


if __name__ == '__main__':
    sample = get_sample()
    results = [
        simulate(sample, R3(0, +1, 0), R3(0, +1, 0), "$++$"),
        simulate(sample, R3(0, +1, 0), R3(0, -1, 0), "$+-$"),
        simulate(sample, R3(0, -1, 0), R3(0, -1, 0), "$--$"),
    ]

    bp.plot_multicurve(results)
    bp.show_or_export()
