import numpy as np


__all__ = ['JustOnceClass', 'just_once', 'Cache']


class JustOnceClass(object):
    '''Base class for classes with methods that should never be executed twice.

       In typically applications, these methods get called many times, but
       only during the first call, an actual computation is carried out. This
       way, the caller can safely call a method, just to make sure that a
       required result is computed.

       All methods in the subclasses that should have this feature, must be
       given the ``just_once`` decoratore, e.g. ::

           class Example(JustOnceClass):
               @just_once
               def do_something():
                   self.foo = self.bar

       When all results are outdated, one can call the ``clear`` method
       to forget which methods were called already.
    '''
    def __init__(self):
        self._done_just_once = set([])

    def __clear__(self):
        self.clear()

    def clear(self):
        self._done_just_once = set([])


def just_once(fn):
    def wrapper(instance):
        if not hasattr(instance, '_done_just_once'):
            raise TypeError('Missing hidden _done_just_once. Forgot to call JustOnceClass.__init__()?')
        if fn.__name__ in instance._done_just_once:
            return
        fn(instance)
        instance._done_just_once.add(fn.__name__)
    wrapper.__doc__ = fn.__doc__
    return wrapper


def _normalize_alloc(alloc):
    '''Normalize the alloc argument of the from_alloc and check_alloc methods'''
    if not hasattr(alloc, '__len__'):
        alloc = (alloc,)
    if len(alloc) == 0:
        raise TypeError('Alloc can not be an empty list')
    return alloc


def _normalize_tags(tags):
    '''Normalize the tags argument of the CacheItem constructor'''
    if tags is None:
        return set([])
    else:
        return set(tags)


class CacheItem(object):
    '''A container for an object stored in a Cache instance'''
    def __init__(self, value, tags=None):
        '''
           **Arguments:**

           value
                The object stored in this container

           **Optional arguments:**

           tags
                Tags to be associated with the object
        '''
        self._value = value
        self._valid = True
        self._tags = _normalize_tags(tags)

    @classmethod
    def from_alloc(cls, alloc, tags):
        alloc = _normalize_alloc(alloc)
        # initialize a floating point array
        array = np.zeros(alloc, float)
        return cls(array, tags=tags)

    def check_alloc(self, alloc):
        alloc = _normalize_alloc(alloc)
        # check if the array has the correct shape and dtype
        if not (isinstance(self._value, np.ndarray) and
                self._value.shape == tuple(alloc) and
                issubclass(self._value.dtype.type, float)):
            raise TypeError('The stored item does not match the given alloc.')

    def check_tags(self, tags):
        tags = _normalize_tags(tags)
        if tags != self._tags:
            raise ValueError('Tags do not match.')

    def _get_value(self):
        if not self._valid:
            raise ValueError('This cached item is not valid.')
        return self._value

    value = property(_get_value)

    def _get_valid(self):
        return self._valid

    valid = property(_get_valid)

    def _get_tags(self):
        return self._tags

    tags = property(_get_tags)

    def clear(self):
        '''Mark the item as invalid and clear the contents of the object.

           **Returns:** A boolean indicating that clearing was successful
        '''
        self._valid = False
        if isinstance(self._value, np.ndarray):
            self._value[:] = 0.0
        elif hasattr(self._value, '__clear__') and callable(self._value.__clear__):
            self._value.__clear__()
        else:
            return False
        return True


class NoDefault(object):
    pass


no_default = NoDefault()


def _normalize_key(key):
    '''Normalize the key argument(s) of the load and dump methods'''
    if hasattr(key, '__len__') and  len(key) == 0:
        raise TypeError('At least one argument needed to specify a key.')
    # upack the key if needed
    while len(key) == 1 and isinstance(key, tuple):
        key = key[0]
    return key


class Cache(object):
    '''Object that stores previously computed results.

       The cache behaves like a dictionary with some extra features that can be
       used to avoid recomputation or reallocation.
    '''
    def __init__(self):
        self._store = {}

    def clear(self, **kwargs):
        '''Clear all items in the cache

           **Optional arguments:**

           dealloc
                When set to True, the items are really removed from memory.

           tags
                Limit the items cleared to those who have at least one tag
                that matches one of the given tags. When this argument is used
                and it contains at least one tag, items with no tags are not
                cleared.
        '''
        # Parse kwargs. This forces the caller to use keywords in order to avoid
        # confusion.
        dealloc = kwargs.pop('dealloc', False)
        tags = kwargs.pop('tags', None)
        if len(kwargs) > 0:
            raise TypeError('Unexpected arguments: %s' % list(kwargs.keys()))
        # actual work
        tags = _normalize_tags(tags)
        for key, item in list(self._store.items()):
            if len(tags) == 0 or len(item.tags & tags) > 0:
                self.clear_item(key, dealloc=dealloc)

    def clear_item(self, *key, **kwargs):
        '''Clear a selected item from the cache

           **Optional arguments:**

           dealloc
                When set to True, the item is really removed from memory.
        '''
        key = _normalize_key(key)
        dealloc = kwargs.pop('dealloc', False)
        if len(kwargs) > 0:
            raise TypeError('Unexpected arguments: %s' % list(kwargs.keys()))
        item = self._store.get(key)
        if item is None:
            return
        cleared = False
        if not dealloc:
            cleared = item.clear()
        if not cleared:
            del self._store[key]

    def load(self, *key, **kwargs):
        '''Get a value from the cache

           **Arguments:**

           key0 [key1 ...]
                All positional arguments are used as keys to identify the cached
                value.

           **Optional arguments:**

           alloc
                Parameters used to allocate a cached value if it is not present
                yet. This argument can take two forms. When integer or a
                tuple of integers is given, an array is allocated.
                Alternatively, a tuple may be given whose first element is a
                constructor, and further elements are arguments for that
                constructor.

           default
                A default value that is returned when the key does not exist in
                the cache. This default value is not stored in the cache.

           tags
                When alloc is used and a new object is thereby created or
                reused, it will get these tags. This argument is only allowed if
                the alloc argument is present. In case no new object is
                allocated, the given tags must match those already present.

           The optional argument alloc and default are both meant to handle
           situations when the key has not associated value. Hence they can not
           be both present.
        '''
        key = _normalize_key(key)

        # parse kwargs
        alloc = kwargs.pop('alloc', None)
        default = kwargs.pop('default', no_default)
        tags = kwargs.pop('tags', None)
        if not (alloc is None or default is no_default):
            raise TypeError('The optional arguments alloc and default can not be used at the same time.')
        if tags is not None and alloc is None:
            raise TypeError('The tags argument is only allowed when the alloc argument is present.')
        if len(kwargs) > 0:
            raise TypeError('Unknown optional arguments: %s' % list(kwargs.keys()))

        # get the item from the store and decide what to do
        item = self._store.get(key)
        # there are three behaviors, depending on the keyword argumentsL
        if alloc is not None:
            # alloc is given. hence two return values: value, new
            if item is None:
                # allocate a new item and store it
                item = CacheItem.from_alloc(alloc, tags)
                self._store[key] = item
                return item.value, True
            elif not item.valid:
                try:
                    # try to reuse the same memroy
                    item.check_alloc(alloc)
                    item._valid = True # as if it is newly allocated
                    item.check_tags(tags)
                except TypeError:
                    # if reuse fails, reallocate
                    item = CacheItem.from_alloc(alloc, tags)
                    self._store[key] = item
                return item.value, True
            else:
                item.check_alloc(alloc)
                item.check_tags(tags)
                return item.value, False
        elif default is not no_default:
            # a default value is given, it is not stored
            if item is None or not item.valid:
                return default
            else:
                return item.value
        else:
            # no optional arguments are given
            if item is None or not item.valid:
                raise KeyError(key)
            else:
                return item.value

    def __contains__(self, key):
        key = _normalize_key(key)
        item = self._store.get(key)
        if item is None:
            return False
        else:
            return item.valid

    def dump(self, *args, **kwargs):
        '''Store an object in the cache.

           **Arguments:**

           key1 [, key2, ...]
                The positional arguments (except for the last) are used as a key
                for the object.

           value
                The object to be stored.

           **Optional argument:**

           tags
                Tags to be associated with the object
        '''
        tags = kwargs.pop('tags', None)
        if len(kwargs) > 0:
            raise TypeError('Unknown optional arguments: %s' % list(kwargs.keys()))
        if len(args) < 2:
            raise TypeError('At least two arguments are required: key1 and value.')
        key = _normalize_key(args[:-1])
        value = args[-1]
        item = CacheItem(value, tags)
        self._store[key] = item

    def __len__(self):
        return sum(item.valid for item in self._store.values())

    def __getitem__(self, key):
        return self.load(key)

    def __setitem__(self, key, value):
        return self.dump(key, value)

    def __iter__(self):
        return iter(self.keys())

    def keys(self, tags=None):
        '''Iterate over the keys of all valid items in the cache.'''
        tags = _normalize_tags(tags)
        for key, item in self._store.items():
            if item.valid and (len(tags) == 0 or len(item.tags & tags) > 0):
                yield key

    def values(self, tags=None):
        '''Iterate over the values of all valid items in the cache.'''
        tags = _normalize_tags(tags)
        for item in self._store.values():
            if item.valid and (len(tags) == 0 or len(item.tags & tags) > 0):
                yield item.value

    def items(self, tags=None):
        '''Iterate over all valid items in the cache.'''
        tags = _normalize_tags(tags)
        for key, item in self._store.items():
            if item.valid and (len(tags) == 0 or len(item.tags & tags) > 0):
                yield key, item.value
