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

__all__ = ['Item', 'Model', 'RFunction']

# Cell
import time
import pandas as pd
import numpy as np
from traitlets import HasTraits
from rpy2 import robjects, rinterface
from rpy2.robjects import numpy2ri
numpy2ri.activate()

# Cell
from ..rpy import Int, Float, SUPERPOWER

# Cell
class Item(HasTraits):
    '''The <code>name</code> property is used by model collections and view collections
    in order to store and retrieve models and views by name'''
    def __init__(self, name=None, **kwargs):
        super().__init__(**kwargs)
        if name:
            self.name = name

    def _repr_pretty_(self, p, cycle):
        if hasattr(self, 'name'):
            name = self.name
        else:
            name = self.__class__.__name__
        with p.group(4, str(name) +'(', ')'):
            p.breakable()
            for trait, value in self._trait_values.items():
                # > paste
                p.text(trait + ': ') # shift right
                p.pretty(value) # shift right
                p.text(',')
                p.breakable()
        """
                        if trait == 'children':
                    with p.group(4, 'children: {', '}'):
                        p.breakable()
                        if hasattr(self, '_size'):
                            for i in range(self._size):
                                if i:
                                    p.text(',')
                                    p.breakable()
                                p.text(str(i) + ': ')
                                p.pretty(self.children[i])
                else:
        """

# Cell
class Model(Item):

    def __init__(self, **kwargs):
        super().__init__(**kwargs)

# Cell
class RFunction():
    ''' The class used to create a model for any R function, and only one R function.
    If you think you need more than one function in a model, you should use a
    ModelCollection.
    '''
    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 = 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 = robjects.r['formals']
        default = formals(self.function)
        args = {}
        for i, name in enumerate(default.names):
            if len(default[i]) > 0 and default[i][0] != 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):
        '''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
        '''
        try:
            startTime = time.time()
            self.result = self.function(**self.kwargs)
            self.resultsString = self.result.r_repr()
            self.packageName = self.function.__rpackagename__
            self.runtime = (time.time() - startTime)
            self.uniqueID = int(round(time.time() * 1000))
            return self.__dict__
        except Exception as e:
            #e = str(e).replace('\n', ' ').replace('\r', '')
            raise Exception(e)

    def attr_dict(self):
        keys = ['name', 'result', 'packageName', 'runtime', 'uniqueID', '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()