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


def get_sample(hasVariance, xi):
    # Define materials
    m_air = ba.HomogeneousMaterial("Air", 0, 0)
    m_particle = ba.HomogeneousMaterial("Particle", 0.0006, 2e-08)
    m_substrate = ba.HomogeneousMaterial("Substrate", 6e-06, 2e-08)

    # Define particles
    R = 2.5*nm
    ff = ba.FormFactorFullSpheroid(R, R)
    particle = ba.Particle(m_particle, ff)

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

    # Define particle layout
    layout = ba.ParticleLayout()
    layout.addParticle(particle)
    layout.setInterferenceFunction(interference)

    # Define layers
    l_air = ba.Layer(m_air)
    l_air.addLayout(layout)
    l_substrate = ba.Layer(m_substrate)

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


def get_simulation(sample):
    beam = ba.Beam(1e8, 0.1*nm, ba.Direction(0.2*deg, 0))

    det = ba.SphericalDetector(400, -2*deg, 2*deg, 400, 0, 3*deg)

    simulation = ba.GISASSimulation(beam, sample, det)
    return simulation


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

    plt.subplot(3, 2, nPlot)

    return ba_plot.plot_colormap(result,
                          title=title,
                          intensity_max=3e7,
                          intensity_min=3e0,
                          zlabel=None,
                          aspect='equal',
                          with_cb=False)


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

    xi1 =  5*deg
    xi2 = 15*deg
    im = run_one(False, 0*deg, 1, "xi=0 deg, fixed positions")
    run_one(True, 0*deg, 2, "position variance 0.3 nm")
    run_one(False, xi1, 3, "xi=5 deg, fixed positions")
    run_one(True, xi1, 4, "position variance 0.3 nm")
    run_one(False, xi2, 5, "xi=15 deg, fixed positions")
    run_one(True, xi2, 6, "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]))
    plt.show()
