# -*- coding: utf-8 -*-
"""
Plotting utilities
"""
from glob import glob
import pandas as pd
import os
import logging
import seaborn as sns
import matplotlib.pyplot as plt
from osos import DATA_DIR, PLOT_DIR


sns.set_style("darkgrid")
logger = logging.getLogger(__name__)


class Plotting:
    """OSOS plotting utilities for auto-plotting of usage data."""

    # common aliases and capitalization conventions for plot titles
    # format is all lower case keys mapped to desired alias in the value
    ALIASES = {'rev': 'reV',
               'revx': 'reVX',
               'pysam': 'PySAM',
               }

    @staticmethod
    def plot_metric(df, metric, ylabel=None, save_path=None, show=True,
                    close=True):
        """Plot a osos metric from an osos dataframe.

        Parameters
        ----------
        df : pd.DataFrame
            OSOS data timeseries dataframe, likely from the repository DATA_DIR
        metric : str
            Metric to plot in the dataframe. Must be one of the column titles.
        ylabel : str
            Optional y axis label. If None the metric will be used.
        save_path : str | None
            Filepath to save figure to (optional).
        show : bool
            Flag to show / display figure (should be false for automated calls)
        close : bool
            Flag to close figure after saving/showing

        Returns
        -------
        fig : matplotlib.axes._subplots.AxesSubplot
            AxesSubplot object generated by seaborn
        """

        if metric not in df:
            msg = ('Could not find "{}", available keys: {}'
                   .format(metric, df.columns.values.tolist()))
            raise KeyError(msg)

        sns_ax = sns.lineplot(data=df[metric])
        fig = sns_ax.get_figure()
        _ = plt.xticks(rotation=45)

        if ylabel is not None:
            sns_ax.set_ylabel(ylabel)

        if save_path is not None:
            fig.savefig(save_path, bbox_inches='tight')
            logger.info('Saved: {}'.format(save_path))

        if show:
            plt.show()

        if close:
            plt.close()

        return sns_ax

    @classmethod
    def auto_plot(cls, metric, ylabel=None, source_dir=DATA_DIR,
                  save_dir=PLOT_DIR):
        """Auto plot a single metric for all osos datasets in the osos
        source_dir path

        Parameters
        ----------
        metric : str
            Metric to plot in the dataframe. Must be one of the column titles.
        ylabel : str
            y axis label for all plots. Can include format string {name} that
            will be replaced with the repo name from the data filename. If
            None, the ylabel will just be the metric.
        source_dir : str
            Directory to look for osos csv's. This function will try to plot
            the desired metric from all csv's in this directory.
        save_dir : str
            Directory to save the plots. Will be created if it doesnt exist.
        """

        fps = glob(source_dir + '/*.csv')
        names = [os.path.basename(fp).replace('.csv', '') for fp in fps]

        if not os.path.exists(save_dir):
            os.makedirs(save_dir)

        for name, fp in zip(names, fps):
            name = name.lower()
            save_path = os.path.join(save_dir, f'{name}_{metric}.png')

            if name in cls.ALIASES:
                name = cls.ALIASES[name]
            else:
                name = name.upper()

            i_ylabel = ylabel
            if isinstance(i_ylabel, str) and '{name}' in i_ylabel:
                i_ylabel = i_ylabel.format(name=name)

            df = pd.read_csv(fp, index_col=0, parse_dates=True)

            try:
                cls.plot_metric(df, metric, save_path=save_path,
                                ylabel=i_ylabel, show=False, close=True)
            except KeyError:
                logger.info(f'Could not plot metric "{metric}" for "{name}".')
