# Mantid Repository : https://github.com/mantidproject/mantid
#
# Copyright &copy; 2018 ISIS Rutherford Appleton Laboratory UKRI,
#   NScD Oak Ridge National Laboratory, European Spallation Source,
#   Institut Laue - Langevin & CSNS, Institute of High Energy Physics, CAS
# SPDX - License - Identifier: GPL - 3.0 +
"""
Plotting functions for use in IPython notebooks that are generated by MantidPlot
"""

import matplotlib.pyplot as plt

# Import Mantid
from mantid.simpleapi import *
import mantid.api as mapi


def _plot_with_options(axes_option, workspace, options_list, plot_number):
    """
    Enable/disable legend, grid, limits according to
    options (ops) for the given axes (ax).
    Plot with or without errorbars.
    """
    ws_plot = ConvertToPointData(workspace)
    if options_list['errorbars']:
        axes_option.errorbar(ws_plot.readX(0), ws_plot.readY(0),
                             yerr=ws_plot.readE(0), label=workspace.name())
    else:
        axes_option.plot(ws_plot.readX(0),
                         ws_plot.readY(0),
                         label=workspace.name())

    axes_option.grid(options_list['grid'])
    axes_option.set_xscale(options_list['xScale'])
    axes_option.set_yscale(options_list['yScale'])
    if options_list['xLimits'] != 'auto':
        axes_option.set_xlim(options_list['xLimits'])
    if options_list['yLimits'] != 'auto':
        axes_option.set_ylim(options_list['yLimits'])

    # If a list of titles was given, use it to title each subplot
    if hasattr(options_list['title'], "__iter__"):
        axes_option.set_title(options_list['title'][plot_number])
    if options_list['legend'] and hasattr(options_list['legendLocation'], "__iter__"):
        axes_option.legend(loc=options_list['legendLocation'][plot_number])
    elif options_list['legend']:
        axes_option.legend(loc=options_list['legendLocation'])


def plots(list_of_workspaces, *args, **kwargs):
    """
    Create a figure with a subplot for each workspace given.
    Workspaces within a group workspace are plotted together in the same subplot.

    Examples:
    plots(rr)
    plots(rr, 'TheGraphTitle')
    plots(rr, 'TheGraphTitle', grid=True, legend=True,
          xScale='linear', yScale='log', xLimits=[0.008, 0.16])
    plots(rr, sharedAxes = False, xLimits = [0, 0.1], yLimits = [1e-5, 2],
          Title='ASF070_07 I=1A T=3K dq/q=2%',
          legend=True, legendLocation=3, errorbars=False)
    """

    if not hasattr(list_of_workspaces, "__iter__"):
        list_of_workspaces = [list_of_workspaces]

    ops = _process_arguments(args, kwargs)

    # Create subplots for workspaces in the list
    fig, axes_handle = plt.subplots(1,
                                    len(list_of_workspaces),
                                    sharey=ops['sharedAxes'],
                                    figsize=(6 * len(list_of_workspaces), 4))

    if not hasattr(axes_handle, "__iter__"):
        axes_handle = [axes_handle]

    for plot_number, workspace in enumerate(list_of_workspaces):
        if isinstance(workspace, mapi.WorkspaceGroup):
            # Plot grouped workspaces on the same axes
            for sub_ws in workspace:
                _plot_with_options(axes_handle[plot_number], sub_ws, ops, plot_number)
        else:
            _plot_with_options(axes_handle[plot_number], workspace, ops, plot_number)

    # If a single title was given, use it to title the whole figure
    if not hasattr(ops['title'], "__iter__"):
        fig.suptitle(ops['title'])
    plt.show()

    return plt.gcf()


def _process_arguments(input_args, input_kwargs):
    """
    Build a dictionary of plotting options
    """
    key_list = ['title', 'grid', 'legend', 'legendLocation',
                'xScale', 'yScale', 'xLimits', 'yLimits', 'sharedAxes', 'errorbars']
    default_values = ['', True, True, 1, 'log', 'log', 'auto', 'auto', True, 'True']

    # Fill ops with the default values
    for i in range(len(input_args)):  # copy in values provided in args
        default_values[i] = input_args[i]
    ops = dict(zip(key_list, default_values))
    for k in ops.keys():  # copy in any key word given arguments
        ops[k] = input_kwargs.get(k, ops[k])

    return ops
