import copy
import warnings
import weakref
from collections import namedtuple
from collections.abc import Mapping, MutableMapping, Sequence
from typing import (Any, Callable, Dict, Iterable, List, NamedTuple, Optional,
					Tuple, Union)

import numpy as np
import tensorlayerx as tlx
from gammagl.data.view import ItemsView, KeysView, ValuesView
from gammagl.utils import is_undirected, coalesce

# Node-types are denoted by a single string, e.g.: `data['paper']`:
NodeType = str

# Edge-types are denotes by a triplet of strings, e.g.:
# `data[('author', 'writes', 'paper')]
EdgeType = Tuple[str, str, str]


class BaseStorage(MutableMapping):
	# Copy from PyG
	# This class wraps a Python dictionary and extends it as follows:
	# 1. It allows attribute assignments, e.g.:
	#    `storage.x = ...` in addition to `storage['x'] = ...`
	# 2. It allows private attributes that are not exposed to the user, e.g.:
	#    `storage._{key} = ...` and accessible via `storage._{key}`
	# 3. It holds an (optional) weak reference to its parent object, e.g.:
	#    `storage._parent = weakref.ref(parent)`
	# 4. It allows iterating over only a subset of keys, e.g.:
	#    `storage.values('x', 'y')` or `storage.items('x', 'y')
	# 5. (Only in PyG) It adds additional PyTorch Tensor functionality, e.g.:
	#    `storage.cpu()`, `storage.cuda()` or `storage.share_memory_()`.
	def __init__(self, _mapping: Optional[Dict[str, Any]] = None, **kwargs):
		super().__init__()
		self._mapping = {}
		for key, value in (_mapping or {}).items():
			setattr(self, key, value)
		for key, value in kwargs.items():
			setattr(self, key, value)
	
	@property
	def _key(self) -> Any:
		return None
	
	def __len__(self) -> int:
		return len(self._mapping)
	
	def __getattr__(self, key: str) -> Any:
		try:
			return self[key]
		except KeyError:
			raise AttributeError(
				f"'{self.__class__.__name__}' object has no attribute '{key}'")
	
	def __setattr__(self, key: str, value: Any):
		if key == '_parent':
			self.__dict__[key] = weakref.ref(value)
		elif key[:1] == '_':
			self.__dict__[key] = value
		else:
			self[key] = value
	
	def __delattr__(self, key: str):
		if key[:1] == '_':
			del self.__dict__[key]
		else:
			del self[key]
	
	def __getitem__(self, key: str) -> Any:
		return self._mapping[key]
	
	def __setitem__(self, key: str, value: Any):
		if value is None and key in self._mapping:
			del self._mapping[key]
		elif value is not None:
			self._mapping[key] = value
	
	def __delitem__(self, key: str):
		if key in self._mapping:
			del self._mapping[key]
	
	def __iter__(self) -> Iterable:
		return iter(self._mapping)
	
	def __copy__(self):
		out = self.__class__.__new__(self.__class__)
		for key, value in self.__dict__.items():
			out.__dict__[key] = value
		out._mapping = copy.copy(out._mapping)
		return out
	
	def __deepcopy__(self, memo):
		out = self.__class__.__new__(self.__class__)
		for key, value in self.__dict__.items():
			out.__dict__[key] = value
		out._mapping = copy.deepcopy(out._mapping, memo)
		return out
	
	def __getstate__(self) -> Dict[str, Any]:
		out = self.__dict__.copy()
		
		_parent = out.get('_parent', None)
		if _parent is not None:
			out['_parent'] = _parent()
		
		return out
	
	def __setstate__(self, mapping: Dict[str, Any]):
		for key, value in mapping.items():
			self.__dict__[key] = value
		
		_parent = self.__dict__.get('_parent', None)
		if _parent is not None:
			self.__dict__['_parent'] = weakref.ref(_parent)
	
	def __repr__(self) -> str:
		return repr(self._mapping)
	
	# Allow iterating over subsets ############################################
	
	# In contrast to standard `keys()`, `values()` and `items()` functions of
	# Python dictionaries, we allow to only iterate over a subset of items
	# denoted by a list of keys `args`.
	# This is especially useful for adding PyTorch Tensor functionality to the
	# storage object, e.g., in case we only want to transfer a subset of keys
	# to the GPU (i.e. the ones that are relevant to the deep learning model).
	
	def keys(self, *args: List[str]) -> KeysView:
		return KeysView(self._mapping, *args)
	
	def values(self, *args: List[str]) -> ValuesView:
		return ValuesView(self._mapping, *args)
	
	def items(self, *args: List[str]) -> ItemsView:
		return ItemsView(self._mapping, *args)
	
	def apply_(self, func: Callable, *args: List[str]):
		r"""Applies the in-place function :obj:`func`, either to all attributes
		or only the ones given in :obj:`*args`."""
		for value in self.values(*args):
			recursive_apply_(value, func)
		return self

	def apply(self, func: Callable, *args: List[str]):
		r"""Applies the function :obj:`func`, either to all attributes or only
		the ones given in :obj:`*args`."""
		for key, value in self.items(*args):
			self[key] = recursive_apply(value, func)
		return self
	
	# Additional functionality ################################################
	
	def to_dict(self) -> Dict[str, Any]:
		r"""Returns a dictionary of stored key/value pairs."""
		return copy.copy(self._mapping)
	
	def to_namedtuple(self) -> NamedTuple:
		r"""Returns a :obj:`NamedTuple` of stored key/value pairs."""
		field_names = list(self.keys())
		typename = f'{self.__class__.__name__}Tuple'
		StorageTuple = namedtuple(typename, field_names)
		return StorageTuple(*[self[key] for key in field_names])

	def clone(self, *args: List[str]):
		r"""Performs a deep-copy of the object."""
		return copy.deepcopy(self)
	
	# def contiguous(self, *args: List[str]):
	#     r"""Ensures a contiguous memory layout, either for all attributes or
	#     only the ones given in :obj:`*args`."""
	#     return self.apply(lambda x: x.contiguous(), *args)
	#
	# def to(self, device: Union[int, str], *args: List[str],
	#        non_blocking: bool = False):
	#     r"""Performs tensor dtype and/or device conversion, either for all
	#     attributes or only the ones given in :obj:`*args`."""
	#     return self.apply(
	#         lambda x: x.to(device=device, non_blocking=non_blocking), *args)
	#
	# def cpu(self, *args: List[str]):
	#     r"""Copies attributes to CPU memory, either for all attributes or only
	#     the ones given in :obj:`*args`."""
	#     return self.apply(lambda x: x.cpu(), *args)
	#
	# def cuda(self, device: Optional[Union[int, str]] = None, *args: List[str],
	#          non_blocking: bool = False):
	#     r"""Copies attributes to CUDA memory, either for all attributes or only
	#     the ones given in :obj:`*args`."""
	#     return self.apply(lambda x: x.cuda(device, non_blocking=non_blocking),
	#                       *args)
	#
	# def pin_memory(self, *args: List[str]):
	#     r"""Copies attributes to pinned memory, either for all attributes or
	#     only the ones given in :obj:`*args`."""
	#     return self.apply(lambda x: x.pin_memory(), *args)
	#
	# def share_memory_(self, *args: List[str]):
	#     r"""Moves attributes to shared memory, either for all attributes or
	#     only the ones given in :obj:`*args`."""
	#     return self.apply(lambda x: x.share_memory_(), *args)
	#
	# def detach_(self, *args: List[str]):
	#     r"""Detaches attributes from the computation graph, either for all
	#     attributes or only the ones given in :obj:`*args`."""
	#     return self.apply(lambda x: x.detach_(), *args)
	#
	# def detach(self, *args: List[str]):
	#     r"""Detaches attributes from the computation graph by creating a new
	#     tensor, either for all attributes or only the ones given in
	#     :obj:`*args`."""
	#     return self.apply(lambda x: x.detach(), *args)
	#
	# def requires_grad_(self, *args: List[str], requires_grad: bool = True):
	#     r"""Tracks gradient computation, either for all attributes or only the
	#     ones given in :obj:`*args`."""
	#     return self.apply(
	#         lambda x: x.requires_grad_(requires_grad=requires_grad), *args)
	
	# def record_stream(self, stream: torch.cuda.Stream, *args: List[str]):
	#     r"""Ensures that the tensor memory is not reused for another tensor
	#     until all current work queued on :obj:`stream` has been completed,
	#     either for all attributes or only the ones given in :obj:`*args`."""
	#     return self.apply_(lambda x: x.record_stream(stream), *args)


class NodeStorage(BaseStorage):
	@property
	def _key(self) -> NodeType:
		key = self.__dict__.get('_key', None)
		if key is None or not isinstance(key, str):
			raise ValueError("'_key' does not denote a valid node type")
		return key
	
	@property
	def can_infer_num_nodes(self):
		keys = set(self.keys())
		num_node_keys = {
			'num_nodes', 'x', 'pos', 'batch', 'adj', 'adj_t', 'edge_index',
			'face'
		}
		if len(keys & num_node_keys) > 0:
			return True
		elif len([key for key in keys if 'node' in key]) > 0:
			return True
		else:
			return False
	
	@property
	def num_nodes(self) -> Optional[int]:
		# We sequentially access attributes that reveal the number of nodes.
		if 'num_nodes' in self:
			return self['num_nodes']
		for key, value in self.items():
			if (tlx.is_tensor(value)
					and (key in {'x', 'pos', 'batch'} or 'node' in key)):
				return value.shape[self._parent().__cat_dim__(key, value, self)]
		# if 'adj' in self and isinstance(self.adj, SparseTensor):
		#     return self.adj.size(0)
		# if 'adj_t' in self and isinstance(self.adj_t, SparseTensor):
		#     return self.adj_t.size(1)
		# warnings.warn(
		#     f"Unable to accurately infer 'num_nodes' from the attribute set "
		#     f"'{set(self.keys())}'. Please explicitly set 'num_nodes' as an "
		#     f"attribute of " +
		#     ("'data'" if self._key is None else f"'data[{self._key}]'") +
		#     " to suppress this warning")
		if 'edge_index' in self and tlx.is_tensor(self.edge_index):
			if self.edge_index.shape[-1] > 0:
				return int(tlx.reduce_max(self.edge_index)) + 1
			else:
				return 0
		return None
	
	@property
	def num_node_features(self) -> int:
		if 'x' in self:
			return 1 if self.x.ndim == 1 else self.x.shape[-1]
		return 0
	
	@property
	def num_features(self) -> int:
		return self.num_node_features

	def is_node_attr(self, key: str) -> bool:
		value = self[key]
		cat_dim = self._parent().__cat_dim__(key, value, self)
		if not tlx.is_tensor(value):
			return False
		if tlx.get_tensor_shape(value)[cat_dim] != self.num_nodes:
			return False
		return True
	
	def is_edge_attr(self, key: str) -> bool:
		return False


class EdgeStorage(BaseStorage):
	r""" Copied from PyG
	We support multiple ways to store edge connectivity in a
	:class:`EdgeStorage` object:
	* :obj:`edge_index`: A :class:`torch.LongTensor` holding edge indices in
	  COO format with shape :obj:`[2, num_edges]` (the default format)
	* :obj:`adj`: A :class:`torch_sparse.SparseTensor` holding edge indices in
	  a sparse format, supporting both COO and CSR format.
	* :obj:`adj_t`: A **transposed** :class:`torch_sparse.SparseTensor` holding
	  edge indices in a sparse format, supporting both COO and CSR format.
	  This is the most efficient one for graph-based deep learning models as
	  indices are sorted based on target nodes.
	"""
	@property
	def _key(self) -> EdgeType:
		key = self.__dict__.get('_key', None)
		if key is None or not isinstance(key, tuple) or not len(key) == 3:
			raise ValueError("'_key' does not denote a valid edge type")
		return key
	
	@property
	def edge_index(self):
		if 'edge_index' in self:
			return self['edge_index']
		# if 'adj' in self and isinstance(self.adj, SparseTensor):
		#     return torch.stack(self.adj.coo()[:2], dim=0)
		# if 'adj_t' in self and isinstance(self.adj_t, SparseTensor):
		#     return torch.stack(self.adj_t.coo()[:2][::-1], dim=0)
		# raise AttributeError(
		#     f"'{self.__class__.__name__}' object has no attribute "
		#     f"'edge_index', 'adj' or 'adj_t'")
	
	@property
	def num_edges(self) -> int:
		# We sequentially access attributes that reveal the number of edges.
		for key, value in self.items():
			if tlx.is_tensor(value) and 'edge' in key:
				return tlx.get_tensor_shape(value)[self._parent().__cat_dim__(key, value, self)]
		# for value in self.values('adj', 'adj_t'):
		#     if isinstance(value, SparseTensor):
		#         return value.nnz()
		return 0
	
	@property
	def num_edge_features(self) -> int:
		if 'edge_attr' in self and tlx.is_tensor(self.edge_attr):
			return 1 if self.edge_attr.ndim == 1 else self.edge_attr.shape[-1]
		return 0
	
	@property
	def num_features(self) -> int:
		return self.num_edge_features
	
	def size(
			self, dim: Optional[int] = None
	) -> Union[Tuple[Optional[int], Optional[int]], Optional[int]]:
		
		if self._key is None:
			raise NameError("Unable to infer 'size' without explicit "
							"'_key' assignment")
		
		size = (self._parent()[self._key[0]].num_nodes,
				self._parent()[self._key[-1]].num_nodes)
		
		return size if dim is None else size[dim]
	
	def is_node_attr(self, key: str) -> bool:
		return False
	
	def is_edge_attr(self, key: str) -> bool:
		value = self[key]
		cat_dim = self._parent().__cat_dim__(key, value, self)
		if not tlx.is_tensor(value):
			return False
		if tlx.get_tensor_shape(value)[cat_dim] != self.num_edges:
			return False
		return True

	def is_coalesced(self) -> bool:
		for value in self.values('adj', 'adj_t'):
			return value.is_coalesced()

		edge_index = self.edge_index
		new_edge_index = coalesce(edge_index)
		return tlx.convert_to_numpy(edge_index).size == tlx.convert_to_numpy(new_edge_index).size

	def coalesce(self, reduce: str = 'sum'):
		for key, value in self.items('adj', 'adj_t'):
			self[key] = value.coalesce(reduce)

		if 'edge_index' in self:
			edge_index = self.edge_index
			if 'edge_attr' in self:
				edge_attr = self.edge_attr
				self.edge_index, self.edge_attr = coalesce(edge_index, edge_attr, reduce=reduce)
			else:
				self.edge_index = coalesce(edge_index, reduce=reduce)

		return self
	
	def has_isolated_nodes(self) -> bool:
		edge_index, num_nodes = self.edge_index, self.size(1)
		if num_nodes is None:
			raise NameError("Unable to infer 'num_nodes'")
		if tlx.is_tensor(edge_index):
			return np.unique(tlx.convert_to_numpy(edge_index[1])).shape[0] < num_nodes
		else:
			return np.unique(edge_index[1]).shape[0] < num_nodes
	
	def has_self_loops(self) -> bool:
		if self.is_bipartite():
			return False
		edge_index = self.edge_index
		if tlx.is_tensor(edge_index):
			return int(tlx.convert_to_numpy((edge_index[0] == edge_index[1])).sum()) > 0
		else:
			return int((edge_index[0] == edge_index[1]).sum()) > 0
	
	def is_undirected(self) -> bool:
		if self.is_bipartite():  # TODO check for inverse storage.
			return False
		# for value in self.values('adj', 'adj_t'):
		# 	return value.is_symmetric()

		edge_index = self.edge_index
		edge_attr = self.edge_attr if 'edge_attr' in self else None
		return is_undirected(edge_index, edge_attr, num_nodes=self.size(0))

	def is_directed(self) -> bool:
		return not self.is_undirected()
	
	def is_bipartite(self) -> bool:
		return self._key is not None and self._key[0] != self._key[-1]


class GlobalStorage(NodeStorage, EdgeStorage):
	@property
	def _key(self) -> Any:
		return None
	
	@property
	def num_features(self) -> int:
		return self.num_node_features
	
	def size(
			self, dim: Optional[int] = None
	) -> Union[Tuple[Optional[int], Optional[int]], Optional[int]]:
		size = (self.num_nodes, self.num_nodes)
		return size if dim is None else size[dim]

	def is_node_attr(self, key: str) -> bool:
		value = self[key]
		cat_dim = self._parent().__cat_dim__(key, value, self)

		num_nodes, num_edges = self.num_nodes, self.num_edges
		if not tlx.is_tensor(value):
			return False
		if tlx.get_tensor_shape(value)[cat_dim] != num_nodes:
			return False
		if num_nodes != num_edges:
			return True
		return 'edge' not in key

	def is_edge_attr(self, key: str) -> bool:
		value = self[key]
		cat_dim = self._parent().__cat_dim__(key, value, self)

		num_edges = self.num_edges
		if not tlx.is_tensor(value):
			return False
		if tlx.get_tensor_shape(value)[cat_dim] != num_edges:
			return False
		return 'edge' in key


def recursive_apply_(data: Any, func: Callable):
	if tlx.is_tensor(data):
		func(data)
	elif isinstance(data, tuple) and hasattr(data, '_fields'):  # namedtuple
		for value in data:
			recursive_apply_(value, func)
	elif isinstance(data, Sequence) and not isinstance(data, str):
		for value in data:
			recursive_apply_(value, func)
	elif isinstance(data, Mapping):
		for value in data.values():
			recursive_apply_(value, func)
	else:
		try:
			func(data)
		except:  # noqa
			pass


def recursive_apply(data: Any, func: Callable) -> Any:
	if tlx.is_tensor(data):
		return func(data)
	# elif isinstance(data, torch.nn.utils.rnn.PackedSequence):
	# 	return func(data)
	elif isinstance(data, tuple) and hasattr(data, '_fields'):  # namedtuple
		return type(data)(*(recursive_apply(d, func) for d in data))
	elif isinstance(data, Sequence) and not isinstance(data, str):
		return [recursive_apply(d, func) for d in data]
	elif isinstance(data, Mapping):
		return {key: recursive_apply(data[key], func) for key in data}
	else:
		try:
			return func(data)
		except:  # noqa
			return data