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

__all__ = ['warnError', 'rPrint', 'R_CONSOLE', 'import_rpy2', 'pkg_name', 'utils', 'lib_loc', 'libPaths', 'SUPERPOWER',
           'types', 'DTYPES', 'AS_DTYPE', 'Result', 'rpy2py_intvector', 'rpy2py_floatvector', 'rpy2py_strvector',
           'rpy2py_boolvector', 'rpy2py_listsexpvector', 'converter', 'rpy2py_langsexpvector',
           'rpy2py_pairlistsexpvector', 'py2rpy_list', 'py2rpy_list', 'py2rpy_none', 'RFunction', 'MatrixMixin']

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

# Cell
from .log import getLogger, WIDGET_HANDLER
from .config import CONFIG

# Cell
import rpy2
from rpy2 import rinterface, robjects, situation
rpy2.rinterface.initr()

from rpy2.robjects import packages, numpy2ri, pandas2ri, conversion
from rpy2.rinterface_lib.embedded import RRuntimeError

# Cell
R_CONSOLE = getLogger('r.console')
def warnError(msg):
    R_CONSOLE.error(msg)

def rPrint(msg):
    R_CONSOLE.info(msg)

rpy2.rinterface_lib.callbacks.consolewrite_warnerror = warnError
rpy2.rinterface_lib.callbacks.consolewrite_print = rPrint

# Cell
def import_rpy2():
    global rpy2
    return rpy2

# Cell
pkg_name = 'Superpower'
utils = rpy2.robjects.packages.importr('utils')
lib_loc = CONFIG['R']['lib_loc']
libPaths = list(rpy2.robjects.r['.libPaths']())
libPaths.insert(0, CONFIG['R']['lib_loc'])
SUPERPOWER = None

for lib_path in libPaths:
    try:
        R_CONSOLE.info('Attempting Superpower import from %s...' % lib_path)
        SUPERPOWER = rpy2.robjects.packages.importr(pkg_name, lib_loc)
        if SUPERPOWER:
            R_CONSOLE.info('Import successful!')
            break
    except:
        pass
if not SUPERPOWER:
    R_CONSOLE.exception('Import failed.')
    R_CONSOLE.info('Using source file %s' % lib_loc)
    R_CONSOLE.info('Attempting install to %s' % lib_loc)
    try:
        kwargs = {}
        kwargs['pkgs'] = CONFIG['R']['source']
        kwargs['repos'] = rpy2.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:
        R_CONSOLE.exception('Install failed with message: %s' % str(e))
        raise Exception('Fatal. Superpower could not be imported or installed.')
    if SUPERPOWER:
        R_CONSOLE.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(str(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
converter = conversion.Converter('superpower converter')

@converter.rpy2py.register(rinterface.IntSexpVector)
def rpy2py_intvector(obj):
    if 'factor' in obj.rclass or len(obj) != 1:
        return robjects.pandas2ri.rpy2py(obj)
    else:
        return obj[0]

@converter.rpy2py.register(rinterface.FloatSexpVector)
def rpy2py_floatvector(obj):
    res = robjects.pandas2ri.rpy2py_floatvector(obj)
    if len(res) == 1:
        return obj[0]
    else:
        return res

@converter.rpy2py.register(rinterface.StrSexpVector)
def rpy2py_strvector(obj):
    if len(obj) == 1:
        return obj[0]
    else:
        return [s for s in obj]

@converter.rpy2py.register(rinterface.BoolSexpVector)
def rpy2py_boolvector(obj):
    if len(obj) == 1:
        return obj[0]
    else:
        return [s for s in obj]

@converter.rpy2py.register(rinterface.ListSexpVector)
def rpy2py_listsexpvector(obj):
    # recursive base case
    if not isinstance(obj, rinterface.ListSexpVector):
        if isinstance(obj, rinterface.NULLType):
            return None
        else:
            return converter.rpy2py(obj)

    if 'data.frame' in obj.rclass: # if data frame
        res = pd.DataFrame(robjects.ListVector(obj)) #robjects.pandas2ri.rpy2py_dataframe(obj)
    elif isinstance(obj.names, rinterface.NULLType): # if list isn't named
        return [rpy2py_listsexpvector(value) for value in obj]
    else: # if named list
        res = {}
        for key, value in zip(obj.names, obj):
            res[key] = rpy2py_listsexpvector(value)
    return res

# this kills creation of ListVectors
#@converter.rpy2py.register(rinterface.NULLType)
#def rpy2py_null(obj):
#    return None

# Cell
@converter.rpy2py.register(rinterface.LangSexpVector)
def rpy2py_langsexpvector(obj):
    ''' This is for when there is an r default like "c('one', 'two')" '''
    obj = [rpy2py_strvector(item) for item in obj if isinstance(item, rinterface.StrSexpVector)]
    return obj

@converter.rpy2py.register(rinterface.PairlistSexpVector)
def rpy2py_pairlistsexpvector(obj):
    ''' This one is important for getting default arguments with the formals function'''
    ret = {}
    for item in obj:
        if item in converter.rpy2py.registry.keys():
            ret.update(rpy2py_listsexpvector(item))
        if item in robjects.default_converter.rpy2py.registry.keys():
            ret.update(rpy2py_listsexpvector(item))
    return ret

# Cell
@converter.py2rpy.register(list)
def py2rpy_list(obj):
    if len(obj) == 1:
        with conversion.localconverter(robjects.default_converter + robjects.numpy2ri.converter + robjects.pandas2ri.converter + converter) as cv:
            return cv.py2rpy(obj[0])
    if len(obj) > 1:
        first = obj[0]
        if all(isinstance(x, type(first)) for x in obj):
            if isinstance(first, bool):
                return rinterface.BoolSexpVector(obj)
            elif isinstance(first, str):
                return rinterface.StrSexpVector(obj)
            elif isinstance(first, int):
                return rinterface.IntSexpVector(obj)
            elif isinstance(first, float):
                return rinterface.FloatSexpVector(obj)
    return robjects.default_converter.py2rpy(obj)

# Cell
@converter.py2rpy.register(bool)
def py2rpy_list(obj):
    return rinterface.BoolSexpVector([obj])

@converter.py2rpy.register(type(None))
def py2rpy_none(obj):
    return robjects.NULL

# 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)
        # converter defined in global scope above
        self.converter = robjects.default_converter + robjects.numpy2ri.converter + robjects.pandas2ri.converter + converter
        self.defaults = self.getDefaultKwargs()
        self.logger = getLogger(__name__)

    def setArgNames(self, argNames=None):
        '''Sets the names of all accepted arguments. Overrides previous.'''
        if argNames is None:
            self.getArgNames(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 getArgNames(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'''
        with conversion.localconverter(self.converter) as cv:
            defaults = self.function.formals()
            ret = {}
            return self.null_to_none(defaults)

    def py2rpy(self):
        self.rargs = {}
        self.rargsString = ''
        with conversion.localconverter(self.converter) as cv:
            for key, value in self.kwargs.items():
                rpy = cv.py2rpy(value)
                self.rargs[key] = rpy
                if hasattr(rpy, 'r_repr'):
                    self.rargsString += '\t%s: %s\n' % (key, rpy.r_repr())
                else:
                    self.rargsString += '\t%s: %s\n' % (key, str(rpy))
            return self.rargsString

    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:
            self.py2rpy()
            ret = {}
            startTime = time.time()
            ret['name'] = self.name
            ret['kwargs'] = self.kwargs
            ret['rargs'] = self.rargs
            self.logger.info('Running %s with R arguments\n%s' % (self.name, self.rargsString))
            with conversion.localconverter(self.converter):
                ret['results'] = self.null_to_none(self.function(**self.rargs))
            ret['runtime'] = (time.time() - startTime)
            ret['id'] = int(round(time.time() * 1000))
            ret['date'] = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            self.logger.info('ID: %i, Date: %s, Runtime: %s' % (ret['id'], str(ret['date']), str(ret['runtime'])))
            return Result(ret)
        except Exception as e:
            if alert:
                estring = str(type(e).__name__) + ': '
                estring += str(e).replace('\n', ' ').replace('\n', '').replace('"','\\"')
                display(Javascript('alert("' + estring + '");'))
                if isinstance(type(e), RRuntimeError):
                    robjects.r('traceback(max.lines=1)')
                else:
                    self.logger.error(e, exc_info=True)
                return None
            else:
                raise e

    # we want the converter to do this but it causes problems
    def null_to_none(self, dictionary):
        ret = {}
        for key, value in dictionary.items():
            if isinstance(value, rinterface.NULLType):
                ret[key] = None
            else:
                ret[key] = value
        return ret

    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 MatrixMixin():

    def matrixRFunction(self, **kwargs):
        func = robjects.r('matrix')
        with robjects.conversion.localconverter(robjects.default_converter + robjects.numpy2ri.converter):
            return func(**kwargs)