"""

This module provides R style pairs plotting functionality.

"""
import matplotlib.cm as cm
import matplotlib.gridspec as gridspec
import matplotlib.pyplot as plt
import numpy as np

# from . import plotting_util
from .plotting_util import LegendEnum, get_color, prepare_pairs_data, make_legend
from ..util import get_module_logger

# .. codeauthor:: jhkwakkel <j.h.kwakkel (at) tudelft (dot) nl>

__all__ = ["pairs_scatter", "pairs_lines", "pairs_density"]
_logger = get_module_logger(__name__)


def pairs_lines(
    experiments,
    outcomes,
    outcomes_to_show=[],
    group_by=None,
    grouping_specifiers=None,
    ylabels={},
    legend=True,
    **kwargs,
):
    """

    Generate a `R style pairs <https://www.stat.psu.edu/~dhunter/R/html/graphics/html/pairs.html>`_
    lines multiplot. It shows the behavior of two outcomes over time against
    each other. The origin is denoted with a circle and the end is denoted
    with a '+'.

    Parameters
    ----------
    experiments : DataFrame
    outcomes : dict
    outcomes_to_show : list of str, optional
                       list of outcome of interest you want to plot.
    group_by : str, optional
               name of the column in the cases array to group results by.
               Alternatively, `index` can be used to use indexing arrays as the
               basis for grouping.
    grouping_specifiers : dict, optional
                          dict of categories to be used as a basis for grouping
                          by. Grouping_specifiers is only meaningful if
                          group_by is provided as well. In case of grouping by
                          index, the grouping  specifiers should be in a
                          dictionary where the key denotes the name of the
                          group.
    ylabels : dict, optional
              ylabels is a dictionary with the outcome names as keys, the
              specified values will be used as labels for the y axis.
    legend : bool, optional
             if true, and group_by is given, show a legend.
    point_in_time : float, optional
                    the point in time at which the scatter is to be made. If
                    None is provided (default), the end states are used.
                    point_in_time should be a valid value on time

    Returns
    -------
    fig
        the figure instance
    dict
        key is tuple of names of outcomes, value is associated axes
        instance

    """

    # unravel return from run_experiments
    _logger.debug("making a pairwise line plot")

    prepared_data = prepare_pairs_data(
        experiments, outcomes, outcomes_to_show, group_by, grouping_specifiers, None
    )
    experiments, outcomes, outcomes_to_show, grouping_labels = prepared_data

    grid = gridspec.GridSpec(len(outcomes_to_show), len(outcomes_to_show))
    grid.update(wspace=0.1, hspace=0.1)

    # the plotting
    figure = plt.figure()
    axes_dict = {}

    combis = [(field1, field2) for field1 in outcomes_to_show for field2 in outcomes_to_show]

    for field1, field2 in combis:
        i = list(outcomes_to_show).index(field1)
        j = list(outcomes_to_show).index(field2)
        ax = figure.add_subplot(grid[i, j])

        axes_dict[(field1, field2)] = ax

        if group_by:
            for x, entry in enumerate(grouping_labels):
                data1 = outcomes[entry][field1]
                data2 = outcomes[entry][field2]
                color = get_color(x)
                if i == j:
                    color = "white"
                simple_pairs_lines(ax, data1, data2, color)
        else:
            data1 = outcomes[field1]
            data2 = outcomes[field2]
            color = "b"
            if i == j:
                color = "white"
            simple_pairs_lines(ax, data1, data2, color)
        do_text_ticks_labels(ax, i, j, field1, field2, ylabels, outcomes_to_show)

    if group_by and legend:
        gs1 = grid[0, 0]

        for ax in figure.axes:
            gs2 = ax._subplotspec
            if all((gs1._gridspec == gs2._gridspec, gs1.num1 == gs2.num1, gs1.num2 == gs2.num2)):
                break

        make_legend(grouping_labels, ax, legend_type=LegendEnum.LINE)

    return figure, axes_dict


def simple_pairs_lines(ax, y_data, x_data, color):
    """

    Helper function for generating a simple pairs lines plot

    Parameters
    ----------
    ax : axes
    data1 : ndarray
    data2 : ndarray
    color : str

    """

    ax.plot(x_data.T, y_data.T, c=color)
    ax.scatter(x_data[:, 0], y_data[:, 0], edgecolor=color, facecolor=color, marker="o")
    ax.scatter(x_data[:, -1], y_data[:, -1], facecolor=color, marker="+")


def pairs_density(
    experiments,
    outcomes,
    outcomes_to_show=[],
    group_by=None,
    grouping_specifiers=None,
    ylabels={},
    point_in_time=-1,
    log=True,
    gridsize=50,
    colormap="coolwarm",
    filter_scalar=True,
):
    """

    Generate a `R style pairs <https://www.stat.psu.edu/~dhunter/R/html/graphics/html/pairs.html>`_
    hexbin density multiplot. In case of time-series data, the end
    states are used.

    hexbin makes hexagonal binning plot of x versus y, where x, y are 1-D
    sequences of the same length, N. If C is None (the default), this is a
    histogram of the number of occurrences of the observations at (x[i],y[i]).
    For further detail see `matplotlib on hexbin <https://matplotlib.sourceforge.net/api/pyplot_api.html#matplotlib.pyplot.hexbin>`_

    Parameters
    ----------
    experiments : DataFrame
    outcomes : dict
    outcomes_to_show : list of str, optional
                       list of outcome of interest you want to plot.
    group_by : str, optional
               name of the column in the cases array to group results by.
               Alternatively, `index` can be used to use indexing arrays as the
               basis for grouping.
    grouping_specifiers : dict, optional
                          dict of categories to be used as a basis for grouping
                          by. Grouping_specifiers is only meaningful if
                          group_by is provided as well. In case of grouping by
                          index, the grouping  specifiers should be in a
                          dictionary where the key denotes the name of the
                          group.
    ylabels : dict, optional
              ylabels is a dictionary with the outcome names as keys, the
              specified values will be used as labels for the y axis.
    point_in_time : float, optional
                    the point in time at which the scatter is to be made. If
                    None is provided (default), the end states are used.
                    point_in_time should be a valid value on time
    log: bool, optional
        indicating whether density should be log scaled. Defaults to True.
    gridsize : int, optional
               controls the gridsize for the hexagonal bining. (default = 50)
    cmap : str
           color map that is to be used in generating the hexbin. For details
           on the available maps, see `pylab <https://matplotlib.sourceforge.net/examples/pylab_examples/show_colormaps.html#pylab-examples-show-colormaps>`_.
           (Defaults = coolwarm)
    filter_scalar: bool, optional
                   remove the non-time-series outcomes. Defaults to True.

    Returns
    -------
    fig
        the figure instance
    dict
        key is tuple of names of outcomes, value is associated axes
        instance

    """
    _logger.debug("generating pairwise density plot")

    prepared_data = prepare_pairs_data(
        experiments,
        outcomes,
        outcomes_to_show,
        group_by,
        grouping_specifiers,
        point_in_time,
        filter_scalar,
    )
    experiments, outcomes, outcomes_to_show, grouping_specifiers = prepared_data

    if group_by:
        # figure out the extents for each combination
        extents = determine_extents(outcomes, outcomes_to_show)

        axes_dicts = {}
        figures = []
        for key, value in outcomes.items():
            figure, axes_dict = simple_pairs_density(
                value,
                outcomes_to_show,
                log,
                colormap,
                gridsize,
                ylabels,
                extents=extents,
                title=key,
            )
            axes_dicts[key] = axes_dict
            figures.append(figure)

        # harmonize the color scaling across figures
        combis = [(field1, field2) for field1 in outcomes_to_show for field2 in outcomes_to_show]
        for combi in combis:
            if combi[0] == combi[1]:
                continue

            vmax = -1
            vmin = 10000
            for entry in axes_dicts.values():
                vmax = max(entry[combi].collections[0].norm.vmax, vmax)
                vmin = min(entry[combi].collections[0].norm.vmin, vmin)
            for entry in axes_dicts.values():
                ax = entry[combi]
                ax.collections[0].set_clim(vmin=vmin, vmax=vmax)
            del vmin, vmax

        return figures, axes_dicts
    else:
        return simple_pairs_density(outcomes, outcomes_to_show, log, colormap, gridsize, ylabels)


def determine_extents(outcomes, outcomes_to_show):
    """
    Helper function used by pairs_density to make sure that multiple groups
    share the same axes extent.

    Parameters
    ----------
    outcomes : dict
    outcomes_to_show : list of str

    Returns
    -------
    dict
        tuple of str as key, and 4-tuple with extent

    """

    limits = {}
    for pol_out in outcomes.values():
        for entry in outcomes_to_show:
            out = pol_out[entry]
            minimum = np.amin(out)
            maximum = np.amax(out)
            try:
                cur = limits[entry]
                new = (min(cur[0], minimum), max(cur[1], maximum))
                limits[entry] = new
            except KeyError:
                limits[entry] = (minimum, maximum)
    extents = {}
    combis = [(field1, field2) for field1 in outcomes_to_show for field2 in outcomes_to_show]
    for field1, field2 in combis:
        limits_1 = limits[field1]
        limits_2 = limits[field2]
        extents[(field1, field2)] = (limits_1[0], limits_1[1], limits_2[0], limits_2[1])
    return extents


def simple_pairs_density(
    outcomes, outcomes_to_show, log, colormap, gridsize, ylabels, extents=None, title=None
):
    """

    Helper function for generating a simple pairs density plot

    Parameters
    ----------
    outcomes : dict
    outcomes_to_show : list of str
    log : bool
    colormap : str
    gridsize : int
    ylabels: dict
    extents : dict, optional
             used to control scaling of plots across figures, If provided, it
             should be a dict with a tuple of outcomes as key and the extend to
             be used as value.
    title : str, optional


    """

    grid = gridspec.GridSpec(len(outcomes_to_show), len(outcomes_to_show))
    grid.update(wspace=0.1, hspace=0.1)

    # the plotting
    figure = plt.figure()

    combis = [(field1, field2) for field1 in outcomes_to_show for field2 in outcomes_to_show]
    axes_dict = {}
    for field1, field2 in combis:
        i = list(outcomes_to_show).index(field1)
        j = list(outcomes_to_show).index(field2)

        ax = figure.add_subplot(grid[i, j])
        axes_dict[(field1, field2)] = ax

        y_data = outcomes[field1]
        x_data = outcomes[field2]

        bins = None
        if log:
            bins = "log"

        extent = None
        if extents:
            extent = extents[(field2, field1)]

        # text and labels
        if i == j:
            # only plot the name in the middle
            pass
            # ax.hexbin(x_data, y_data, bins=bins, gridsize=gridsize,
            #           cmap=cm.__dict__[colormap], alpha=0, edgecolor='white',
            #           linewidths=1, extent=extent)
        else:
            ax.hexbin(
                x_data,
                y_data,
                bins=bins,
                gridsize=gridsize,
                cmap=cm.__dict__[colormap],
                edgecolor="black",
                linewidths=1,
                extent=extent,
                mincnt=1,
            )
        do_text_ticks_labels(ax, i, j, field1, field2, ylabels, outcomes_to_show)

    return figure, axes_dict


def pairs_scatter(
    experiments,
    outcomes,
    outcomes_to_show=[],
    group_by=None,
    grouping_specifiers=None,
    ylabels={},
    legend=True,
    point_in_time=-1,
    filter_scalar=False,
    **kwargs,
):
    """

    Generate a `R style pairs <https://www.stat.psu.edu/~dhunter/R/html/graphics/html/pairs.html>`_
    scatter multiplot. In case of time-series data, the end states are used.

    Parameters
    ----------
    experiments : DataFrame
    outcomes : dict
    outcomes_to_show : list of str, optional
                       list of outcome of interest you want to plot.
    group_by : str, optional
               name of the column in the cases array to group results by.
               Alternatively, `index` can be used to use indexing arrays as the
               basis for grouping.
    grouping_specifiers : dict, optional
                          dict of categories to be used as a basis for grouping
                          by. Grouping_specifiers is only meaningful if
                          group_by is provided as well. In case of grouping by
                          index, the grouping  specifiers should be in a
                          dictionary where the key denotes the name of the
                          group.
    ylabels : dict, optional
              ylabels is a dictionary with the outcome names as keys, the
              specified values will be used as labels for the y axis.
    legend : bool, optional
             if true, and group_by is given, show a legend.
    point_in_time : float, optional
                    the point in time at which the scatter is to be made. If
                    None is provided (default), the end states are used.
                    point_in_time should be a valid value on time
    filter_scalar: bool, optional
                   remove the non-time-series outcomes. Defaults to True.

    Returns
    -------
    fig : Figure instance
          the figure instance
    axes  : dict
            key is tuple of names of outcomes, value is associated axes
            instance

    .. note:: the current implementation is limited to seven different
              categories in case of column, categories, and/or
              discretesize. This limit is due to the colors specified
              in COLOR_LIST.

    """

    _logger.debug("generating pairwise scatter plot")

    prepared_data = prepare_pairs_data(
        experiments,
        outcomes,
        outcomes_to_show,
        group_by,
        grouping_specifiers,
        point_in_time,
        filter_scalar,
    )
    experiments, outcomes, outcomes_to_show, grouping_labels = prepared_data

    grid = gridspec.GridSpec(len(outcomes_to_show), len(outcomes_to_show))
    grid.update(wspace=0.1, hspace=0.1)

    # the plotting
    figure = plt.figure()
    axes_dict = {}

    combis = [(field1, field2) for field1 in outcomes_to_show for field2 in outcomes_to_show]

    for field1, field2 in combis:
        i = list(outcomes_to_show).index(field1)
        j = list(outcomes_to_show).index(field2)
        ax = figure.add_subplot(grid[i, j])
        axes_dict[(field1, field2)] = ax

        if group_by:
            for x, group in enumerate(grouping_labels):
                y_data = outcomes[group][field1]
                x_data = outcomes[group][field2]

                facecolor = get_color(x)
                edgecolor = "k"
                if i == j:
                    facecolor = "white"
                    edgecolor = "white"
                ax.scatter(x_data, y_data, facecolor=facecolor, edgecolor=edgecolor)
        else:
            y_data = outcomes[field1]
            x_data = outcomes[field2]

            facecolor = get_color(0)
            edgecolor = "k"
            if i == j:
                facecolor = "white"
                edgecolor = "white"
            ax.scatter(x_data, y_data, facecolor=facecolor, edgecolor=edgecolor)
        do_text_ticks_labels(ax, i, j, field1, field2, ylabels, outcomes_to_show)

    if group_by and legend:
        gs1 = grid[0, 0]

        for ax in figure.axes:
            gs2 = ax._subplotspec
            if all((gs1._gridspec == gs2._gridspec, gs1.num1 == gs2.num1, gs1.num2 == gs2.num2)):
                break

        make_legend(grouping_labels, ax, legend_type=LegendEnum.SCATTER)

    return figure, axes_dict


def do_text_ticks_labels(ax, i, j, field1, field2, ylabels, outcomes_to_show):
    """

    Helper function for setting the tick labels on the axes correctly on and
    off

    Parameters
    ----------
    ax : axes
    i : int
    j : int
    field1 : str
    field2 : str
    ylabels : dict, optional
    outcomes_to_show : str


    """

    # text and labels
    if i == j:
        # only plot the name in the middle
        if ylabels:
            text = ylabels[field1]
        else:
            text = field1
        ax.text(
            0.5,
            0.5,
            text,
            horizontalalignment="center",
            verticalalignment="center",
            transform=ax.transAxes,
        )

    # are we at the end of the row?
    if i != len(outcomes_to_show) - 1:
        # xaxis off
        ax.set_xticklabels([])
    else:
        if ylabels:
            try:
                ax.set_xlabel(ylabels.get(field2))
            except KeyError:
                _logger.info("no label specified for " + field2)
        else:
            ax.set_xlabel(field2)

    # are we at the end of the column?
    if j != 0:
        # yaxis off
        ax.set_yticklabels([])
    else:
        if ylabels:
            try:
                ax.set_ylabel(ylabels.get(field1))
            except KeyError:
                _logger.info("no label specified for " + field1)
        else:
            ax.set_ylabel(field1)
