#!/usr/bin/env python3
"""
In this example we demonstrate how to plot simulation results with
axes in different units (nbins, mm, degs and QyQz).
"""
import bornagain as ba
from bornagain import angstrom, ba_plot as bp, deg, nm
from matplotlib import rcParams


def get_sample():
    # 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, 2*deg)
    interference = ba.Interference2DLattice(lattice)
    interference_pdf = ba.Profile2DCauchy(50*nm, 50*nm, 0)
    interference.setDecayFunction(interference_pdf)

    # 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.Sample()
    sample.addLayer(l_air)
    sample.addLayer(l_substrate)
    return sample


def transformed_plot(i, result, title):
    """
    Plots simulation results for different detectors.
    """

    bp.plt.subplot(2, 2, i)
    bp.plot_simres(result,
                   # xlabel=r'$x \;({\rm mm})$',
                   # ylabel=r'$y \;({\rm mm})$',
                   zlabel=None,
                   with_cb=False)
    bp.plt.title(title, loc='left')


if __name__ == '__main__':
    beam = ba.Beam(1e9, 1*angstrom, 0.5*deg)
    n = 500
    width = 170  # nm
    detector = ba.FlatDetector(n, n, width, width, beam, ba.FlatDetector.R, 2000.)
    simulation = ba.ScatteringSimulation(beam, get_sample(), detector)
    result = simulation.simulate()

    # setup plot
    rcParams['image.aspect'] = 'auto'
    bp.plt.figure(figsize=(10, 10))

    transformed_plot(1, result, "Real-space detector coordinates")
    transformed_plot(2, detector.field2bins(result), "Bin indices")
    transformed_plot(3, detector.field2angles(result), "Deflection angles")
    transformed_plot(4, detector.field2q(result), "Q space")

    # finish plot
    bp.plt.subplots_adjust(
        left=0.07,
        right=0.97,
        top=0.9,
        bottom=0.1,
        hspace=0.35,
        wspace=0.,
    )
    bp.plt.show()
