#!/usr/bin/env python3
"""
Reflectivity with and without polarization analyzer.
"""
import bornagain as ba
from bornagain import deg, nm, R3
from math import sin, cos


def get_sample():
    # Materials
    Bmag = 1e8
    Bangle = 60*deg
    B = R3(Bmag*cos(Bangle), Bmag*sin(Bangle), 0)
    material_layer = ba.MaterialBySLD("Layer", 0.0001, 1e-08, B)
    material_substrate = ba.MaterialBySLD("Substrate", 7e-05, 2e-06)
    vacuum = ba.MaterialBySLD("Vacuum", 0, 0)

    # 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 run_simulation(polarizer_vec, analyzer_vec=None):
    sample = get_sample()

    n = 50
    scan = ba.AlphaScan(n, 5*deg/n, 5*deg)
    scan.setWavelength(0.15*nm)

    # adding polarizer and analyzer operator
    scan.setPolarization(polarizer_vec)
    if analyzer_vec:
        scan.setAnalyzer(analyzer_vec)

    simulation = ba.SpecularSimulation(scan, sample)

    result = simulation.simulate()

    return result.axis(0).binCenters(), result.npArray()


if __name__ == '__main__':
    q, results_pp = run_simulation(R3(0, +1, 0), R3(0, +1, 0))
    q, results_mm = run_simulation(R3(0, -1, 0), R3(0, -1, 0))

    q, results_pm = run_simulation(R3(0, +1, 0), R3(0, -1, 0))
    q, results_mp = run_simulation(R3(0, -1, 0), R3(0, +1, 0))

    r_plus = results_pp + results_pm
    r_minus = results_mm + results_mp

    # same result, but need half the computational time
    q, results_p = run_simulation(R3(0, +1, 0))
    q, results_m = run_simulation(R3(0, -1, 0))

    if not ba.checkRelativeDifference(results_pp + results_pm, results_p, 1e-14):
        exit(1)
    if not ba.checkRelativeDifference(results_mp + results_mm, results_m, 1e-14):
        exit(1)
    exit(0)
