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

__all__ = ['remove_folder_or_file', 'get_pointer_names', 'global_delete_from_reference', 'global_set_from_reference',
           'global_delete_from_object', 'global_set_from_object', 'SimpleCacher', 'EstimatorCacher']

# Cell
from shutil import rmtree
from pathlib import Path
from warnings import warn
import sys
import os
import gc
from functools import partial
import time


from sklearn.base import BaseEstimator, clone
#import pickle library used
import pickle
import joblib

# Cell
def remove_folder_or_file(path):
    '''
    recursively remove dir and files in dir given by direcotry(path)
    '''
    path = Path(path)
    if path.exists():
        if path.is_dir():
            rmtree(path)
        else:
            path.unlink()
    else:
        pass

    return

# Cell
def get_pointer_names(obj):
    '''
    returns all names in scope() (globals or locals) pointing to object
    '''
    return [k for k,v in globals().items() if v is obj]

def global_delete_from_reference(varname):
    '''
    deletes varname from scope() (globals or locals) and runs gc.collect()
    '''
    globals().pop(varname)
    gc.collect()
    return

def global_set_from_reference(varname, value):
    '''
    sets value to globals()[varname] and runs gc.collect()
    '''
    globals()[varname] = value
    gc.collect()
    return

def global_delete_from_object(obj):
    '''
    removes object references from globals().
    '''

    names = get_pointer_names(obj)

    for name in names:
        global_delete_from_reference(name)

def global_set_from_object(obj, value):
    '''
    sets value to references from globals().
    '''
    names = get_pointer_names(obj)

    for name in names:
        global_set_from_reference(name, value)

# Cell
class SimpleCacher():

    def __init__(self, obj, serializer, dirpath = './_skdag_cache', override_on_initialization = False, remove_after_gc = True, dump_kwargs = {}, load_kwargs = {}):

        #avoid_caching a cacher and creating nested structures
        while isinstance(obj, self.__class__):
            obj = obj.load()
        #
        self.obj = None
        self.serializer = serializer
        self.dirpath = Path(dirpath)
        self.override_on_initialization = False
        self.remove_after_gc = remove_after_gc
        self.dump_kwargs = dump_kwargs
        self.load_kwargs = load_kwargs
        #extra attrs
        self._name = f'{obj.__class__.__name__}{id(self)}'
        self.path = self.dirpath/f'{self._name}.skdag_cache'
        #dump obj
        self._safe_serialize(obj, override = override_on_initialization)
        return

    def _safe_serialize(self, obj, override = False):
        '''
        checks if file exists prior to serializing
        '''
        if not override:
            if self.path.exists():
                raise FileExistsError(f'A file already exists under {self.path.absolute()}. To override, set override = True.')

        #creates parent folders
        if not self.path.parent.exists():
            os.makedirs(dirpath)

        self.serializer.dump(obj, self.path, **self.dump_kwargs)
        return

    def __del__(self,):
        '''
        deletes serialized objects in dict when self is deleted
        '''
        if hasattr(self, 'remove_after_gc'):
            if self.remove_after_gc:
                if hasattr(self, 'path'):
                    if self.path.exists():
                        self.path.unlink()
        return

    def __repr__(self,):
        name = self._name
        path = self.path
        return f'{name} at <{path.absolute()}>'

    def dump(self, obj, override = False):
        #pickle
        self._safe_serialize(obj, override = override)
        return self.path

    def load(self,):
        #unpickle
        return self.serializer.load(self.path, **self.load_kwargs)

    def remove_file(self,):
        '''
        removes file and returns object
        '''
        obj = self.load()
        remove_folder_or_file(self.path)
        return obj


class EstimatorCacher(BaseEstimator):

    def _unnest_cacher(self, obj):

        sklearn_estim = obj

        while isinstance(sklearn_estim, (EstimatorCacher, SimpleCacher)):
            if isinstance(sklearn_estim, EstimatorCacher):
                sklearn_estim = sklearn_estim.estimator.load()
            else:
                sklearn_estim = sklearn_estim.load()

        return sklearn_estim

    def __init__(self, cached_estimator, serializer = joblib, dirpath = './_skdag_cache', override_on_initialization = False, remove_after_gc = True, load_kwargs = {}, dump_kwargs = {}):

        #avoid caching a Cacher
        self.cached_estimator = self._unnest_cacher(cached_estimator)
        self.serializer = serializer
        self.dirpath = dirpath
        self.override_on_initialization = override_on_initialization
        self.remove_after_gc = remove_after_gc
        self.load_kwargs = load_kwargs
        self.dump_kwargs = dump_kwargs
        #extra
        self.__initialized = True
        self.estimator_ = None
        return

    def fit(self, X, y = None, **kwargs):

        estim = clone(self._unnest_cacher(self.cached_estimator))
        self.estimator_ = SimpleCacher(estim, self.serializer, self.dirpath, self.override_on_initialization, self.remove_after_gc, self.load_kwargs, self.dump_kwargs)
        estim = self.estimator_.load()
        estim.fit(X, y, **kwargs)
        self.estimator_.dump(estim, override = True)
        return self

    def __getattr__(self, attr):

        try:
            #checks if is instaitated
            super().__getattribute__(f'_{self.__class__.__name__}__initialized')
            #deserializes
            obj = super().__getattribute__('estimator_').load()
            return self._safe_access(obj, attr)

        #case where object is not instatiated yet
        except AttributeError:
            return super().__getattribute__(attr)


    def _safe_run(self, obj, method, *args, **kwargs):
        '''
        dumps updated object after running a method or accessing a property
        '''
        result = getattr(obj, method)(*args, **kwargs)
        self.estimator_.dump(obj, override = True)
        return result


    def _safe_access(self, obj, attr):

        #handles cases where accessing attribute might change object state
        if callable(getattr(obj.__class__, attr)):
            return partial(self._safe_run, obj, attr)

        if isinstance(getattr(obj.__class__, attr), property):
            return partial(self._safe_run, obj, attr)()
        else:
            return getattr(obj,attr)