#  **************************************************************************  #
"""
#   BornAgain: simulate and fit reflection and scattering
#
#   @file      Wrap/Python/ba_plot.py
#   @brief     Python extensions of the SWIG-generated Python module bornagain.
#
#   @homepage  http://apps.jcns.fz-juelich.de/BornAgain
#   @license   GNU General Public License v3 or higher (see COPYING)
#   @copyright Forschungszentrum Juelich GmbH 2016
#   @authors   Scientific Computing Group at MLZ (see CITATION, AUTHORS)
"""
#  **************************************************************************  #

try:
    import os, sys
    import bornagain as ba
    import numpy as np
    import matplotlib as mpl
    from matplotlib import rc, pyplot as plt
except Exception as e:
    print(f"Import failure in ba_plot.py: {e}")

plotargs_default = {
    'label_fontsize': 16,
    'legendloc': "upper right"
}
plotargs_default['cmap'] = os.environ.setdefault('CMAP', "inferno")

# Arguments from plotargs_default that are not used by all plot functions.
# Plot functions should ignore these (not reject them as unknown).
_passthrough_args = {'figfile', 'label_fontsize', 'legendloc', 'cmap'}

#  **************************************************************************  #
#  environment variables, plot parameters, command-line arguments
#  **************************************************************************  #

def _env_to_bool(varname):
    if varname not in os.environ:
        return False
    value = os.environ[varname].lower()
    if value in ('false', 'off', 'n', 'no', '0'):
        return False
    if value in ('true', 'on', 'y', 'yes', '1'):
        return True
    raise Exception(
        f"Environment variable {varname} has ambiguous value {value}.")

if _env_to_bool('USETEX'):
    rc('text', usetex=True)

# Set consistent font for all text elements (axis labels, tick labels, title).
# Use a font that supports Greek and mathematical symbols and is widely available.
plt.rcParams['font.family'] = "DejaVu Sans"
plt.rcParams['mathtext.fontset'] = 'dejavusans'

def _check_unknown_args(kwargs):
    """Raise if kwargs contains keys not in _passthrough_args."""
    unknown = set(kwargs.keys()) - _passthrough_args
    if unknown:
        raise Exception(f'Unknown argument(s): {", ".join(sorted(unknown))}')

def parse_commandline():
    """ parse the arguments given on the command-line """

    plotargs = dict(plotargs_default)

    try:
        get_ipython()  # raises NameError outside IPython/Jupyter
        return plotargs
    except NameError:
        pass

    for arg in sys.argv[1:]:
        s = arg.split('=')
        if len(s) != 2:
            raise Exception(
                f"command-line argument '{arg}' does not have form key=value"
            )
        try:
            plotargs[s[0]] = int(s[1])
        except:
            plotargs[s[0]] = s[1]

    return plotargs

#  **************************************************************************  #
#  internal functions
#  **************************************************************************  #

def _setup_colorbar(fig, im, cax, label, fontsize):
    """
    Configure colorbar with consistent styling.
    """
    cb = fig.colorbar(im, cax=cax)
    if label:
        cb.set_label(label, size=fontsize)
    cb.ax.tick_params(labelsize=fontsize)
    cb.ax.yaxis.set_minor_locator(mpl.ticker.LogLocator(subs='all', numticks=20))
    cb.ax.yaxis.set_minor_formatter(mpl.ticker.NullFormatter())
    return cb

def _get_axes_limits(result):
    """
    Returns axes range as expected by pyplot.imshow.
    :param result: Datafield object from a Simulation
    :return: axes ranges as a flat list
    """
    limits = []
    for i in range(result.rank()):
        ax = result.axis(i)
        if ax.size() == 1:
            raise Exception(f"Axis {i} '{ax.axisLabel()}' has size 1:" +
                            " rather plot <datafield>.flat()")
        ami = ax.min()
        ama = ax.max()
        assert ami < ama, f"Datafield has invalid axis {i}, extending from {ami} to {ama}"
        limits.append(ami)
        limits.append(ama)

    return limits

def _translate_axis_label(label):
    """
    Formats an axis label into a LaTeX representation
    :param label: text representation of the axis label
    :return: LaTeX representation
    """
    label_dict = {
        'X (nbins)': r'$X \; $(bins)',
        'X (mm)': r'$X \; $(mm)',
        'Y (nbins)': r'$Y \; $(bins)',
        'Y (mm)': r'$Y \; $(mm)',
        'phi_f (rad)': r'$\varphi_f \; $(rad)',
        'phi_f (deg)': r'$\varphi_f \;(^{\circ})$',
        'alpha_i (rad)': r'$\alpha_{\rm i} \; $(rad)',
        'alpha_i (deg)': r'$\alpha_{\rm i} \;(^{\circ})$',
        'alpha_f (rad)': r'$\alpha_{\rm f} \; $(rad)',
        'alpha_f (deg)': r'$\alpha_{\rm f} \; (^{\circ})$',
        'qx (1/nm)': r'$q_x \; $(nm$^{-1}$)',
        'qy (1/nm)': r'$q_y \; $(nm$^{-1}$)',
        'qz (1/nm)': r'$q_z \; $(nm$^{-1}$)',
        'q (1/nm)': r'$q \; $(nm$^{-1}$)',
        'lambda (nm)': r'$\lambda \; $(nm)',
        'Position (nm)': "Position (nm)"
    }
    if label in label_dict.keys():
        return label_dict[label]
    return label

def _get_axes_labels(result):
    """
    Returns axes range as expected by pyplot.imshow.
    :param result: Datafield object from a Simulation
    :return: axes ranges as a flat list
    Used internally and in Examples/fit/specular/RealLifeReflectometryFitting.py.
    """
    labels = []
    for i in range(result.rank()):
        labels.append(_translate_axis_label(result.axis(i).axisLabel()))

    return labels

def _plot_curve(xarray, yarray, **kwargs):
    """
    Used internally.
    """
    title = kwargs.pop('title', None)
    xlabel = kwargs.pop('xlabel', None)
    ylabel = kwargs.pop('ylabel', None)
    fontsize = kwargs.pop('label_fontsize',
                          plotargs_default['label_fontsize'])

    if xlabel:
        plt.xlabel(xlabel, fontsize=fontsize)
    if ylabel:
        plt.ylabel(ylabel, fontsize=fontsize)
    if title:
        plt.title(title, fontsize=fontsize)

    inside_ticks()

    plt.plot(xarray, yarray)

def _plot_specular_curve(result, **plotargs):
    """
    Plots intensity data for specular simulation result
    :param result: Datafield from SpecularSimulation
    Used internally.
    """
    pfield = result.plottableField()
    intensity = pfield.intensities()
    x_axis = pfield.xCenters()

    xlabel = plotargs.pop('xlabel', _get_axes_labels(pfield)[0])
    ylabel = plotargs.pop('ylabel', "Intensity")
    title = plotargs.pop('title', None)
    label_fontsize = plotargs.pop('label_fontsize', None)
    ymax = plotargs.pop('intensity_max', np.amax(np.amax(intensity)*2))
    ymin = plotargs.pop('intensity_min',
                        max(np.amin(intensity)*0.5, 1e-18*ymax))
    _check_unknown_args(plotargs)

    plt.yscale('log')
    plt.ylim([ymin, ymax])

    kwargs = {'xlabel': xlabel, 'ylabel': ylabel}
    if title is not None:
        kwargs['title'] = title
    if label_fontsize is not None:
        kwargs['label_fontsize'] = label_fontsize
    _plot_curve(x_axis, intensity, **kwargs)

#  **************************************************************************  #
#  multiple frames in one plot
#  **************************************************************************  #

class _MultiPlot:
    """
    Used internally.
    """

    def __init__(self, n, ncol, zlabel, fontsize, spacing, frame_height=1.0):
        self.n = n
        self.ncol = ncol
        self.nrow = 1 + (self.n - 1) // self.ncol
        self.zlabel = zlabel

        # Parameters as fraction of subfig size.
        xskip = spacing[0]
        yskip = spacing[1]
        bottomskip = yskip
        topskip = yskip/2
        leftskip = xskip
        rightskip = 0.28 + ncol*0.03
        xtot = self.ncol*1.0 + (self.ncol - 1)*xskip + leftskip + rightskip
        ytot = self.nrow*frame_height + (self.nrow - 1)*yskip + bottomskip + topskip

        # Spacing for subplots_adjust (fraction of subplot size).
        self.xskip = xskip
        self.yskip = yskip
        # Margins as fraction of total fig size.
        self.leftskip = leftskip/xtot
        self.rightskip = rightskip/xtot
        self.bottomskip = bottomskip/ytot
        self.topskip = topskip/ytot
        # For colorbar positioning (fraction of figure size).
        self.xskip_fig = xskip/xtot

        # Set total figure dimensions.
        ftot = 5
        if fontsize:
            self.fontsize = fontsize
        else:
            self.fontsize = plotargs_default['label_fontsize']
        # Create the figure 'fig' and its subplots axes ('tmp'->'axes').
        self.fig, tmp = plt.subplots(self.nrow,
                                     self.ncol,
                                     figsize=(ftot*xtot, ftot*ytot))
        if n > 1:
            self.axes = tmp.flat
        else:
            self.axes = [tmp]

        # Adjust whitespace around and between subfigures.
        plt.subplots_adjust(wspace=self.xskip,
                            hspace=self.yskip,
                            left=self.leftskip,
                            right=1 - self.rightskip,
                            bottom=self.bottomskip,
                            top=1 - self.topskip)

    def plot_colorlegend(self, im):
        # Force layout calculation to get actual subplot positions.
        self.fig.canvas.draw()
        # Get the bounding box of the rendered images (not the subplot allocation)
        # so the colorbar height matches the actual data area, even when imshow
        # uses a non-unit aspect ratio that shrinks the axes within the allocation.
        renderer = self.fig.canvas.get_renderer()
        fig_h_px = self.fig.get_size_inches()[1] * self.fig.dpi
        images = [ax.get_images()[0] for ax in self.axes[:self.n]
                  if ax.get_images()]
        if images:
            y0 = min(img.get_window_extent(renderer).y0 for img in images) / fig_h_px
            y1 = max(img.get_window_extent(renderer).y1 for img in images) / fig_h_px
        else:
            y0 = min(ax.get_position().y0 for ax in self.axes[:self.n])
            y1 = max(ax.get_position().y1 for ax in self.axes[:self.n])
        # Plot the color legend.
        cbar_ax = self.fig.add_axes([
            1 - self.rightskip + 0.4*self.xskip_fig, y0,
            0.25*self.xskip_fig, y1 - y0
        ])
        _setup_colorbar(self.fig, im, cbar_ax, self.zlabel, self.fontsize)

#  **************************************************************************  #
#  versatile plot calls
#  **************************************************************************  #

def inside_ticks():
    """
    Ticks settings for xy plots: on all axes and pointing inside.
    Used internally and in a few examples.
    """
    ax = plt.gca()
    ax.yaxis.set_ticks_position('both')
    ax.xaxis.set_ticks_position('both')
    ax.tick_params(which='both', direction='in',
                   labelsize=plotargs_default['label_fontsize'])

def _compute_aspect(data_shape, axes_limits, frame_aspect=None, unit_aspect=None):
    """
    Compute the matplotlib aspect parameter for imshow.

    :param data_shape: tuple (nrows, ncols) of the data array
    :param axes_limits: [xmin, xmax, ymin, ymax]
    :param frame_aspect: if given, desired width:height ratio of plot frame
    :param unit_aspect: if given, aspect based on axis units (1 = equal scaling)
    :return: aspect parameter for imshow
    """
    if frame_aspect is not None and unit_aspect is not None:
        raise Exception("Cannot specify both frame_aspect and unit_aspect")

    extent_width = axes_limits[1] - axes_limits[0]
    extent_height = axes_limits[3] - axes_limits[2]

    if frame_aspect is not None:
        if frame_aspect <= 0:
            raise Exception(f'frame_aspect must be positive, got {frame_aspect}')
        # frame_aspect = width:height of the plot frame
        return extent_width / (extent_height * frame_aspect)

    if unit_aspect is not None:
        if unit_aspect <= 0:
            raise Exception(f'unit_aspect must be positive, got {unit_aspect}')
        # unit_aspect=1 means equal scaling: 1 data unit in x = 1 data unit in y
        return 1 / unit_aspect

    # Default: square pixels
    return (extent_width / extent_height) * (data_shape[0] / data_shape[1])

def plot_heatmap(result, **kwargs):
    """
    Plots intensity data as heat map
    :param result: Datafield from GISAS/OffspecSimulation
    Used internally and in a few examples.
    """

    pfield = result.plottableField()
    axes_limits = _get_axes_limits(pfield)
    axes_labels = _get_axes_labels(pfield)

    if 'xlabel' not in kwargs:
        kwargs['xlabel'] = axes_labels[0]
    if 'ylabel' not in kwargs:
        kwargs['ylabel'] = axes_labels[1]

    array = result.intensities()
    assert len(array.shape) == 2
    assert array.shape[0] > 0
    assert array.shape[1] > 0
    zmax = kwargs.pop('intensity_max', np.amax(array))
    zmin = kwargs.pop('intensity_min', 1e-6*zmax)

    if zmin == zmax == 0.0:
        norm = mpl.colors.Normalize(0, 1)
    else:
        norm = mpl.colors.LogNorm(zmin, zmax)

    xlabel = kwargs.pop('xlabel', None)
    ylabel = kwargs.pop('ylabel', None)
    zlabel = kwargs.pop('zlabel', "Intensity")
    title = kwargs.pop('title', None)
    cmap = kwargs.pop('cmap', plotargs_default['cmap'])
    withCBar = kwargs.pop('with_cb', True)
    frame_aspect = kwargs.pop('frame_aspect', None)
    unit_aspect = kwargs.pop('unit_aspect', None)
    _check_unknown_args(kwargs)

    aspect = _compute_aspect(array.shape, axes_limits, frame_aspect, unit_aspect)

    fig = plt.gcf()
    ax = plt.gca()
    im = ax.imshow(array,
                   origin='lower',
                   cmap=cmap,
                   norm=norm,
                   aspect=aspect,
                   extent=axes_limits)

    fontsize = plotargs_default['label_fontsize']
    if xlabel:
        plt.xlabel(xlabel, fontsize=fontsize)
    if ylabel:
        plt.ylabel(ylabel, fontsize=fontsize)
    if title:
        plt.title(title, fontsize=fontsize)
    ax.tick_params(axis='both', which='major', labelsize=fontsize)

    if withCBar:
        # Force layout calculation to get actual image position
        fig.canvas.draw()
        # Get image position in figure coordinates
        bbox = ax.get_position()
        cbar_width = 0.03
        cbar_pad = 0.02
        cax = fig.add_axes([bbox.x1 + cbar_pad, bbox.y0, cbar_width, bbox.height])
        _setup_colorbar(fig, im, cax, zlabel, fontsize)

    return im

#  **************************************************************************  #
#  standard user calls
#  **************************************************************************  #

def export(**plotargs):
    _figfile = plotargs.pop('figfile', None)
    if _figfile:
        plt.savefig(_figfile, bbox_inches='tight')

def plot_datafield(result, **plotargs):
    """
    Draws simulation result and (optionally) shows the plot.
    """
    data = result.intensities()
    if len(data.shape) == 1:
        # 1D data => assume specular simulation
        _plot_specular_curve(result, **plotargs)
        return
    if len(data.shape) != 2:
        raise Exception("Unexpected dimension of data array")
    if data.shape[0] == 1 or data.shape[1] == 1:
        _plot_specular_curve(result.flat(), **plotargs)
        return
    plot_heatmap(result, **plotargs)

def plot_ff_to_row(results, **plotargs):
    plotargs.setdefault('zlabel', r'$\left|F(q)\right|^2/V^{\,2}$')
    plotargs.setdefault('log_range', 8)
    plotargs.setdefault('normalized_cb', True)
    plot2d_to_row(results, **plotargs)

def plot2d_to_row(results, **plotargs):
    plot2d_to_grid(results, len(results), **plotargs)

def plot2d_to_grid(results, ncol, **plotargs):
    """
    Make a plot consisting of one detector image for each Result in results,
    plus one common color legend.

    :param results: List of simulation results
    :param ncol:    Maximum number of plot frames per row
    """
    pfields = [result.plottableField() for result in results]
    zlabel = plotargs.pop('zlabel', "Intensity")
    fontsize = plotargs.pop('fontsize', None)
    spacing = plotargs.pop('spacing', (0.18, 0.2))
    cmap = plotargs.pop('cmap', plotargs_default['cmap'])

    # Colorbar range: max = 1
    log_range = plotargs.pop('log_range', 8)
    lr_factor = 10**(-log_range)
    norm = mpl.colors.LogNorm(lr_factor, 1)

    # Colorbar range: max = peak intensity value
    if not plotargs.pop('normalized_cb', False):
        global_max = np.amax([result.maxVal() for result in results])
        zmax = plotargs.pop('intensity_max', global_max)
        zmin = plotargs.pop('intensity_min', lr_factor*global_max)
        norm = mpl.colors.LogNorm(zmin, zmax)

    frame_aspect = plotargs.pop('frame_aspect', None)
    unit_aspect = plotargs.pop('unit_aspect', None)
    _check_unknown_args(plotargs)

    # Compute normalized frame height for figure sizing.
    # When a non-default aspect is requested, the frames are not unit-square;
    # scale the figure height so that each frame fills its subplot allocation.
    frame_height = 1.0
    if pfields and (unit_aspect is not None or frame_aspect is not None):
        lim = _get_axes_limits(pfields[0])
        ew = lim[1] - lim[0]
        eh = lim[3] - lim[2]
        if unit_aspect is not None and ew > 0:
            frame_height = max(eh / (ew * unit_aspect), 0.1)
        elif frame_aspect is not None and frame_aspect > 0:
            frame_height = 1.0 / frame_aspect
    mp = _MultiPlot(len(pfields), ncol, zlabel, fontsize, spacing,
                    frame_height=frame_height)

    # Plot the subfigures.
    for i, pfield in enumerate(pfields):
        ax = mp.axes[i]
        axes_limits = _get_axes_limits(pfield)
        axes_labels = _get_axes_labels(pfield)
        data = pfield.intensities()
        if len(data.shape) != 2 or data.shape[0] <= 1 or data.shape[1] <= 1:
            raise Exception(f"plot2d_to_grid requires 2D data, got shape {data.shape}")

        aspect = _compute_aspect(data.shape, axes_limits, frame_aspect, unit_aspect)

        im = ax.imshow(data,
                       origin='lower',
                       cmap=cmap,
                       norm=norm,
                       extent=axes_limits,
                       aspect=aspect)

        ax.set_xlabel(axes_labels[0], fontsize=mp.fontsize)
        if i % ncol == 0:
            ax.set_ylabel(axes_labels[1], fontsize=mp.fontsize)
        if pfield.title() != "":
            ax.set_title(pfield.title(), fontsize=mp.fontsize)
        ax.tick_params(axis='both',
                       which='major',
                       labelsize=mp.fontsize)

    mp.plot_colorlegend(im)

def plot_multicurve(results, **plotargs):
    legendloc = plotargs.pop('legendloc', "upper right")
    _check_unknown_args(plotargs)

    plt.yscale('log')
    pfields = [result.plottableField() for result in results]

    legend = []
    for pfield in pfields:
        x = pfield.xCenters()
        y = pfield.intensities()
        legend.append(pfield.title())
        plt.plot(x, y)

    inside_ticks()

    fontsize = plotargs_default['label_fontsize']
    plt.xlabel(_get_axes_labels(pfields[0])[0], fontsize=fontsize)
    plt.ylabel("Intensity", fontsize=fontsize)

    plt.legend(legend, loc=legendloc, fontsize=fontsize)
