import torch
from typing import *

from deepsnap.graph import Graph


class Batch(Graph):
    r"""
    A plain old python object modeling a batch of :class:`deepsnap.graph.Graph` objects 
    as one big (disconnected) graph, with :class:`torch_geometric.data.Data` being the 
    base class, all its methods can also be used here.
    In addition, single graphs can be reconstructed via the assignment vector
    :obj:`batch`, which maps each node to its respective graph identifier.
    """
    def __init__(self, batch=None, **kwargs):
        super(Batch, self).__init__(**kwargs)

        self.batch = batch
        self.__data_class__ = Graph
        self.__slices__ = None

    @staticmethod
    def collate(follow_batch=[], transform=None, **kwargs):
        return lambda batch: Batch.from_data_list(
            batch, follow_batch, transform, **kwargs)

    @staticmethod
    def from_data_list(data_list: List[Graph], follow_batch: List = None,
                       transform: Callable = None, **kwargs):
        r"""
        Constructs A :class:`deepsnap.batch.Batch` object from a python list holding
        :class:`torch_geometric.data.Data` objects.
        The assignment vector :obj:`batch` is created on the fly.
        Additionally, creates assignment batch vectors for each key in
        :obj:`follow_batch`.

        Args:
            data_list (list): A list of `deepsnap.graph.Graph` objects.
            follow_batch (list, optional): Creates assignment batch vectors for each key.
            transform: If apply transform when batching.
            **kwargs: Other parameters.
        """
        if follow_batch is None:
            follow_batch = []
        if transform is not None:
            data_list = [data.apply_transform(
                            transform,
                            deep_copy=True,
                            **kwargs)
                         for data in data_list]
        keys = [set(data.keys) for data in data_list]
        keys = list(set.union(*keys))
        assert 'batch' not in keys

        batch, cumsum = Batch._init_batch_fields(keys, follow_batch)
        batch.__data_class__ = data_list[0].__class__
        batch.batch = []
        for i, data in enumerate(data_list):
            # Note: in heterogeneous graph, __inc__ logic is different
            Batch._collate_dict(data, cumsum,
                                batch.__slices__, batch,
                                data,
                                follow_batch)
            num_nodes = data.num_nodes
            if num_nodes is not None:
                item = torch.full((num_nodes, ), i, dtype=torch.long)
                batch.batch.append(item)

        if num_nodes is None:
            batch.batch = None

        Batch._dict_list_to_tensor(batch, data_list[0])

        return batch.contiguous()

    @staticmethod
    def _init_batch_fields(keys, follow_batch):
        batch = Batch()
        batch.__slices__ = {key: [0] for key in keys}

        for key in keys:
            batch[key] = []

        for key in follow_batch:
            batch['{}_batch'.format(key)] = []

        cumsum = {key: 0 for key in keys}
        return batch, cumsum

    @staticmethod
    def _collate_dict(curr_dict, cumsum: Dict[str, int], slices, batched_dict, graph, follow_batch):
        r""" Called in from_data_list to collate a dictionary.
        This can also be applied to Graph object, since it has support for
        keys and __getitem__().

        Args:
            curr_dict: current dictionary to be added to the collated dictionary.
            cumsum: cumulative sum to be used for indexing.
            slices: a dictionary of the same structure as batched_dict,
                slices[key] indicates the indices to slice batch[key] into 
                tensors for all graphs in the batch.
            batched_dict: the batched dictionary of the same structure as curr_dict.
                But all graph data are batched together.
        """
        if isinstance(curr_dict, dict):
            keys = curr_dict.keys()
        else:
            keys = curr_dict.keys
        for key in keys:
            item = curr_dict[key]
            if isinstance(item, dict):
                # recursively collate every key in the dictionary
                if isinstance(batched_dict[key], list):
                    # nested dictionary not initialized yet
                    assert len(batched_dict[key]) == 0
                    # initialize the nested dictionary for batch
                    cumsum[key] = {inner_key: 0 for inner_key in item.keys()}
                    slices[key] = {inner_key: [0] for inner_key in item.keys()}
                    batched_dict[key] = {}
                    for inner_key in item.keys():
                        batched_dict[key][inner_key] = []
                    for inner_key in follow_batch:
                        batched_dict[key]['{}_batch'.format(key)] = []
                Batch._collate_dict(item, cumsum[key],
                                    slices[key], batched_dict[key],
                                    graph, follow_batch)
                continue
            if torch.is_tensor(item) and item.dtype != torch.bool:
                item = item + cumsum[key]
            if torch.is_tensor(item):
                size = item.size(graph.__cat_dim__(key, curr_dict[key]))
            else:
                size = 1
            slices[key].append(size + slices[key][-1])
            cumsum[key] = cumsum[key] + graph.__inc__(key, item)
            batched_dict[key].append(item)

            if key in follow_batch:
                item = torch.full((size, ), i, dtype=torch.long)
                batched_dict['{}_batch'.format(key)].append(item)

    @staticmethod
    def _dict_list_to_tensor(dict_of_list, graph):
        r"""Convert a dict/Graph with list as values to a dict/Graph with
        concatenated/stacked tensor as values.
        """
        if isinstance(dict_of_list, dict):
            keys = dict_of_list.keys()
        else:
            keys = dict_of_list.keys
        for key in keys:
            if isinstance(dict_of_list[key], dict):
                # recursively convert the dictionary of list to dict of tensor
                Batch._dict_list_to_tensor(dict_of_list[key], graph)
                continue
            item = dict_of_list[key][0]
            if torch.is_tensor(item):
                if Graph._is_graph_attribute(key) and item.ndim == 1 and (
                        not item.dtype == torch.long):
                    # special consideration: 1D tensor for graph attribute (classification)
                    # batch by stacking the first dim
                    dict_of_list[key] = torch.stack(
                        dict_of_list[key],
                        dim=0)
                else:
                    # concat at the __cat_dim__
                    dict_of_list[key] = torch.cat(
                        dict_of_list[key],
                        dim=graph.__cat_dim__(key, item))
            elif isinstance(item, (float, int)):
                dict_of_list[key] = torch.tensor(dict_of_list[key])

    def to_data_list(self):
        r"""
        Reconstructs the list of :class:`torch_geometric.data.Data` objects
        from the batch object.
        The batch object must have been created via :meth:`from_data_list` in
        order to be able reconstruct the initial objects.
        """
        if self.__slices__ is None:
            raise RuntimeError(
                ('Cannot reconstruct data list from batch because the batch '
                 'object was not created using Batch.from_data_list()'))

        keys = [key for key in self.keys if key[-5:] != 'batch']
        cumsum = {key: 0 for key in keys}
        data_list = []
        for i in range(len(self.__slices__[keys[0]]) - 1):
            # i: from 0 up to num graphs in the batch
            data = self.__data_class__()
            self._reconstruct_dict(
                i, keys, data, cumsum, self.__slices__, self, data)
            data_list.append(data)

        return data_list

    def _reconstruct_dict(
            self, graph_idx: int, keys, data_dict,
            cumsum: Dict[str, int], slices, batched_dict, graph):

        for key in keys:
            if isinstance(batched_dict[key], dict):
                # recursively unbatch the dict
                data_dict[key] = {}
                inner_keys = [inner_key for inner_key in batched_dict[key].keys()
                              if inner_key[-5:] != 'batch']
                inner_cumsum = {inner_key: 0 for inner_key in inner_keys}
                inner_slices = slices[key]
                self._reconstruct_dict(graph_idx, inner_keys,
                                       data_dict[key], inner_cumsum,
                                       inner_slices,
                                       batched_dict[key], graph)
                continue

            if torch.is_tensor(batched_dict[key]):
                data_dict[key] = batched_dict[key].narrow(
                    graph.__cat_dim__(key, batched_dict[key]),
                    slices[key][graph_idx],
                    slices[key][graph_idx + 1] - slices[key][graph_idx])
                if batched_dict[key].dtype != torch.bool:
                    data_dict[key] = data_dict[key] - cumsum[key]
            else:
                data_dict[key] = batched_dict[key][
                        slices[key][graph_idx]:slices[key][graph_idx + 1]]
            cumsum[key] = cumsum[key] + graph.__inc__(key, data_dict[key])

    @property
    def num_graphs(self) -> int:
        r"""
        Returns the number of graphs in the batch.

        Returns:
            int: The number of graphs in the batch.
        """
        return self.batch[-1].item() + 1

    def apply_transform(self, transform, update_tensor: bool = True,
                        update_graph: bool = False,
                        deep_copy: bool = False, **kwargs):
        r"""
        Applies a transformation to each graph object in parallel by first
        calling `to_data_list`, applying the transform, and then perform
        re-batching again to a `Batch`.
        A transform should edit the graph object,
        including changing the graph structure, and adding node/edge/graph attributes.
        The rest are automatically handled by the :class:`deepsnap.graph.Graph` object,
        including everything ended with index.

        Args:
            transform: Transformation function applied to each graph object.
            update_tensor: Whether use nx graph to update tensor attributes.
            update_graph: Whether use tensor attributes to update nx graphs.
            deep_copy: :obj:`True` if a new deep copy of batch is returned.
                This option allows modifying the batch of graphs without
                changing the graphs in the original dataset.
            kwargs: Parameters used in transform function in :class:`deepsnap.graph.Graph` objects.
        """
        # TODO: transductive setting, assert update_tensor == True
        return self.from_data_list([
            Graph(graph).apply_transform(
                transform, update_tensor, update_graph, deep_copy, **kwargs)
            for graph in self.G])

    def apply_transform_batched(self, transform):
        r"""
        A transform that directly operates on batched graphs
        User customized apply for batched graphs (expert-only)

        Args:
            transform: Transformation function applied to each graph object.
        """
        raise NotImplementedError
