#!/usr/bin/env python3
"""
"""
import bornagain as ba
from bornagain import ba_plot as bp, deg, nm
from matplotlib import pyplot as plt


def get_sample(hasVariance, xi):
    # Materials
    material_air = ba.RefractiveMaterial("Air", 0, 0)
    material_particle = ba.RefractiveMaterial("Particle", 0.0006, 2e-08)
    material_substrate = ba.RefractiveMaterial("Substrate", 6e-06, 2e-08)

    # Particles
    R = 2.5*nm
    ff = ba.Spheroid(R, R)
    particle = ba.Particle(material_particle, ff)

    # Interference function
    lattice = ba.SquareLattice2D(10*nm, xi)
    interference = ba.Interference2DLattice(lattice)
    interference_pdf = ba.Profile2DCauchy(500*nm, 500*nm, 0)
    interference.setDecayFunction(interference_pdf)
    if hasVariance:
        interference.setPositionVariance(0.3*nm)

    # Particle layout
    layout = ba.ParticleLayout()
    layout.addParticle(particle)
    layout.setInterference(interference)

    # Layers
    l_air = ba.Layer(material_air)
    l_air.addLayout(layout)
    l_substrate = ba.Layer(material_substrate)

    # Sample
    sample = ba.MultiLayer()
    sample.addLayer(l_air)
    sample.addLayer(l_substrate)
    return sample


def get_simulation(sample):
    n = 11
    beam = ba.Beam(1e8, 0.1*nm, 0.2*deg)
    det = ba.Detector2D(4*deg, 3*deg, n, n, 0, 1.5*deg)
    return ba.ScatteringSimulation(beam, sample, det)


def run_one(hasVariance, xi, nPlot, title):
    sample = get_sample(hasVariance, xi)
    simulation = get_simulation(sample)
    result = simulation.simulate()

    plt.subplot(3, 2, nPlot)

    return bp.plot_simres(result,
                          title=title,
                          intensity_max=3e7,
                          intensity_min=3e0,
                          aspect='equal',
                          with_cb=False)


if __name__ == '__main__':
    fig, axs = plt.subplots(3, 2, figsize=(10, 13))

    xi1 = 5*deg
    xi2 = 15*deg
    im = run_one(False, 0*deg, 1, r"$\xi=0^\circ$, fixed positions")
    run_one(True, 0*deg, 2, r"position variance 0.3 nm")
    run_one(False, xi1, 3, r"$\xi=5^\circ$, fixed positions")
    run_one(True, xi1, 4, r"position variance 0.3 nm")
    run_one(False, xi2, 5, r"$\xi=15^\circ$, fixed positions")
    run_one(True, xi2, 6, r"position variance 0.3 nm")

    plt.subplots_adjust(bottom=0.05, left=0.1)
    plt.colorbar(im, cax=plt.axes([0.93, 0.36, 0.03, 0.21]))

    bp.show_or_export()
