#!/usr/bin/env python3
"""
Reflectivity with and without polarization analyzer.
"""
import unittest
import bornagain as ba
from bornagain import deg, nm, R3
from bornagain.numpyutil import Arrayf64Converter as dac

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.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.axis(0).binCenters(), dac.npArray(result.dataArray())


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()
