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


def get_sample(hasVariance, xi):
    # Materials
    air_color = (0.90, 0.93, 0.97)
    air_mat = ba.RefractiveMaterial("Air", air_color, 0, 0)
    particle_color = (0.86, 0.24, 0.18)
    particle_mat = ba.RefractiveMaterial("Particle", particle_color, 0.0006, 2e-08)
    substrate_color = (0.28, 0.57, 0.82)
    substrate_mat = ba.RefractiveMaterial("Substrate", substrate_color, 6e-06, 2e-08)

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

    # Interference function
    lattice = ba.SquareLattice2D(10*nm, xi)
    layout = ba.Crystal2D(particle, lattice)
    profile = ba.Profile2DCauchy(500*nm, 500*nm, 0)
    layout.setDecayFunction(profile)
    if hasVariance:
        layout.setLateralPositionVariance(0.3*nm)

    # Layers
    l_air = ba.Layer(air_mat)
    l_air.deposit2D(layout)
    l_substrate = ba.Layer(substrate_mat)

    # Sample
    sample = ba.Sample()
    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.SphericalDetector(n, -2*deg, 2*deg, n, 0, 3*deg)
    return ba.ScatteringSimulation(beam, sample, det)


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

    bp.plt.sca(ax)
    return bp.plot_heatmap(result,
                           title=title,
                           intensity_max=3e7,
                           intensity_min=3e0,
                           with_cb=False,
                           unit_aspect=1)


if __name__ == '__main__':
    fig, axs = bp.plt.subplots(3, 2, figsize=(10, 11.5), constrained_layout=False)

    # Margins (left, right, bottom, top) are relative to the figure;
    # spacing (wspace, hspace) is relative to allocated average subplot size:
    fig.subplots_adjust(left=0.08, right=0.82, bottom=0.04, top=0.97,
                        wspace=0.22, hspace=0.2)

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

    # Position colorbar to match middle row data area
    import matplotlib.ticker as ticker
    bbox = axs[1, 1].get_position()
    cax = fig.add_axes([0.88, bbox.y0, 0.025, bbox.height])
    fontsize = bp.plotargs_default['label_fontsize']
    cb = fig.colorbar(im, cax=cax)
    cb.set_label("Intensity", fontsize=fontsize)
    cb.ax.tick_params(labelsize=fontsize)
    cb.ax.yaxis.set_minor_locator(ticker.LogLocator(subs='all', numticks=20))
    cb.ax.yaxis.set_minor_formatter(ticker.NullFormatter())

