'''
Utilities for dealing with simple state dicts as npz files instead of pth files.
'''

import torch
from collections.abc import MutableMapping, Mapping

def load_from_numpy_dict(model, numpy_dict, prefix='', examples=None):
    '''
    Loads a model from numpy_dict using load_state_dict.
    Converts numpy types to torch types using the current state_dict
    of the model to determine types and devices for the tensors.
    Supports loading a subdict by prepending the given prefix to all keys.
    '''
    if prefix:
        if not prefix.endswith('.'):
            prefix = prefix + '.'
        numpy_dict = PrefixSubDict(numpy_dict, prefix)
    if examples is None:
        exampels = model.state_dict()
    torch_state_dict = TorchTypeMatchingDict(numpy_dict, examples)
    model.load_state_dict(torch_state_dict)

def save_to_numpy_dict(model, numpy_dict, prefix=''):
    '''
    Saves a model by copying tensors to numpy_dict.
    Converts torch types to numpy types using `t.detach().cpu().numpy()`.
    Supports saving a subdict by prepending the given prefix to all keys.
    '''
    if prefix:
        if not prefix.endswith('.'):
            prefix = prefix + '.'
    for k, v in model.numpy_dict().items():
        if isinstance(v, torch.Tensor):
            v = v.detach().cpu().numpy()
        numpy_dict[prefix + k] = v

class TorchTypeMatchingDict(Mapping):
    '''
    Provides a view of a dict of numpy values as torch tensors, where the
    types are converted to match the types and devices in the given
    dict of examples.
    '''
    def __init__(self, data, examples):
        self.data = data
        self.examples = examples
        self.cached_data = {}
    def __getitem__(self, key):
        if key in self.cached_data:
            return self.cached_data[key]
        val = self.data[key]
        if key not in self.examples:
            return val
        example = self.examples.get(key, None)
        example_type = type(example)
        if example is not None and type(val) != example_type:
            if isinstance(example, torch.Tensor):
                val = torch.from_numpy(val)
            else:
                val = example_type(val)
        if isinstance(example, torch.Tensor):
            val = val.to(dtype=example.dtype, device=example.device)
        self.cached_data[key] = val
        return val
    def __iter__(self):
        return self.data.keys()
    def __len__(self):
        return len(self.data)

class PrefixSubDict(MutableMapping):
    '''
    Provides a view of the subset of a dict where string keys begin with
    the given prefix.  The prefix is stripped from all keys of the view.
    '''
    def __init__(self, data, prefix=''):
        self.data = data
        self.prefix = prefix
        self._cached_keys = None
    def __getitem__(self, key):
        return self.data[self.prefix + key]
    def __setitem__(self, key, value):
        pkey = self.prefix + key
        if self._cached_keys is not None and pkey not in self.data:
            self._cached_keys = None
        self.data[pkey] = value
    def __delitem__(self, key):
        pkey = self.prefix + key
        if self._cached_keys is not None and pkey in self.data:
            self._cached_keys = None
        del self.data[pkey]
    def __cached_keys(self):
        if self._cached_keys is None:
            plen = len(self.prefix)
            self._cached_keys = list(k[plen:] for k in self.data
                    if k.startswith(self.prefix))
        return self._cached_keys
    def __iter__(self):
        return iter(self.__cached_keys())
    def __len__(self):
        return len(self.__cached_keys())