from .utility import r_ggsave

import sys
import pandas as pd
import rpy2.robjects as ro
import rpy2.robjects.lib.ggplot2 as ggplot2


class _Pasta:
    # contains private stuff for RatPasta and future *Pasta classes
    def __init__(self):
        self._ratpasta = self._import_ratpasta()

    @staticmethod
    def _import_ratpasta():
        try:
            return ro.packages.importr('ratPASTA')
        except ro.packages.PackageNotInstalledError:
            print("ERROR: The R ratPASTA package isn't installed!\n\n"
                  "You can install it manually, or you can use the r_install_package()\n"
                  "convenience function provided in the .utils subpackage\n"
                  "to have R install the package. For more information, see:\n\n"
                  "    help(", __name__, ".utility.r_install_package)\n",
                  sep="", file=sys.stderr)
            raise

    @staticmethod
    def _convert_r2py(x):
        from rpy2.robjects import pandas2ri
        with ro.conversion.localconverter(ro.default_converter + pandas2ri.converter):
            return ro.conversion.rpy2py(x)


class RatPasta(_Pasta):
    """The RatPasta class used for wrapping the ratPASTA functionality of the ratPASTA R package.

    This is a Python class which wraps all ratPASTA R functionality. When instantiated, all
    parameters are passed to the load_startle_data() method, which stores the working data
    as an R dataframe in self.df.
    The class also stores its methods' outputs into appropriate attributes, listed below.
    The last generated plot is stored in self.last_plot. All plots may be modified by appending
    rpy2.robjects.lib.ggplot2 functions' outputs to them using the '+' operator (for more information
    on available parameters, please see rpy2's ggplot2 library documentation).
    All parameters are passed to the load_startle_data() method's R function - only some are explained here.
    For more information, please refer to ratPASTA documentation.

    :param addhead: Correction factor for the animal's latency (before the pulse).
    :type addhead: float
    :param addtail: Correction factor for the animal's latency (after the pulse).
    :type addtail: float
    :param *args: Arguments, passed verbatim to the respective R function.
    :param *kwargs: Keyword arguments, passed verbatim to the respective R function.
    :param *args: Arguments, passed verbatim to the load_startle_data() method.
    :param *kwargs: Keyword arguments, passed verbatim to the load_startle_data() method.

    :var df: R dataframe of raw .pasta data imported using load_startle_data(), passed
             by default to the summary and plotting methods.
    :vartype df: rpy2.robjects.vectors.DataFrame
    :var summary: A dictionary containing three pandas dataframes: "values", "ratio" and "ratio_s",
                  automatically (re)generated on every data load.
    :vartype summary: dict of pandas.core.frame.DataFrame
    :var plots: A dictionary containing plots generated by the last invocation of each of the
                plotting methods. Unlike "basic_startle_plot" and "startle_plot" keys in this dict
                which contain a ggplot, the "latency_plot" key is a nested dictionary containing
                "latency_vs_cycle" and "latency_vs_group" keys, which contain the corresponding
                latency plots.
                All values default to None if the respective method hasn't been invoked yet.
    :vartype plots: nested dict of rpy2.robjects.lib.ggplot2.GGPlot
    :var last_plot: Contains the last generated plot by any of the plotting methods.
    :vartype last_plot: rpy2.robjects.lib.ggplot2.GGPlot

    >>> rp = pastawrap.RatPasta(addhead=0.5, addtail=0.5)

    >>> rp.startle_plot(type=2)
    >>> import rpy2.robjects.lib.ggplot2 as ggplot2
    >>> rp.last_plot += ggplot2.theme_light()
    """
    def __init__(self, *args, **kwargs):
        super().__init__()

        self._grdevices = ro.packages.importr("grDevices")

        self.summary = {"values": None,
                        "ratio": None,
                        "ratio_s": None}
        self.plots = {"basic_startle_plot": None,
                      "startle_plot": None,
                      "latency_plot": {"latency_vs_cycle": None,
                                       "latency_vs_group": None}
                      }
        self.last_plot = None

        self.df = self.load_startle_data(*args, **kwargs)
        self.summarise_startle()

    def load_startle_data(self, *args, **kwargs):
        """Load startle data. Calls ratPASTA's loadStartleData()

        Clears all instance attributes on load and regenerates self.summary so as to
        avoid confusing freshly loaded data with plots and summaries of old data.
        All parameters are passed to the respective R function - only some are explained here.
        For more information, please refer to ratPASTA documentation.

        :param addhead: Correction factor for the animal's latency (before the pulse).
        :type addhead: float
        :param addtail: Correction factor for the animal's latency (after the pulse).
        :type addtail: float
        :param *args: Arguments, passed verbatim to the respective R function.
        :param *kwargs: Keyword arguments, passed verbatim to the respective R function.

        :return: R dataframe which can be passed to pastaWRAP plotting methods.
        :rtype: rpy2.robjects.vectors.DataFrame"""
        self.summary = {"values" : None,
                        "ratio" : None,
                        "ratio_s" : None}
        self.plots = {"basic_startle_plot" : None,
                      "startle_plot": None,
                      "latency_plot": {"latency_vs_cycle": None,
                                       "latency_vs_group": None}
                      }
        self.last_plot = None
        self.df = self._ratpasta.loadStartleData(*args, **kwargs)
        self.summarise_startle()
        return self.df

    def summarise_startle(self, df=None, *args, **kwargs):
        """Generate a summary. Calls ratPASTA's summariseStartle()

        All parameters are passed to the respective R function.
        For more information, please refer to ratPASTA documentation.

        :param df: R dataframe to be processed, defaults to self.df
                   (auto-generated using load_startle_data() when instantiating the RatPasta class)
        :param *args: Arguments, passed verbatim to the respective R function.
        :param *kwargs: Keyword arguments, passed verbatim to the respective R function.

        :return: A dictionary containing three pandas dataframes: "values", "ratio" and "ratio_s",
                 which is also stored in self.summary.
        :rtype: dict of pandas.core.frame.DataFrame"""
        if not df:
            df = self.df

        self.summary["values"], self.summary["ratio"], self.summary["ratio_s"] = \
            [pd.DataFrame(self._convert_r2py(x)) # convert to a pandas df
             for x
             in self._ratpasta.summariseStartle(df, *args, **kwargs)]

        return self.summary

    def basic_startle_plot(self, df=None, *args, **kwargs):
        """Plotting method. Calls ratPASTA's basicStartlePlot()

        All parameters are passed to the respective R function.
        For more information, please refer to ratPASTA documentation.

        :param df: R dataframe to be plotted, defaults to self.df
                   (auto-generated using load_startle_data() when instantiating the RatPasta class)
        :param *args: Arguments, passed verbatim to the respective R function.
        :param *kwargs: Keyword arguments, passed verbatim to the respective R function.

        :return: A ggplot, which is also stored in self.plots and self.last_plot.
        :rtype: rpy2.robjects.lib.ggplot2.GGPlot"""
        if not df:
            df = self.df
        self.plots["basic_startle_plot"] = \
            ggplot2.GGPlot(self._ratpasta.basicStartlePlot(df, *args, **kwargs))
        self.last_plot = self.plots["basic_startle_plot"]
        return self.last_plot

    def startle_plot(self, df=None, type=None, *args, **kwargs):
        """Plotting method. Calls ratPASTA's startlePlot()

        All parameters are passed to the respective R function.
        For more information, please refer to ratPASTA documentation.

        :param df: R dataframe to be plotted, defaults to self.df
                   (auto-generated using load_startle_data() when instantiating the RatPasta class)
        :type df: rpy2.robjects.vectors.DataFrame
        :param type: Type of startle plot. Please refer to ratPASTA documentation for available types.
        :type type: int
        :param *args: Arguments, passed verbatim to the respective R function.
        :param *kwargs: Keyword arguments, passed verbatim to the respective R function.

        :return: A ggplot, which is also stored in self.plots and self.last_plot.
        :rtype: rpy2.robjects.lib.ggplot2.GGPlot"""
        if not df:
            df = self.df
        if not type:
            raise ValueError("Please specify plot type!")
        self.plots["startle_plot"] = \
            ggplot2.GGPlot(self._ratpasta.startlePlot(df, type, *args, **kwargs))
        self.last_plot = self.plots["startle_plot"]
        return self.last_plot

    def latency_plot(self, df=None, *args, **kwargs):
        """Plotting method. Calls ratPASTA's latencyPlot()

        All parameters are passed to the respective R function.
        For more information, please refer to ratPASTA documentation.

        :param df: R dataframe to be plotted, defaults to self.df
                   (auto-generated using load_startle_data() when instantiating the RatPasta class)
        :type df: rpy2.robjects.vectors.DataFrame
        :param *args: Arguments, passed verbatim to the respective R function.
        :param *kwargs: Keyword arguments, passed verbatim to the respective R function.

        :return: A dict containing two ggplots: "latency_vs_cycle" and "latency_vs_group"
        :rtype: dict of rpy2.robjects.lib.ggplot2.GGPlot"""
        if not df:
            df = self.df
        self.plots["latency_plot"]["latency_vs_cycle"], self.plots["latency_plot"]["latency_vs_group"] = \
            [ggplot2.GGPlot(x)
             for x
             in self._ratpasta.latencyPlot(df, *args, **kwargs)]
        self.last_plot = self.plots["latency_plot"]
        return self.last_plot

    def plot_show(self, plot_key=None):
        """Shows the ggplot2 plotting window.

        Displays the last plot generated (stored in self.last_plot) in
        a ggplot2 plotting window.

        :param plot_key: Currently only used when displaying a latency plot, in which case
                    it can be either "latency_vs_cycle" or "latency_vs_group", defaults to None.
        :type plot_key: str

        :return: None

        >>> rp.plot_show() # usual usage
        >>> rp.plot_show(plot_key="latency_vs_group") # for latency plots
        """
        if not self.last_plot:
            raise ValueError("last_plot is empty: Nothing to plot.\n"
                             "Please run one of the plot methods first.")
        if type(self.last_plot) is dict:
            if not plot_key:
                raise KeyError("There are multiple plots. Please specify plot_key!")
            else:
                self.last_plot[plot_key].plot()
        else:
            self.last_plot.plot()
        return

    def plot_close(self):
        """Closes the ggplot2 plotting window."""
        self._grdevices.dev_off()
        return

    def plot_save(self, plot=None, plot_key=None, *args, **kwargs):
        """Saves a plot using ggsave.

        Saves the last plot generated, or the one specified by 'plot', using
        ggsave. For more information, please refer to ggplot2's ggsave documentation.

        :param plot: Plot to save, defaults to the last generated plot (self.last_plot).
        :type plot: str

        :return: None

        >>> rp.plot_save(filename="test.png", width=10, height=8, dpi=150, units="in", device="png")
        """
        if not plot:
            if not self.last_plot:
                raise ValueError("last_plot is empty and no plot is specified: Nothing to save.\n"
                                 "Please run one of the plot methods first.")
            elif type(self.last_plot) is dict:
                if not plot_key:
                    raise KeyError("There are multiple plots in last_plot. Please specify plot_key!")
                else:
                    plot = self.last_plot[plot_key]
            else:
                plot = self.last_plot
        r_ggsave(plot=plot, *args, **kwargs)
        return
