"""Views of DGLGraph."""
from __future__ import absolute_import

from collections import namedtuple
from collections.abc import MutableMapping

import numpy as np

from .base import ALL, is_all, DGLError
from . import backend as F

NodeSpace = namedtuple('NodeSpace', ['data'])
EdgeSpace = namedtuple('EdgeSpace', ['data'])

class NodeView(object):
    """A NodeView class to act as G.nodes for a DGLGraph.

    Can be used to get a list of current nodes and get and set node data.

    See Also
    --------
    dgl.DGLGraph.nodes
    """
    __slots__ = ['_graph']

    def __init__(self, graph):
        self._graph = graph

    def __len__(self):
        return self._graph.number_of_nodes()

    def __getitem__(self, nodes):
        if isinstance(nodes, slice):
            # slice
            if not (nodes.start is None and nodes.stop is None
                    and nodes.step is None):
                raise DGLError('Currently only full slice ":" is supported')
            return NodeSpace(data=NodeDataView(self._graph, ALL))
        else:
            return NodeSpace(data=NodeDataView(self._graph, nodes))

    def __call__(self):
        """Return the nodes."""
        return F.copy_to(F.arange(0, len(self)), F.cpu())

class NodeDataView(MutableMapping):
    """The data view class when G.nodes[...].data is called.

    See Also
    --------
    dgl.DGLGraph.nodes
    """
    __slots__ = ['_graph', '_nodes']

    def __init__(self, graph, nodes):
        self._graph = graph
        self._nodes = nodes

    def __getitem__(self, key):
        return self._graph.get_n_repr(self._nodes)[key]

    def __setitem__(self, key, val):
        if isinstance(val, np.ndarray):
            val = F.zerocopy_from_numpy(val)
        self._graph.set_n_repr({key : val}, self._nodes)

    def __delitem__(self, key):
        if not is_all(self._nodes):
            raise DGLError('Delete feature data is not supported on only a subset'
                           ' of nodes. Please use `del G.ndata[key]` instead.')
        self._graph.pop_n_repr(key)

    def __len__(self):
        return len(self._graph._node_frame)

    def __iter__(self):
        return iter(self._graph._node_frame)

    def __repr__(self):
        data = self._graph.get_n_repr(self._nodes)
        return repr({key : data[key] for key in self._graph._node_frame})

class EdgeView(object):
    """A EdgeView class to act as G.edges for a DGLGraph.

    Can be used to get a list of current edges and get and set edge data.

    See Also
    --------
    dgl.DGLGraph.edges
    """
    __slots__ = ['_graph']

    def __init__(self, graph):
        self._graph = graph

    def __len__(self):
        return self._graph.number_of_edges()

    def __getitem__(self, edges):
        if isinstance(edges, slice):
            # slice
            if not (edges.start is None and edges.stop is None
                    and edges.step is None):
                raise DGLError('Currently only full slice ":" is supported')
            return EdgeSpace(data=EdgeDataView(self._graph, ALL))
        else:
            return EdgeSpace(data=EdgeDataView(self._graph, edges))

    def __call__(self, *args, **kwargs):
        """Return all the edges."""
        return self._graph.all_edges(*args, **kwargs)

class EdgeDataView(MutableMapping):
    """The data view class when G.edges[...].data is called.

    See Also
    --------
    dgl.DGLGraph.edges
    """
    __slots__ = ['_graph', '_edges']

    def __init__(self, graph, edges):
        self._graph = graph
        self._edges = edges

    def __getitem__(self, key):
        return self._graph.get_e_repr(self._edges)[key]

    def __setitem__(self, key, val):
        if isinstance(val, np.ndarray):
            val = F.zerocopy_from_numpy(val)
        self._graph.set_e_repr({key : val}, self._edges)

    def __delitem__(self, key):
        if not is_all(self._edges):
            raise DGLError('Delete feature data is not supported on only a subset'
                           ' of nodes. Please use `del G.edata[key]` instead.')
        self._graph.pop_e_repr(key)

    def __len__(self):
        return len(self._graph._edge_frame)

    def __iter__(self):
        return iter(self._graph._edge_frame)

    def __repr__(self):
        data = self._graph.get_e_repr(self._edges)
        return repr({key : data[key] for key in self._graph._edge_frame})

class LayerView(object):
    """A LayerView class to act as nflow.layers for a NodeFlow.

    Can be used to get a list of current nodes and get and set node data.
    """
    __slots__ = ['_graph']

    def __init__(self, graph):
        self._graph = graph

    def __len__(self):
        return self._graph.num_layers()

    def __getitem__(self, layer):
        if not isinstance(layer, int):
            raise DGLError('Currently we only support the view of one layer')
        return NodeSpace(data=LayerDataView(self._graph, layer))

    def __call__(self):
        """Return the nodes."""
        return F.arange(0, len(self))

class LayerDataView(MutableMapping):
    """The data view class when G.layers[...].data is called.
    """
    __slots__ = ['_graph', '_layer']

    def __init__(self, graph, layer):
        self._graph = graph
        self._layer = layer

    def __getitem__(self, key):
        return self._graph._node_frames[self._layer][key]

    def __setitem__(self, key, val):
        self._graph._node_frames[self._layer][key] = val

    def __delitem__(self, key):
        del self._graph._node_frames[self._layer][key]

    def __len__(self):
        return len(self._graph._node_frames[self._layer])

    def __iter__(self):
        return iter(self._graph._node_frames[self._layer])

    def __repr__(self):
        data = self._graph._node_frames[self._layer]
        return repr({key : data[key] for key in data})

class BlockView(object):
    """A BlockView class to act as nflow.blocks for a NodeFlow.

    Can be used to get a list of current edges and get and set edge data.
    """
    __slots__ = ['_graph']

    def __init__(self, graph):
        self._graph = graph

    def __len__(self):
        return self._graph.num_blocks

    def __getitem__(self, flow):
        if not isinstance(flow, int):
            raise DGLError('Currently we only support the view of one flow')
        return EdgeSpace(data=BlockDataView(self._graph, flow))

    def __call__(self, *args, **kwargs):
        """Return all the edges."""
        return self._graph.all_edges(*args, **kwargs)

class BlockDataView(MutableMapping):
    """The data view class when G.blocks[...].data is called.
    """
    __slots__ = ['_graph', '_flow']

    def __init__(self, graph, flow):
        self._graph = graph
        self._flow = flow

    def __getitem__(self, key):
        return self._graph._edge_frames[self._flow][key]

    def __setitem__(self, key, val):
        self._graph._edge_frames[self._flow][key] = val

    def __delitem__(self, key):
        del self._graph._edge_frames[self._flow][key]

    def __len__(self):
        return len(self._graph._edge_frames[self._flow])

    def __iter__(self):
        return iter(self._graph._edge_frames[self._flow])

    def __repr__(self):
        data = self._graph._edge_frames[self._flow]
        return repr({key : data[key] for key in data})

class HeteroNodeView(object):
    """A NodeView class to act as G.nodes for a DGLHeteroGraph."""
    __slots__ = ['_graph', '_typeid_getter']

    def __init__(self, graph, typeid_getter):
        self._graph = graph
        self._typeid_getter = typeid_getter

    def __getitem__(self, key):
        if isinstance(key, slice):
            # slice
            if not (key.start is None and key.stop is None
                    and key.step is None):
                raise DGLError('Currently only full slice ":" is supported')
            nodes = ALL
            ntype = None
        elif isinstance(key, tuple):
            nodes, ntype = key
        elif key is None or isinstance(key, str):
            nodes = ALL
            ntype = key
        else:
            nodes = key
            ntype = None
        ntid = self._typeid_getter(ntype)
        return NodeSpace(data=HeteroNodeDataView(self._graph, ntype, ntid, nodes))

    def __call__(self, ntype=None):
        """Return the nodes."""
        ntid = self._typeid_getter(ntype)
        return F.copy_to(F.arange(0, self._graph._graph.number_of_nodes(ntid),
                                  dtype=self._graph.idtype),
                         self._graph.device)

class HeteroNodeDataView(MutableMapping):
    """The data view class when G.ndata[ntype] is called."""
    __slots__ = ['_graph', '_ntype', '_ntid', '_nodes']

    def __init__(self, graph, ntype, ntid, nodes):
        self._graph = graph
        self._ntype = ntype
        self._ntid = ntid
        self._nodes = nodes

    def __getitem__(self, key):
        if isinstance(self._ntype, list):
            ret = {}
            for (i, ntype) in enumerate(self._ntype):
                value = self._graph._get_n_repr(self._ntid[i], self._nodes).get(key, None)
                if value is not None:
                    ret[ntype] = value
            return ret
        else:
            return self._graph._get_n_repr(self._ntid, self._nodes)[key]

    def __setitem__(self, key, val):
        if isinstance(self._ntype, list):
            assert isinstance(val, dict), \
                'Current HeteroNodeDataView has multiple node types, ' \
                'please passing the node type and the corresponding data through a dict.'

            for (ntype, data) in val.items():
                ntid = self._graph.get_ntype_id(ntype)
                self._graph._set_n_repr(ntid, self._nodes, {key : data})
        else:
            assert isinstance(val, dict) is False, \
                'The HeteroNodeDataView has only one node type. ' \
                'please pass a tensor directly'
            self._graph._set_n_repr(self._ntid, self._nodes, {key : val})

    def __delitem__(self, key):
        if isinstance(self._ntype, list):
            for ntid in self._ntid:
                if self._graph._get_n_repr(ntid, ALL).get(key, None) is None:
                    continue
                self._graph._pop_n_repr(ntid, key)
        else:
            self._graph._pop_n_repr(self._ntid, key)

    def __len__(self):
        assert isinstance(self._ntype, list) is False, \
            'Current HeteroNodeDataView has multiple node types, ' \
            'can not support len().'
        return len(self._graph._node_frames[self._ntid])

    def __iter__(self):
        assert isinstance(self._ntype, list) is False, \
            'Current HeteroNodeDataView has multiple node types, ' \
            'can not be iterated.'
        return iter(self._graph._node_frames[self._ntid])

    def keys(self):
        return self._graph._node_frames[self._ntid].keys()

    def values(self):
        return self._graph._node_frames[self._ntid].values()

    def __repr__(self):
        if isinstance(self._ntype, list):
            ret = {}
            for (i, ntype) in enumerate(self._ntype):
                data = self._graph._get_n_repr(self._ntid[i], self._nodes)
                value = {key : data[key]
                         for key in self._graph._node_frames[self._ntid[i]]}
                ret[ntype] = value
            return repr(ret)
        else:
            data = self._graph._get_n_repr(self._ntid, self._nodes)
            return repr({key : data[key]
                         for key in self._graph._node_frames[self._ntid]})

class HeteroEdgeView(object):
    """A EdgeView class to act as G.edges for a DGLHeteroGraph."""
    __slots__ = ['_graph']

    def __init__(self, graph):
        self._graph = graph

    def __getitem__(self, key):
        if isinstance(key, slice):
            # slice
            if not (key.start is None and key.stop is None
                    and key.step is None):
                raise DGLError('Currently only full slice ":" is supported')
            edges = ALL
            etype = None
        elif key is None:
            edges = ALL
            etype = None
        elif isinstance(key, tuple):
            if len(key) == 3:
                edges = ALL
                etype = key
            else:
                edges = key
                etype = None
        elif isinstance(key, str):
            edges = ALL
            etype = key
        else:
            edges = key
            etype = None
        return EdgeSpace(data=HeteroEdgeDataView(self._graph, etype, edges))

    def __call__(self, *args, **kwargs):
        """Return all the edges."""
        return self._graph.all_edges(*args, **kwargs)

class HeteroEdgeDataView(MutableMapping):
    """The data view class when G.edata[etype] is called."""
    __slots__ = ['_graph', '_etype', '_etid', '_edges']

    def __init__(self, graph, etype, edges):
        self._graph = graph
        self._etype = etype
        self._etid = [self._graph.get_etype_id(t) for t in etype] \
                     if isinstance(etype, list) \
                     else self._graph.get_etype_id(etype)
        self._edges = edges

    def __getitem__(self, key):
        if isinstance(self._etype, list):
            ret = {}
            for (i, etype) in enumerate(self._etype):
                value = self._graph._get_e_repr(self._etid[i], self._edges).get(key, None)
                if value is not None:
                    ret[etype] = value
            return ret
        else:
            return self._graph._get_e_repr(self._etid, self._edges)[key]

    def __setitem__(self, key, val):
        if isinstance(self._etype, list):
            assert isinstance(val, dict), \
                'Current HeteroEdgeDataView has multiple edge types, ' \
                'please pass the edge type and the corresponding data through a dict.'

            for (etype, data) in val.items():
                etid = self._graph.get_etype_id(etype)
                self._graph._set_e_repr(etid, self._edges, {key : data})
        else:
            assert isinstance(val, dict) is False, \
                'The HeteroEdgeDataView has only one edge type. ' \
                'please pass a tensor directly'
            self._graph._set_e_repr(self._etid, self._edges, {key : val})

    def __delitem__(self, key):
        if isinstance(self._etype, list):
            for etid in self._etid:
                if self._graph._get_e_repr(etid, ALL).get(key, None) is None:
                    continue
                self._graph._pop_e_repr(etid, key)
        else:
            self._graph._pop_e_repr(self._etid, key)

    def __len__(self):
        assert isinstance(self._etype, list) is False, \
            'Current HeteroEdgeDataView has multiple edge types, ' \
            'can not support len().'
        return len(self._graph._edge_frames[self._etid])

    def __iter__(self):
        assert isinstance(self._etype, list) is False, \
            'Current HeteroEdgeDataView has multiple edge types, ' \
            'can not be iterated.'
        return iter(self._graph._edge_frames[self._etid])

    def keys(self):
        return self._graph._edge_frames[self._etid].keys()

    def values(self):
        return self._graph._edge_frames[self._etid].values()

    def __repr__(self):
        if isinstance(self._etype, list):
            ret = {}
            for (i, etype) in enumerate(self._etype):
                data = self._graph._get_e_repr(self._etid[i], self._edges)
                value = {key : data[key]
                         for key in self._graph._edge_frames[self._etid[i]]}
                ret[etype] = value
            return repr(ret)
        else:
            data = self._graph._get_e_repr(self._etid, self._edges)
            return repr({key : data[key]
                         for key in self._graph._edge_frames[self._etid]})
