"""
Cache
"""
import os
import pickle
import yaml
from .helpers import actual_kwargs, mkdir, rmd


def _itemize(data):
    """Itemize numpy arrays in data"""
    if hasattr(data, 'shape') and len(data.shape) > 1:
        return data
    try:
        new_data = {}
        for key in data:
            if hasattr(data[key], 'item') and data[key].size == 1:
                new_data[key] = data[key].item()
            else:
                new_data[key] = data[key]
    except TypeError:
        # Also do it if data itself is a numpy array
        if hasattr(data, 'item') and data.size == 1:
            new_data = data.item()
        else:
            new_data = data
    return new_data

class Cache:
    """
    Cache results of function evaluation on the basis of its name and
    keyword arguments
    """

    def __init__(self, path):
        """
        :param path: path to cache folder
        """
        self.path = path

    def _storage(self, name):
        return os.path.join(self.path, name)

    def setup(self, name, **kwargs):
        """Setup the cache for `name` with `kwargs`"""
        path = self._storage(name)
        mkdir(path)
        with open(os.path.join(path, 'arguments.pkl'), 'wb') as fh:
            pickle.dump(kwargs, fh)
        # For quick inspection, we also provide a yaml file
        # which will be ignored by default by parsers
        with open(os.path.join(path, '.arguments.yaml'), 'w') as fh:
            yaml.dump(_itemize(kwargs), fh)

    def write(self, name, results):
        """Write the function cache results for `name`"""
        path = self._storage(name)
        with open(os.path.join(path, 'results.pkl'), 'wb') as fh:
            pickle.dump(results, fh)
        with open(os.path.join(path, '.results.yaml'), 'w') as fh:
            # Itemize all numpy arrays
            yaml.dump(_itemize(results), fh)

    def read(self, name):
        """Return the function cache results"""
        path = self._storage(name)
        with open(os.path.join(path, 'results.pkl'), 'rb') as fh:
            results = pickle.load(fh)
        return results

    def found(self, name):
        """Return True is `name` if found in cache"""
        path = self._storage(name)
        return os.path.exists(os.path.join(path, 'results.pkl'))

    def clear(self, name):
        """Clear cache for `name`"""
        rmd(self._storage(name))


default_cache = Cache('.pantarei')
