#!/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 pyplot as plt
from matplotlib import rcParams


def get_sample():
    from bornagain import std_samples
    return std_samples.cylinders()


def get_simulation(sample):
    beam = ba.Beam(1e9, 1*angstrom, 0.2*deg)
    n = <%= sm ? 11 : 500 %>
    detector_distance = 2000.0  # in mm
    width = 170  # nm
    height = width
    detector = ba.RectangularDetector(n, width, n, height)
    detector.setPerpendicularToSampleX(detector_distance, width/2., 0)

    simulation = ba.ScatteringSimulation(beam, sample, detector)
    return simulation


def plot(result):
    """
    Plots simulation results for different detectors.
    """
    # set plotting parameters
    rcParams['image.aspect'] = 'auto'

    plt.figure(figsize=(10, 10))

    plt.subplot(2, 2, 1)
    # default units for rectangular detector are millimeters

    bp.plot_simres(result,
                   xlabel=r'$x \;({\rm mm})$',
                   ylabel=r'$y \;({\rm mm})$',
                   zlabel=None,
                   with_cb=False)
    plt.title("Default: real-space detector coordinates", loc='left')

    plt.subplot(2, 2, 2)
    bp.plot_simres(result,
                   units=ba.Coords_NBINS,
                   xlabel=r'$X_{\rm bin}$',
                   ylabel=r'$Y_{\rm bin}$',
                   zlabel=None,
                   with_cb=False)
    plt.title("Bin indices", loc='left')

    plt.subplot(2, 2, 3)
    bp.plot_simres(result,
                   units=ba.Coords_DEGREES,
                   xlabel=r'$\varphi_{\rm f} \;(^{\circ})$',
                   ylabel=r'$\alpha_{\rm f} \;(^{\circ})$',
                   zlabel=None,
                   with_cb=False)
    plt.title("Deflection angles", loc='left')

    plt.subplot(2, 2, 4)
    bp.plot_simres(result,
                   units=ba.Coords_QSPACE,
                   xlabel=r'$Q_{y} \;(1/{\rm nm})$',
                   ylabel=r'$Q_{z} \;(1/{\rm nm})$',
                   zlabel=None,
                   with_cb=False)
    plt.title("Q space", loc='left')

    plt.subplots_adjust(
        left=0.07,
        right=0.97,
        top=0.9,
        bottom=0.1,
        hspace=0.35,
        wspace=0.,
    )


if __name__ == '__main__':
    simulation = get_simulation(get_sample())
    result = simulation.simulate()

    plot(result)
    bp.show_or_export()
