# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/05_rpy.ipynb (unless otherwise specified).

__all__ = ['logger', 'RConsole', 'R_CONSOLE', 'rpy2', 'lib_loc', 'lib_paths', 'pkg_name', 'utils', 'types', 'DTYPES',
           'AS_DTYPE', 'Result', 'RFunction', 'RMixins']

# Cell
import time, os, datetime
import pandas as pd
import numpy as np
from IPython.display import display, Javascript, HTML
from collections import UserDict

# Cell
from .log import HANDLERS
from .config import CONFIG, APP_DIR
logger = HANDLERS.getLogger(__name__)

# Cell
class RConsole():

    def __init__(self):
        self.logger = HANDLERS.getLogger('r.console')
        self.setup()

    def print(self, x):
        self.logger.info(x)

    def warnError(self, x):
        self.logger.error(x)

    def setup(self):
        # output from the R console will now be handled by the above functions
        import rpy2
        from rpy2 import rinterface, robjects, situation
        rpy2.rinterface.initr()
        from rpy2.robjects import packages, numpy2ri, pandas2ri, conversion
        rpy2.robjects.numpy2ri.activate()
        rpy2.robjects.numpy2ri.activate()
        self.rpy2 = rpy2
        rpy2.rinterface_lib.callbacks.consolewrite_warnerror = self.warnError
        rpy2.rinterface_lib.callbacks.consolewrite_print = self.print
        self.cat = rpy2.rinterface.baseenv['cat']
        self.cat('RConsole initialized.')


R_CONSOLE = RConsole()
R_CONSOLE.setup()
rpy2 = R_CONSOLE.rpy2

# Cell
lib_loc = CONFIG['R']['lib_loc']
lib_paths = list(rpy2.robjects.r['.libPaths']())
if lib_loc:
    lib_loc = os.path.expanduser(lib_loc)
    lib_paths.insert(0, lib_loc)

# Cell
pkg_name='Superpower'
utils = rpy2.robjects.packages.importr('utils')
for lib_path in lib_paths:
    try:
        logger.info('Attempting Superpower import from %s...' % lib_path)
        SUPERPOWER = rpy2.robjects.packages.importr(pkg_name, lib_loc)
        if SUPERPOWER:
            logger.info('Import successful!')
            break
    except Exception as e:
        logger.exception('Import failed with message: %s' % str(e))
        logger.info('Attempting install from source to %s' % lib_loc)
        if CONFIG['R']['source']:
            source_path = os.path.expanduser(CONFIG['R']['source'])
            if not os.path.isabs(source_path):
                source_path = os.path.join(APP_DIR, source_path)
            if os.path.isfile(source_path):
                logger.info('Using source file %s' % source_path)
            else:
                raise Exception("Source file not found at %s" % source_path)
        else:
            raise Exception('Superpower not installed and source file not provided.')
        try:
            kwargs = {}
            kwargs['pkgs'] = source_path
            kwargs['repos'] = rinterface.NULL
            kwargs['type'] = 'source'
            kwargs['lib'] = lib_loc
            utils.install_packages(**kwargs)
            SUPERPOWER = rpy2.robjects.packages.importr(pkg_name, lib_loc)
        except Exception as e:
            logger.exception('Install failed with message: %s' % str(e))
            raise Exception('Fatal. Superpower could not be imported or installed.')
        if SUPERPOWER:
            logger.info('Install successful!')
        else:
            raise Exception('Install failed.')

# Cell
types = (('int', rpy2.rinterface.IntSexpVector([])), ('bool', rpy2.rinterface.BoolSexpVector([])), ('string', rpy2.rinterface.StrSexpVector([])), ('float', rpy2.rinterface.FloatSexpVector([])))
DTYPES = {}
AS_DTYPE = {}
for py_type, r_vector in types:
    np_array = np.asarray(r_vector)
    DTYPES[py_type] = np_array.dtype
    AS_DTYPE[py_type] = getattr(np, str(np_array.dtype))

# Cell
class Result(UserDict):

    def _repr_pretty_(self, p, cycle):
        with p.group(4, self.data['name'] +'([', '])'):
            p.breakable()
            for key, value in self.data.items():
                if key != 'name':
                    p.text(key  + ': ')
                    if isinstance(value, dict):
                        with p.group(4, '{', '}'):
                            for k, v in value.items():
                                p.text(k + ': ')
                                p.pretty(v)
                                p.text(', ')
                                p.breakable()
                    elif isinstance(value, pd.DataFrame):
                        p.pretty('See above data frame.')
                        display(HTML(value.to_html()))
                    else:
                        p.pretty(value)
                        p.text(',')
                    p.breakable()

# Cell
class RFunction():
    ''' The class used to create a model for any R function, and only one R function.
    This class abstracts away all rpy2 calls to make interaction with the SUPERPOWER library
    easy.
    '''

    def __init__(self, name, argNames=None):
        self.name = name
        self.function = getattr(SUPERPOWER, name)
        self.defaults = self.getDefaultKwargs()
        self.setArgNames(argNames)

    def setArgNames(self, argNames=None):
        '''Sets the names of all accepted arguments. Overrides previous.'''
        if argNames is None:
            self.getDefaultArgNames(self)
        else:
            if all(name in self.argNames for name in argNames):
                self.argNames = argNames
            else:
                raise Exception('All argNames must be valid arguments for the given R function')

    def setKwargs(self, **kwargs):
        """Adds kwargs if and only if the key is is a valid function argument.
        """
        self.kwargs = {k:v for (k, v) in kwargs.items() if k in self.argNames}

    def setDefaultArg(self, argName):
        self.kwargs[argName] = self.defaults[argName]

    def getDefaultArgNames(self):
        '''Gets the names of all accepted arguments'''
        formals = rpy2.robjects.r['formals']
        default = formals(self.function)
        return default.names

    def getDefaultKwargs(self):
        '''This method gets a dictionary of default arguments provided by the R function'''
        formals = rpy2.robjects.r['formals']
        default = formals(self.function)
        args = {}
        for i, name in enumerate(default.names):
            if len(default[i]) > 0 and default[i][0] != rpy2.robjects.r("NULL"):
                ary = np.asarray(default[i])
                if len(ary) > 1:
                    args[name] = ary
                else:
                    args[name] = ary[0]
        return args

    def run(self, alert=True, **kwargs):
        '''This method runs the R function associated with the model. The run will be
        timed and the runtime and results are added to the results dictionary. You should
        not have to modify this function to run your custom Model
        '''
        if kwargs:
            self.setKwargs(**kwargs)
        try:
            logger.info('Running %s with kwargs...' % self.name)
            for key, value in self.kwargs.items():
                logger.info('\t%s: %s' % (key, str(value).replace("\n", "")))
            ret = {}
            startTime = time.time()
            ret['name'] = self.name
            ret['kwargs'] = self.kwargs
            ret['df'] = self.function(**self.kwargs)
            if isinstance(ret['df'], rpy2.robjects.vectors.ListVector):
                ret['df'] = self.listVector_to_DataFrame(ret['df'])
            ret['df'] = self.convert_to_pandas(ret['df'])
            ret['runtime'] = (time.time() - startTime)
            ret['id'] = int(round(time.time() * 1000))
            ret['date'] = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            logger.info('ID: %i, Date: %s, Runtime: %s' % (ret['id'], str(ret['date']), str(ret['runtime'])))
            logger.info('Results:\n%s' % ret['df'].to_string())
            return Result(ret)
        except Exception as e:
            if alert:
                e = str(e).replace('\n', ' ').replace('\n', '').replace('"','\\"')
                display(Javascript('alert("' + e + '");'))
                return None
            else:
                raise e

    def listVector_to_DataFrame(self, result):
        df_dict = {}
        for i in range(len(result.names)):
            df_dict[result.names[i]] = result[i]
        return rpy2.robjects.DataFrame(df_dict)

    def convert_to_pandas(self, df):
        with rpy2.robjects.conversion.localconverter(rpy2.robjects.default_converter + rpy2.robjects.pandas2ri.converter):
            return rpy2.robjects.conversion.rpy2py(df)

    def attr_dict(self):
        keys = ['name', 'result', 'packageName', 'runtime', 'id', 'kwargs', 'defaults']
        return {key: value for key, value in self.__dict__.items() if key in keys}

    def _repr_pretty_(self, p, cycle):
        with p.group(4, str(self.name) +'([', '])'):
            p.breakable()
            for key, value in self.attr_dict().items():
                if key != 'name':
                    p.text(key  + ': ')
                    if isinstance(value, dict):
                        with p.group(4, '{', '}'):
                            for k, v in value.items():
                                p.text(k + ': ')
                                p.pretty(v)
                                p.text(',')
                                p.breakable()
                    else:
                        p.pretty(value)
                        p.text(',')
                        p.breakable()
                    p.breakable()

# Cell
class RMixins():

    def __init__(self):
        self.matrix_func = rpy2.robjects.baseenv['matrix']