#!/usr/bin/env python3
"""
Reflectivity with and without polarization analyzer.
"""
import unittest
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)
    layer_color = (0.05, 0.62, 0.55)
    layer_mat = ba.SLDMaterial("Layer", layer_color, 0.0001, 1e-08, B)
    substrate_color = (0.55, 0.55, 0.55)
    substrate_mat = ba.SLDMaterial("Substrate", substrate_color, 7e-05, 2e-06)
    vacuum = ba.Vacuum()

    # 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 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.xCenters(), result.intensities()


class PolarizedNoAnalyzerTest(unittest.TestCase):

    def testPolarizedNoAnalyzer(self):
        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))

        self.assertTrue(ba.checkRelVecDifference(results_pp + results_pm, results_p, 1e-14))
        self.assertTrue(ba.checkRelVecDifference(results_mp + results_mm, results_m, 1e-14))


if __name__ == '__main__':
    unittest.main()
