from __future__ import absolute_import, division, print_function

from .callbacks import Callback
from timeit import default_timer
from numbers import Number
import sys

try:
    import cachey
except ImportError:
    pass

overhead = sys.getsizeof(1.23) * 4 + sys.getsizeof(()) * 4


class Cache(Callback):
    """ Use cache for computation

    Examples
    --------

    >>> cache = Cache(1e9)  # available bytes

    >>> with cache:         # use as a context manager around get/compute calls
    ...     result = x.compute()  # doctest: +SKIP

    >>> cache.register()    # or use globally
    >>> cache.unregister()
    """

    def __init__(self, cache, *args, **kwargs):
        if isinstance(cache, Number):
            cache = cachey.Cache(cache, *args, **kwargs)
        else:
            assert not args and not kwargs
        self.cache = cache
        self.starttimes = dict()

    def _start(self, dsk):
        self.durations = dict()
        overlap = set(dsk) & set(self.cache.data)
        for key in overlap:
            dsk[key] = self.cache.data[key]

    def _pretask(self, key, dsk, state):
        self.starttimes[key] = default_timer()

    def _posttask(self, key, value, dsk, state, id):
        duration = default_timer() - self.starttimes[key]
        deps = state['dependencies'][key]
        if deps:
            duration += max(self.durations.get(k, 0) for k in deps)
        self.durations[key] = duration
        nb = cachey.nbytes(value) + overhead + sys.getsizeof(key) * 4
        self.cache.put(key, value, cost=duration / nb / 1e9, nbytes=nb)

    def _finish(self, dsk, state, errored):
        self.starttimes.clear()
        self.durations.clear()

@@ -167,61 +167,8 @@
 ys%0A%0A
-try:%0A    import cachey%0Aexcept ImportError:%0A    pass%0A%0A
 over
@@ -341,60 +341,67 @@
 9)  
-# available bytes%0A%0A    %3E%3E%3E with cache:         # use
+        # doctest: +SKIP%0A%0A    The cache can be used locally
  as 
@@ -429,76 +429,216 @@
 und 
-get/
+%60%60
 compute
- calls%0A    ...     result = x.compute()  # doctest: +SKIP
+%60%60 or%0A    %60%60get%60%60 calls:%0A%0A    %3E%3E%3E with cache:                 # doctest: +SKIP%0A    ...     result = x.compute()%0A%0A    You can also register a cache globally, so that it works for all%0A    computations:
 %0A%0A  
@@ -667,25 +667,32 @@
     
-# or use globally
+        # doctest: +SKIP
 %0A   
@@ -714,16 +714,42 @@
 gister()
+          # doctest: +SKIP
 %0A    %22%22%22
@@ -798,16 +798,259 @@
 wargs):%0A
+        try:%0A            import cachey%0A        except ImportError as ex:%0A            raise ImportError('Cache requires cachey, %22%7Bex%7D%22 problem '%0A                              'importing'.format(ex=str(ex)))%0A        self._nbytes = cachey.nbytes%0A
         
@@ -1815,23 +1815,22 @@
    nb = 
-cachey.
+self._
 nbytes(v

Move cachey import to local import