#!/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.Vacuum()
    layer_color = (0.05, 0.62, 0.55)
    layer_mat = ba.SLDMaterial("Layer", layer_color, 0.0001, 1e-08, B)
    substrate_color = (0.28, 0.57, 0.82)
    substrate_mat = ba.SLDMaterial("Substrate", substrate_color, 7e-05, 2e-06)

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

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

    return sample


def simulate(sample, polarizer_vec, analyzer_vec, title):
    n = <%= test_mode ? 50 : 500 %>
    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()
    <%- unless test_mode || figure_mode -%>
    ba.showSample3D(sample, sample_size=80*nm, seed=0)
    <%- end -%>
    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), "$--$"),
    ]
<%= include_partial("template/main/plot_multicurve.py") -%>
