"""
The idea is to make a generic ggplot-ish engine. Then, you can run "to_mpl()", "to_plotly()", etc. to
generate the graphic.

The goal is to provide a flexible plotting framework.
>>> from uplot import *
>>> plot(x, y)
"""

# https://altair-viz.github.io/getting_started/overview.html

from IPython import get_ipython
from IPython.display import display, HTML, SVG
from copy import deepcopy
import numpy as np
import os.path
import pandas as pd
import sys
from tempfile import NamedTemporaryFile
from typelike import ArrayLike, NumberLike
import yaml

# test

# TODO facet wrap (?) pandas does this, maybe back into it that way
# this is not a high priority for me
"""
>>> import uplot as u
>>> # facet_wrap => for every column in data, create line
>>> u.figure(data) + u.line() + u.facet_wrap(n_col=4)  
>>> # facet_grid => for every unique label in column, create and xy line
>>> u.figure(data) + u.line() + u.facet_grid(column='label', x='x', y='y', n_col=4)    
"""

# Include directory
include_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), '_include')

# TODO lazy load this please
with open(os.path.join(include_dir, 'mpl_markers.yml'), 'r') as stream:
    mpl_markers = yaml.safe_load(stream.read())

with open(os.path.join(include_dir, 'mpl_line_styles.yml'), 'r') as stream:
    mpl_line_styles = yaml.safe_load(stream.read())


# Style defaults
# TODO make this customizable; allow user to specify different style in options
# https://matplotlib.org/3.2.2/gallery/style_sheets/style_sheets_reference.html
def set_mpl_theme():
    import matplotlib.pyplot as plt
    plt.style.use(os.path.join(include_dir, 'uplot.mplstyle'))


_style_defaults = {
    'legend': False
}


# ujet
def _jet(n=256):
    ones = np.ones(n)
    zeros = np.zeros(n)
    incr = np.linspace(0., 1., n)
    decr = incr[::-1]

    b2t = np.vstack([zeros, incr, ones]).T  # 001 -> 011
    t2g = np.vstack([zeros, ones, decr]).T  # 011 -> # 010
    g2y = np.vstack([incr, ones, zeros]).T  # 010 -> 110
    y2r = np.vstack([ones, decr, zeros]).T  # 110 -> 100

    return np.vstack([b2t, t2g, g2y, y2r])


jet = _jet(n=256)


# https://ggplot2-book.org/polishing.html
# TODO one day think about SVG bounding box https://matplotlib.org/3.2.2/api/backend_svg_api.html
class Figure:
    def __init__(self, data=None, style=None):
        self._data = data
        self._style = style if style is not None else {}
        self._geometry_objects = []

    # Add a new Geometry
    def __add__(self, other):
        # Deepcopy self
        obj = deepcopy(self)

        # Type check
        if not isinstance(other, Geometry):
            raise AttributeError('must be Geometry instance')

        # Add Geometry
        obj.add_geometry(other)

        # Return
        return obj

    def __repr__(self):
        if not _is_pycharm():
            self.to_mpl(show=True)
        return super().__repr__()

    @property
    def geometry_objects(self):
        """
        Access the Geometry objects of the Figure.

        Returns
        -------
        list
        """

        # If there is nothing, return points
        if not self._geometry_objects:
            return [Line(style={'line_style': 'none', 'marker': 'circle'})]

        # Otherwise, return our saved geometry objects
        return self._geometry_objects

    # Add Geometry
    def add_geometry(self, geometry):
        """
        Add a child instance of Geometry to the Figure in place.

        Parameters
        ----------
        geometry : Geometry
            Child instance of Geometry.
        """

        # Make sure we have an instance of Geometry
        if not isinstance(geometry, Geometry):
            raise AttributeError('must be child of Geometry')

        # Tell the geometry who owns it
        geometry._figure = self

        # Add the figure object
        self._geometry_objects.append(geometry)

    # Get style from the dictionary
    def get_style(self, style, default=None, index=None):
        if style in self._style:
            result = self._style[style]
        else:
            result = _style_defaults.get(style, default)

        if isinstance(result, ArrayLike) and index is not None:
            result = result[index]
        return result

    # Show figure
    def show(self):
        """
        Show Figure using the default backend.
        """

        self.to_mpl(show=True)

    # Convert figure to matplotlib
    # noinspection PyShadowingNames
    def to_mpl(self, show=False, save_as=None):
        """
        Convert to matplotlib objects.

        Parameters
        ----------
        show : bool
            Should the figure be displayed?
        save_as : None or str
            If str is provided, filename to save the figure as.

        Returns
        -------
        matplotlib.pyplot.Figure, matplotlib.pyplot.Axis
            Returned only if `show` = False.
        """

        # Make sure pyplot is loaded
        import matplotlib.pyplot as plt
        set_mpl_theme()

        # Create the figure and axis
        height, width = self.get_style('height'), self.get_style('width')
        figsize = None
        if height is not None and width is not None:
            figsize = (width, height)
        figure = plt.figure(figsize=figsize)  # type: plt.Figure
        axis = figure.add_subplot()

        # Iterate through geometry objects and draw
        for geometry in self.geometry_objects:
            # noinspection PyProtectedMember
            geometry._to_mpl(figure, axis)

        # Canvas
        # figure.patch.set_facecolor(self.get_style('background'))

        # Set plot elements
        figure.suptitle(self.get_style('title'))
        axis.set_xlabel(self.get_style('x_title'))
        axis.set_ylabel(self.get_style('y_title'))
        axis.set_xlim(self.get_style('x_min'), self.get_style('x_max'))
        axis.set_ylim(self.get_style('y_min'), self.get_style('y_max'))
        # TODO https://matplotlib.org/3.1.0/api/dates_api.html#matplotlib.dates.IndexDateFormatter set up date format
        # n_minor_ticks, n_major_ticks = self.get_style('n_minor_ticks'), self.get_style('n_major_ticks')
        # if n_minor_x_ticks is not None:
        #     # axis.tick_params(which='minor', length=n_minor_ticks)
        #     axis.xaxis.set_minor_locator(plt.MaxNLocator(n_minor_ticks))
        # if n_x_major_x_ticks is not None:
        #     # axis.tick_params(which='major', length=n_major_ticks)
        #     axis.xaxis.set_major_locator(plt.MaxNLocator(n_major_ticks))
        axis.tick_params(axis='x', labelrotation=self.get_style('x_rotation'))

        # Legend
        if self.get_style('legend'):
            axis.legend(bbox_to_anchor=(1., 0.5), loc='center left')

        # Should we save?
        if save_as is not None:
            # TODO there will have to be a way to include some arguments here
            figure.savefig(save_as)

        # Return
        if show:
            if _is_jupyter():
                # _display_svg(figure)
                return None
            else:
                figure.show()
            plt.close(figure)
            return None
        else:
            return figure, axis

    def to_plotnine(self):
        # https: // stackoverflow.com / questions / 19921842 / plotting - multiple - time - series - on - the - same - plot - using - ggplot
        pass

    # noinspection PyShadowingNames
    def to_plotly(self, height='4.8in', width='6.4in'):
        import plotly.graph_objects as go
        figure = go.Figure()
        for geometry in self.geometry_objects:
            # noinspection PyProtectedMember
            geometry._to_plotly(figure)
        display(HTML(figure.to_html(default_height=height, default_width=width)))


class Geometry(Figure):
    def __init__(self, data=None, style=None):
        super().__init__(data, style)
        self._figure = None

    def __add__(self, other):
        raise NotImplementedError

    def add_geometry(self, geometry):
        raise NotImplementedError

    # Get data from the Geometry first, then its parent
    def _get_data(self):
        # Get data from the Geometry first, then its parent
        data = self._data
        if data is None and isinstance(self._figure, Figure):
            data = self._figure._data

        # If there is still no data, throw an error
        if data is None:
            raise AttributeError('data does not exist')

        # If the data is not a DataFrame, throw an error
        if not isinstance(data, pd.DataFrame):
            raise AttributeError('data must be pandas DataFrame')

        # Return
        return data


# class Style:
#     def __init__(self, **kwargs):
#         self._kwargs = kwargs


class Bar(Geometry):
    def __init__(self, data=None, style=None):
        super().__init__(data, style)

    # noinspection PyShadowingNames
    def _to_mpl(self, figure, axis):
        # Get data
        data = self._get_data()
        x = data.index.values

        # Loop over all columns
        for i, column in enumerate(self._data.columns):
            # Get y for column
            y = self._data[column].values
            label = self.get_style('label', default=column, index=i)
            axis.bar(x, y, label=label)


class Error(Geometry):
    def __init__(self, data=None, style=None):
        super().__init__(data, style)

    # noinspection PyShadowingNames
    def _to_mpl(self, figure, axis):
        # Get data
        data = self._get_data()
        x = data.index.values

        # Loop over all columns
        for i, column in enumerate(data.columns):
            # Get y for column
            y = data[column].values
            label = self.get_style('label', default=column, index=i)
            color = self.get_style('color', index=i)
            line_style = self.get_style('line_style', index=i)
            if line_style is not None:
                line_style = mpl_line_styles[line_style]
            marker = self.get_style('marker', index=i)
            if marker is not None:
                marker = mpl_markers[marker]
            axis.plot(x, y, label=label, color=color, linestyle=line_style, marker=marker)


class Line(Geometry):  # noqa
    def __init__(self, data=None, style=None):
        super().__init__(data, style)

    # noinspection PyShadowingNames
    def _to_mpl(self, figure, axis):
        # Get data
        data = self._get_data()
        x = data.index.values

        # Loop over all columns
        for i, column in enumerate(data.columns):
            # Get y for column
            y = data[column].values

            label = self.get_style('label', default=column, index=i)
            color = self.get_style('color', index=i)

            # TODO are we in love with camel_case here?
            line_width = self.get_style('line_width', index=i)

            line_style = self.get_style('line_style', index=i)
            if line_style is not None:
                line_style = mpl_line_styles[line_style]

            marker = self.get_style('marker', index=i)
            if marker is not None:
                marker = mpl_markers[marker]

            axis.plot(x, y, label=label, color=color, linestyle=line_style, linewidth=line_width, marker=marker)

    # noinspection PyShadowingNames
    def _to_plotly(self, figure):
        import plotly.graph_objects as go

        # Get data
        data = self._get_data()
        x = data.index.values

        # Loop over all columns
        for i, column in enumerate(data.columns):
            # Get y for column
            y = data[column].values
            label = self.get_style('label', default=column, index=i)
            # marker = self.get_style('marker', index=i)
            # if marker is not None:
            #     marker = markers_mpl[marker]
            figure.add_trace(go.Scatter(x=x, y=y, mode='lines', name=label))


#
# class Point(Geometry):
#     def __init__(self, data=None, style=None):
#         super().__init__(data, style)
#
#     # noinspection PyShadowingNames
#     def to_mpl(self, figure, axis):
#         x = self._data.index.values
#         for column in self._data.columns:
#             y = self._data[column].values
#             # TODO label override from style
#             axis.plot(x, y, 'o', label=column)


def figure(data=None, style=None):
    """
    Convenience function to create a Fine object.

    Parameters
    ----------
    data : pandas.DataFrame
        If provided, uses the data in `data` for the figure.
    style : dict
        If provided, list of style elements.

    Returns
    -------
    Figure
        Instance of Figure object.
    """

    # Create Figure
    element = Figure(data, style)

    # Set figure attributes
    # ...

    # Return
    return element


def bar(x=None, y=None, style=None):
    data = _coerce_x_y(x, y)
    style = _coerce_style(style, defaults={'line_style': 'solid'})
    geometry = Bar(data, style)
    return geometry


def error(x=None, y=None, y_err=None, style=None):
    pass


def line(x=None, y=None, style=None, **kwargs):
    """
    Convenience function to create a Line object.

    Parameters
    ----------
    x : ArrayLike
        (Optional) If present, `x` dimension for plotting.
    y : ArrayLike
        (Optional) If present, `y` dimension for plotting. If this is an array of arrays, every interior array will be
        treated as a dependent variable to `x`.
    style : dict
        If provided, list of style elements.

    Returns
    -------
    Line
        Instance of Line object.
    """

    data = _coerce_x_y(x, y)
    style = _coerce_style(style, defaults={'line_style': 'solid'})
    element = Line(data, style)
    return element


def hline(y=None, style=None):
    if not isinstance(y, NumberLike):
        raise AttributeError('y must be a number')
    style = _coerce_style(style, defaults={'line_style': 'solid'})
    # return HorizontalLine(y, style)
    pass


def point(x=None, y=None, style=None):
    data = _coerce_x_y(x, y)
    style = _coerce_style(style, defaults={'marker': 'circle'})
    element = Line(data, style=style)
    return element


# Extract label from pandas if possible
def _get_label(x, default='x'):
    label = default
    if isinstance(x, pd.Series):
        label = x.name
    return label


# Coerce x and y into expected types and formats
# TODO is there a need for this and _coerce_data_x_y?? I don't think so.
def _coerce_x_y(x, y):
    # Sanity check
    if x is None:
        # raise AttributeError('cannot parse data')
        return None

    # Create data
    data = None
    if y is None:
        data = x
        x = None

    # Coerce data, x, and y
    return _coerce_data_x_y(data, x, y)


# Coerce data and x and y into expected types and formats
def _coerce_data_x_y(data, x, y):
    # If data is not set
    if data is None:
        # If both of x and y are arrays, build DataFrame
        if isinstance(x, ArrayLike) and isinstance(y, ArrayLike):
            # Make sure y is an array of arrays (for simplicity)
            if isinstance(y, pd.Series) or not isinstance(y[0], ArrayLike):
                y = [y]

            # Build DataFrame
            x_label = _get_label(x, default='x')
            data = pd.DataFrame({x_label: x}).set_index(x_label)
            len_y = len(y)
            for i, yi in enumerate(y):
                if len_y == 1:
                    y_label_default = 'y'
                else:
                    y_label_default = 'y' + str(i)
                y_label = _get_label(yi, default=y_label_default)
                data[y_label] = yi if not isinstance(yi, pd.Series) else yi.values

        # If data is not set, and x and y are not arrays, we don't know what we're doing
        elif x is not None or y is not None:
            raise AttributeError('must pass data or x and y arrays')

    # Otherwise, if data is a DataFrame, extract out x and y columns
    elif isinstance(data, pd.DataFrame):
        # Create a copy of the DataFrame
        data = data.copy()

        # Set the index
        if x is not None:
            data = data.set_index(x)

        # Extract columns (need to create a copy so we're not working on a slice)
        if y is not None:
            data = data[y].copy()

    # Otherwise, if data is a Series, convert to DataFrame
    elif isinstance(data, pd.Series):
        data = data.to_frame()

    # Else
    else:
        raise AttributeError('cannot process input')

    # Return the DataFrame
    return data


def _coerce_style(style, defaults=None):
    if style is None:
        style = {}
    style = {key.lower(): value for key, value in style.items()}
    if defaults is not None:
        for key in defaults:
            if key not in style:
                style[key] = defaults[key]
    return style


#
# (
#     figure(style={'xtitle': '$x$', 'ytitle': '$y$', 'legend': True})
#     + line([1, 2, 3], [4, 5, 6], style={'marker': 'circle', 'label': 'y1'})
#     + line([1, 2, 3], [6, 5, 4], style={'marker': 'circle', 'label': 'y2'})
# ).to_mpl(show=True)

# Check if we are in PyCharm
def _is_pycharm():
    return '_pydev_imps' in sys.modules


# Check if we are in a Jupyter window
def _is_jupyter():
    """

    Exclude QtConsole:
    'qtconsole' not in sys.modules

    Exclude PyCharm:
    '_pydev_imps' not in sys.modules

    Returns
    -------

    """
    return get_ipython() and 'qtconsole' not in sys.modules and not _is_pycharm()


def _data_or_x(data_or_x):
    if isinstance(data_or_x, pd.DataFrame):
        data = data_or_x
        x = None

    else:
        data = None
        x = data_or_x

    return data, x


# Display SVG in IPython (for matplotlib)
def _display_svg(fig):
    with NamedTemporaryFile(delete=False) as tempfile:
        filename = str(tempfile.name) + '.svg'
    fig.savefig(fname=filename, transparent=True)
    display(SVG(filename))
    os.remove(filename)
