import logging
import math
from typing import (
    Any,
    Callable,
    Iterator,
    List,
    NamedTuple,
    Optional,
    Tuple,
    Union,
)

import numpy as np
import torch
from torch import Tensor
from tqdm import tqdm

from torch_geometric.data import Data
from torch_geometric.typing import SparseTensor
from torch_geometric.utils import get_ppr, is_undirected, subgraph

try:
    import numba
    WITH_NUMBA = True
except ImportError:  # pragma: no cover
    WITH_NUMBA = False


class OutputNodes(NamedTuple):
    seed_id: Tensor
    auxiliary_id: Tensor


class _IBMBBaseLoader(torch.utils.data.DataLoader):
    def __init__(self, data: Data, **kwargs):
        kwargs.pop('collate_fn', None)
        batch_size = kwargs.get('batch_size', 1)

        output_nodes = self.get_output_nodes(self)

        if batch_size == 1:  # Pre-process subgraphs:
            data_list = ...
            super().__init__(data_list, collate_fn=self._cache_fn, **kwargs)
        else:
            self.data = data
            super().__init__(output_nodes, collate_fn=self._collate_fn,
                             **kwargs)

    def get_output_nodes(self) -> List[OutputNodes]:
        raise NotImplementedError

    def _cache_fn(self, data_list: List[Data]) -> Data:
        assert len(data_list) == 1
        return data_list[0]

    def _collate_fn(self, output_nodes: List[OutputNodes]) -> Data:
        raise NotImplementedError

    def __repr__(self) -> str:
        return f'{self.__class__.__name__}()'


###############################################################################


def get_partitions(
    edge_index: Union[Tensor, SparseTensor],
    num_partitions: int,
    indices: Tensor,
    num_nodes: int,
    output_weight: Optional[float] = None,
) -> List[Tensor]:
    assert isinstance(
        edge_index,
        (torch.LongTensor,
         SparseTensor)), f'Unsupported edge_index type {type(edge_index)}'
    if isinstance(edge_index, torch.LongTensor):
        edge_index = SparseTensor.from_edge_index(
            edge_index, sparse_sizes=(num_nodes, num_nodes))

    if output_weight is not None and output_weight != 1:
        node_weight = torch.ones(num_nodes)
        node_weight[indices] = output_weight
    else:
        node_weight = None

    _, partptr, perm = edge_index.partition(num_parts=num_partitions,
                                            recursive=False, weighted=False,
                                            node_weight=node_weight)

    partitions = []
    for i in range(len(partptr) - 1):
        partitions.append(perm[partptr[i]:partptr[i + 1]])

    return partitions


def get_pair_wise_distance(
    ys: List,
    num_classes: int,
    dist_type: str = 'kl',
) -> np.ndarray:
    num_batches = len(ys)

    counts = np.zeros((num_batches, num_classes), dtype=np.int32)
    for i in range(num_batches):
        unique, count = np.unique(ys[i], return_counts=True)
        counts[i, unique] = count

    counts += 1
    counts = counts / counts.sum(1).reshape(-1, 1)
    pairwise_dist = np.zeros((num_batches, num_batches), dtype=np.float64)

    for i in range(0, num_batches - 1):
        for j in range(i + 1, num_batches):
            if dist_type == 'l1':
                pairwise_dist[i, j] = np.sum(np.abs(counts[i] - counts[j]))
            elif dist_type == 'kl':

                def kl_divergence(p: np.ndarray, q: np.ndarray):
                    return (p * np.log(p / q)).sum()

                pairwise_dist[i, j] = kl_divergence(counts[i],
                                                    counts[j]) + kl_divergence(
                                                        counts[j], counts[i])
            else:
                raise ValueError

    pairwise_dist += pairwise_dist.T
    pairwise_dist += 1e-5  # for numerical stability
    np.fill_diagonal(pairwise_dist, 0.)

    return pairwise_dist


def indices_complete_check(
    loader: List[Tuple[Union[Tensor, np.ndarray], Union[Tensor, np.ndarray]]],
    output_indices: Union[Tensor, np.ndarray],
):
    if isinstance(output_indices, Tensor):
        output_indices = output_indices.cpu().numpy()

    outs = []
    for out, aux in loader:
        if isinstance(out, Tensor):
            out = out.cpu().numpy()
        if isinstance(aux, Tensor):
            aux = aux.cpu().numpy()

        assert np.all(np.in1d(out,
                              aux)), "Not all output nodes are in aux nodes!"
        outs.append(out)

    outs = np.sort(np.concatenate(outs))
    assert np.all(
        outs == np.sort(output_indices)), "Output nodes missing or duplicate!"


def get_subgraph(
    out_indices: Tensor,
    graph: Data,
    return_edge_index_type: str,
    adj: SparseTensor,
    **kwargs,
):
    if return_edge_index_type == 'adj':
        assert adj is not None

    if return_edge_index_type == 'adj':
        subg = Data(x=graph.x[out_indices], y=graph.y[out_indices],
                    edge_index=adj[out_indices, :][:, out_indices])
    elif return_edge_index_type == 'edge_index':
        edge_index, edge_attr = subgraph(out_indices, graph.edge_index,
                                         graph.edge_attr, relabel_nodes=True,
                                         num_nodes=graph.num_nodes,
                                         return_edge_mask=False)
        subg = Data(x=graph.x[out_indices], y=graph.y[out_indices],
                    edge_index=edge_index, edge_attr=edge_attr)
    else:
        raise NotImplementedError

    for k, v in kwargs.items():
        subg[k] = v

    return subg


def define_sampler(
    batch_order: str,
    ys: List[Union[Tensor, np.ndarray, List]],
    num_classes: int,
    dist_type: str = 'kl',
):
    if batch_order == 'rand':
        logging.info("Running with random order")
        sampler = torch.utils.data.RandomSampler(ys)
    elif batch_order in ['order', 'sample']:
        kl_div = get_pair_wise_distance(ys, num_classes, dist_type=dist_type)
        if batch_order == 'order':
            from python_tsp.heuristics import solve_tsp_simulated_annealing
            best_perm, _ = solve_tsp_simulated_annealing(kl_div)
            logging.info(f"Running with given order: {best_perm}")
            sampler = IBMBOrderedSampler(best_perm)
        else:
            logging.info("Running with weighted sampling")
            sampler = IBMBWeightedSampler(kl_div)
    else:
        raise ValueError

    return sampler


def create_batchwise_out_aux_pairs(
    adj: SparseTensor,
    partitions: List[Union[torch.LongTensor, np.ndarray]],
    prime_indices: Union[torch.LongTensor, np.ndarray],
    topk: int,
    num_outnodeset_per_batch: int = 50,
    alpha: float = 0.2,
    ppr_iterations: int = 50,
) -> List[Tuple[np.ndarray, np.ndarray]]:
    def ppr_power_method(
        adj: SparseTensor,
        batch: List[Union[np.ndarray, torch.LongTensor]],
        topk: int,
        num_iter: int,
        alpha: float,
    ) -> List[np.ndarray]:

        topk_neighbors = []
        logits = torch.zeros(
            adj.size(0), len(batch),
            device=adj.device())  # each column contains a set of output nodes
        for i, tele_set in enumerate(batch):
            logits[tele_set, i] = 1. / len(tele_set)

        new_logits = logits.clone()
        for i in range(num_iter):
            new_logits = adj @ new_logits * (1 - alpha) + alpha * logits

        inds = new_logits.argsort(0)
        nonzeros = (new_logits > 0).sum(0)
        nonzeros = torch.minimum(
            nonzeros,
            torch.tensor([topk], dtype=torch.int64, device=adj.device()))
        for i in range(new_logits.shape[1]):
            topk_neighbors.append(inds[-nonzeros[i]:, i].cpu().numpy())

        return topk_neighbors

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    if isinstance(prime_indices, Tensor):
        prime_indices = prime_indices.cpu().numpy()

    adj = adj.to(device)

    cur_output_nodes = []
    loader = []

    pbar = tqdm(range(len(partitions)))
    pbar.set_description("Processing topic-sensitive PPR batches")
    for n in pbar:
        part = partitions[n]
        if isinstance(part, Tensor):
            part = part.cpu().numpy()

        primes_in_part, *_ = np.intersect1d(part, prime_indices,
                                            assume_unique=True,
                                            return_indices=True)
        if len(primes_in_part):  # no output nodes in this partition
            cur_output_nodes.append(primes_in_part)

        # accumulate enough output nodes to make good use of GPU memory
        if len(cur_output_nodes
               ) >= num_outnodeset_per_batch or n == len(partitions) - 1:
            topk_neighbors = ppr_power_method(adj, cur_output_nodes, topk,
                                              ppr_iterations, alpha)
            for i in range(len(cur_output_nodes)):
                # force output nodes to be aux nodes
                auxiliary_nodes = np.union1d(cur_output_nodes[i],
                                             topk_neighbors[i])
                loader.append((cur_output_nodes[i], auxiliary_nodes))
            cur_output_nodes = []

        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    return loader


def get_pairs(ppr_mat: Any) -> np.ndarray:
    ppr_mat = ppr_mat + ppr_mat.transpose()

    ppr_mat = ppr_mat.tocoo()
    row, col, data = ppr_mat.row, ppr_mat.col, ppr_mat.data
    mask = (row > col)  # lu

    row, col, data = row[mask], col[mask], data[mask]
    sort_arg = np.argsort(data)[::-1]
    # sort_arg = parallel_sort.parallel_argsort(data)[::-1]

    # map prime_nodes to arange
    ppr_pairs = np.vstack((row[sort_arg], col[sort_arg])).T
    return ppr_pairs


_prime_orient_merge_numba: Optional[Callable] = None


def prime_orient_merge(
    ppr_pairs: np.ndarray,
    primes_per_batch: int,
    num_nodes: int,
):
    if not WITH_NUMBA:  # pragma: no cover
        raise ImportError("'prime_orient_merge' requires the 'numba' package")

    global _prime_orient_merge_numba
    if _prime_orient_merge_numba is None:
        _prime_orient_merge_numba = numba.njit(cache=True)(_prime_orient_merge)

    return _prime_orient_merge_numba(ppr_pairs, primes_per_batch, num_nodes)


def _prime_orient_merge(
    ppr_pairs: np.ndarray,
    primes_per_batch: int,
    num_nodes: int,
):
    id_primes_list = list(np.arange(num_nodes, dtype=np.int32).reshape(-1, 1))
    node_id_list = np.arange(num_nodes, dtype=np.int32)
    placeholder = np.zeros(0, dtype=np.int32)

    for i, j in ppr_pairs:
        id1, id2 = node_id_list[i], node_id_list[j]
        if id1 > id2:
            id1, id2 = id2, id1

        if id1 != id2 and len(id_primes_list[id1]) + len(
                id_primes_list[id2]) <= primes_per_batch:
            id_primes_list[id1] = np.concatenate(
                (id_primes_list[id1], id_primes_list[id2]))
            node_id_list[id_primes_list[id2]] = id1
            id_primes_list[id2] = placeholder

    prime_lst = list()
    ids = np.unique(node_id_list)

    for _id in ids:
        prime_lst.append(list(id_primes_list[_id]))

    return list(prime_lst)


def prime_post_process(loader, merge_max_size):
    from heapq import heapify, heappop, heappush

    h = [(
        len(p),
        p,
    ) for p in loader]
    heapify(h)

    while len(h) > 1:
        len1, p1 = heappop(h)
        len2, p2 = heappop(h)
        if len1 + len2 <= merge_max_size:
            heappush(h, (len1 + len2, p1 + p2))
        else:
            heappush(h, (
                len1,
                p1,
            ))
            heappush(h, (
                len2,
                p2,
            ))
            break

    new_batch = []

    while len(h):
        _, p = heappop(h)
        new_batch.append(p)

    return new_batch


def topk_ppr_matrix(
    edge_index: Tensor,
    num_nodes: int,
    alpha: float,
    eps: float,
    output_node_indices: Union[np.ndarray, torch.LongTensor],
    topk: int,
    normalization='row',
) -> Tuple[Any, List[np.ndarray]]:
    neighbors, weights = get_ppr(edge_index, alpha, eps, output_node_indices,
                                 num_nodes)

    _, neighbor_counts = neighbors[0].unique(return_counts=True)

    ppr_matrix = SparseTensor(
        row=torch.arange(
            len(output_node_indices)).repeat_interleave(neighbor_counts),
        col=neighbors[1], value=weights,
        sparse_sizes=(len(output_node_indices),
                      num_nodes)).to_scipy(layout='csr')

    neighbors = [
        n.cpu().numpy()
        for n in torch.split(neighbors[1],
                             neighbor_counts.cpu().tolist(), dim=0)
    ]
    weights = [
        n.cpu().numpy()
        for n in torch.split(weights,
                             neighbor_counts.cpu().tolist(), dim=0)
    ]

    def sparsify(neighbors: List[np.ndarray], weights: List[np.ndarray],
                 topk: int):
        new_neighbors = []
        for n, w in zip(neighbors, weights):
            idx_topk = np.argsort(w)[-topk:]
            new_neighbor = n[idx_topk]
            new_neighbors.append(new_neighbor)

        return new_neighbors

    neighbors = sparsify(neighbors, weights, topk)
    neighbors = [
        np.union1d(nei, pr) for nei, pr in zip(neighbors, output_node_indices)
    ]

    _, out_degree = torch.unique(edge_index[0], sorted=True,
                                 return_counts=True)
    if normalization == 'sym':
        # Assume undirected (symmetric) adjacency matrix
        deg_sqrt = np.sqrt(np.maximum(out_degree, 1e-12))
        deg_inv_sqrt = 1. / deg_sqrt

        row, col = ppr_matrix.nonzero()
        ppr_matrix.data = deg_sqrt[output_node_indices[row]] * \
            ppr_matrix.data * \
            deg_inv_sqrt[col]
    elif normalization == 'col':
        # Assume undirected (symmetric) adjacency matrix
        deg_inv = 1. / np.maximum(out_degree, 1e-12)

        row, col = ppr_matrix.nonzero()
        ppr_matrix.data = out_degree[output_node_indices[row]] * \
            ppr_matrix.data * \
            deg_inv[col]
    elif normalization == 'row':
        pass
    else:
        raise ValueError(f"Unknown PPR normalization: {normalization}")

    return ppr_matrix, neighbors


class IBMBBaseLoader(torch.utils.data.DataLoader):
    def __init__(
        self,
        data_list: Union[List[Data], List[Tuple]],
        graph: Data,
        adj: SparseTensor,
        return_edge_index_type: str,
        **kwargs,
    ):
        self.graph = graph
        self.adj = adj
        self.return_edge_index_type = return_edge_index_type
        if 'collate_fn' in kwargs:
            del kwargs['collate_fn']
        super().__init__(data_list, collate_fn=self.collate_fn, **kwargs)

    def create_loader(self, *args, **kwargs):
        raise NotImplementedError

    @classmethod
    def prepare_cache(
        cls,
        graph: Data,
        batch_wise_out_aux_pairs: List[Tuple[np.ndarray, np.ndarray]],
        adj: Optional[SparseTensor],
        return_edge_index_type: str,
    ):
        subgraphs = []

        pbar = tqdm(batch_wise_out_aux_pairs)
        pbar.set_description(
            f"Caching data with type {return_edge_index_type}")

        if return_edge_index_type == 'adj':
            assert adj is not None

        for out, aux in pbar:
            mask = torch.from_numpy(np.in1d(aux, out))
            if isinstance(aux, np.ndarray):
                aux = torch.from_numpy(aux)
            subg = get_subgraph(aux, graph, return_edge_index_type, adj,
                                output_node_mask=mask)
            subgraphs.append(subg)

        return subgraphs

    @classmethod
    def create_adj_from_edge_index(
        cls,
        edge_index: Tensor,
        num_nodes: int,
        normalization: str,
    ):
        assert normalization in ['sym', 'rw']
        adj = SparseTensor.from_edge_index(
            edge_index,
            sparse_sizes=(num_nodes, num_nodes),
        )
        adj = adj.fill_value(1.)
        degree = adj.sum(0)

        degree[degree == 0.] = 1e-12
        deg_inv = 1 / degree

        if normalization == 'sym':
            deg_inv_sqrt = deg_inv**0.5
            adj = adj * deg_inv_sqrt.reshape(1, -1)
            adj = adj * deg_inv_sqrt.reshape(-1, 1)
        elif normalization == 'rw':
            adj = adj * deg_inv.reshape(-1, 1)

        return adj

    def collate_fn(self, data_list: List[Union[Data, Tuple]]):
        if len(data_list) == 1 and isinstance(data_list[0], Data):
            return data_list[0]

        out, aux = zip(*data_list)
        out = np.concatenate(out)
        aux = np.unique(np.concatenate(aux))
        mask = torch.from_numpy(np.in1d(aux, out))
        aux = torch.from_numpy(aux)

        subg = get_subgraph(aux, self.graph, self.return_edge_index_type,
                            self.adj, output_node_mask=mask)
        return subg

    def __repr__(self) -> str:
        return f'{self.__class__.__name__}()'


class IBMBBatchLoader(IBMBBaseLoader):
    r"""The batch-wise influence-based data loader from the
    `"Influence-Based Mini-Batching for Graph Neural Networks"
    <https://arxiv.org/abs/2212.09083>`__ paper.

    First, the METIS graph partitioning algorithm separates the graph into
    :obj:`num_partitions` many partitions.
    Afterwards, input/seed nodes and their auxiliary nodes (found via
    topic-sensitive PageRank) are used to form a mini-batch.

    If :obj:`batch_size` is set to :obj:`1`, mini-batches are pre-calculated
    and cached in memory.
    Otherwise, only input nodes and their auxiliary nodes are pre-computed, and
    mini-batches are collated on-the-fly.

    Args:
        data (torch_geometric.data.Data): A
            :class:`~torch_geometric.data.Data` object.
        batch_order (str): A string indicating the batch order type (one of
            :obj:`"order"`, :obj:`"sample"` or :obj:`"rand"`).
            If :obj:`"order"`, calculates the pair-wise KL divergence between
            every two batches to organize an optimal order.
            If :obj:`"sample"`, samples the next batch w.r.t. the last one in
            which a batch with higher KL divergence score is more likely to be
            sampled.
            If :obj:`"rand"`, batches are generated randomly.
        num_partitions (int): The number of partitions.
        input_nodes (torch.Tensor): A vector containing the set of seed
            nodes.
        batch_expand_ratio (float, optional): The ratio between the returned
            batch size and the original partition size. For example, set it to
            :obj:`2.0` in case you would like the batch to have double the
            number of nodes as the size of its partition.
            (default: :obj:`1.0`)
        metis_input_node_weight (float, optional): The weights on the input
            nodes for METIS graph partitioning. (default: :obj:`None`)
        alpha (float, optional): The teleport probability of the PageRank
            calculation. (default: :obj:`0.2`)
        approximate_ppr_iterations (int, optional): The number of power
            iterations for PageRank calculation. (default: :obj:`50`)
        return_edge_index_type (str, optional): A string indicating the output
            type of edge indices (one of :obj:`"edge_index"` or :obj:`"adj"`).
            If set to :obj:`"adj"`, the :obj:`edge_index` of the batch will
            be a :class:`torch_sparse.SparseTensor`, otherwise a
            :class:`torch.Tensor`. (default: :obj:`"edge_index"`)
        **kwargs (optional): Additional arguments of
            :class:`torch.utils.data.DataLoader`, such as :obj:`batch_size`,
            :obj:`shuffle`, :obj:`drop_last` or :obj:`num_workers`.
    """
    def __init__(
        self,
        data: Data,
        batch_order: str,
        num_partitions: int,
        input_nodes: Tensor,
        batch_expand_ratio: Optional[float] = 1.0,
        metis_input_node_weight: Optional[float] = None,
        alpha: Optional[float] = 0.2,
        approximate_ppr_iterations: Optional[int] = 50,
        return_edge_index_type: str = 'edge_index',
        **kwargs,
    ):
        self.subgraphs = []
        self.batch_wise_out_aux_pairs = []

        assert is_undirected(
            data.edge_index,
            num_nodes=data.num_nodes), "Assume the graph to be undirected"
        assert batch_order in ['rand', 'sample', 'order'
                               ], f"Unsupported batch order: {batch_order}"

        adj = self.create_adj_from_edge_index(
            data.edge_index,
            data.num_nodes,
            normalization='rw',
        )

        self.cache_data = kwargs['batch_size'] == 1
        self.num_partitions = num_partitions
        self.output_indices = input_nodes
        assert return_edge_index_type in ['adj', 'edge_index']
        self.return_edge_index_type = return_edge_index_type
        self.batch_expand_ratio = batch_expand_ratio
        self.metis_output_weight = metis_input_node_weight
        self.num_outnodeset_per_batch = 50
        self.alpha = alpha
        self.approximate_ppr_iterations = approximate_ppr_iterations

        self.create_loader(data, adj)

        if len(self.batch_wise_out_aux_pairs) > 2:  # <= 2 order makes no sense
            ys = [
                data.y[out].numpy() for out, _ in self.batch_wise_out_aux_pairs
            ]
            sampler = define_sampler(batch_order, ys, data.y.max().item() + 1)
        else:
            sampler = None

        if not self.cache_data:
            cached_data = data  # need to cache the original graph
            if return_edge_index_type == 'adj':
                cached_adj = adj
            else:
                cached_adj = None
        else:
            cached_data = None
            cached_adj = None

        super().__init__(
            self.subgraphs
            if self.cache_data else self.batch_wise_out_aux_pairs,
            cached_data,
            cached_adj,
            return_edge_index_type,
            sampler=sampler,
            **kwargs,
        )

    def create_loader(self, graph: Data, adj: SparseTensor):
        partitions = get_partitions(
            adj,
            self.num_partitions,
            self.output_indices,
            graph.num_nodes,
            self.metis_output_weight,
        )

        # get output - auxiliary node pairs
        topk = math.ceil(self.batch_expand_ratio * graph.num_nodes /
                         self.num_partitions)
        batch_wise_out_aux_pairs = create_batchwise_out_aux_pairs(
            adj, partitions, self.output_indices, topk,
            self.num_outnodeset_per_batch, self.alpha,
            self.approximate_ppr_iterations)

        indices_complete_check(batch_wise_out_aux_pairs, self.output_indices)
        self.batch_wise_out_aux_pairs = batch_wise_out_aux_pairs

        if self.cache_data:
            self.subgraphs = self.prepare_cache(
                graph,
                batch_wise_out_aux_pairs,
                adj,
                self.return_edge_index_type,
            )


class IBMBNodeLoader(IBMBBaseLoader):
    r"""The node-wise influence-based data loader from the
    `"Influence-Based Mini-Batching for Graph Neural Networks"
    <https://arxiv.org/abs/2212.09083>`__ paper.

    First, the Personalized PageRank (PPR) score for each input node is
    computed, for which the :obj:`k` nodes with the highest scores are taken
    auxiliary nodes.
    Afterwards, input nodes are merged according to their pair-wise PPR scores.

    Similar to :class:`~torch_geometric.loader.IBMBBatchLoader`, subgraphs are
    cached in memory for :obj:`batch_size = 1`, and collated on-the-fly
    otherwise.

    Args:
        data (torch_geometric.data.Data): A
            :class:`~torch_geometric.data.Data` object.
        batch_order (str): A string indicating the batch order type (one of
            :obj:`"order"`, :obj:`"sample"` or :obj:`"rand"`).
            If :obj:`"order"`, calculates the pair-wise KL divergence between
            every two batches to organize an optimal order.
            If :obj:`"sample"`, samples the next batch w.r.t. the last one in
            which a batch with higher KL divergence score is more likely to be
            sampled.
            If :obj:`"rand"`, batches are generated randomly.
        input_nodes (torch.Tensor): A vector containing the set of seed
            nodes.
        num_auxiliary_nodes (int): The number of auxiliary nodes per input
            node.
        num_nodes_per_batch (int): The number of seed nodes per batch.
        alpha (float, optional): The teleport probability of the PageRank
            calculation. (default: :obj:`0.2`)
        eps (float, optional): The threshold for stopping the PPR calculation
            The smaller :obj`eps` is, the more accurate are the results of
            PPR calculation, but it also takes longer.
            (default: :obj:`1e-5`)
        return_edge_index_type (str, optional): A string indicating the output
            type of edge indices (one of :obj:`"edge_index"` or :obj:`"adj"`).
            If set to :obj:`"adj"`, the :obj:`edge_index` of the batch will
            be a :class:`torch_sparse.SparseTensor`, otherwise a
            :class:`torch.Tensor`. (default: :obj:`"edge_index"`)
        **kwargs (optional): Additional arguments of
            :class:`torch.utils.data.DataLoader`, such as :obj:`batch_size`,
            :obj:`shuffle`, :obj:`drop_last` or :obj:`num_workers`.
    """
    def __init__(
        self,
        data: Data,
        batch_order: str,
        input_nodes: torch.Tensor,
        num_auxiliary_nodes: int,
        num_nodes_per_batch: int,
        alpha: float = 0.2,
        eps: float = 1e-5,
        return_edge_index_type: str = 'edge_index',
        **kwargs,
    ):
        self.subgraphs = []
        self.node_wise_out_aux_pairs = []

        assert is_undirected(
            data.edge_index,
            num_nodes=data.num_nodes), "Assume the graph to be undirected"
        assert batch_order in ['rand', 'sample', 'order'
                               ], f"Unsupported batch order: {batch_order}"

        if return_edge_index_type == 'adj':
            adj = self.create_adj_from_edge_index(data.edge_index,
                                                  data.num_nodes,
                                                  normalization='rw')
        else:
            adj = None

        self.cache_data = kwargs['batch_size'] == 1
        self._batchsize = kwargs['batch_size']
        self.output_indices = input_nodes.numpy()
        assert return_edge_index_type in ['adj', 'edge_index']
        self.return_edge_index_type = return_edge_index_type
        self.num_auxiliary_node_per_output = num_auxiliary_nodes
        self.num_output_nodes_per_batch = num_nodes_per_batch
        self.alpha = alpha
        self.eps = eps

        self.create_loader(data, adj)

        if len(self.node_wise_out_aux_pairs) > 2:  # <= 2 order makes no sense
            ys = [
                data.y[out].numpy() for out, _ in self.node_wise_out_aux_pairs
            ]
            sampler = define_sampler(batch_order, ys, data.y.max().item() + 1)
        else:
            sampler = None

        if not self.cache_data:
            cached_graph = data  # need to cache the original graph
            cached_adj = adj
        else:
            cached_graph = None
            cached_adj = None

        super().__init__(
            self.subgraphs
            if self.cache_data else self.node_wise_out_aux_pairs,
            cached_graph,
            cached_adj,
            return_edge_index_type,
            sampler=sampler,
            **kwargs,
        )

    def create_loader(self, graph: Data, adj: SparseTensor):
        logging.info("Start PPR calculation")
        ppr_matrix, neighbors = topk_ppr_matrix(
            graph.edge_index, graph.num_nodes, self.alpha, self.eps,
            torch.from_numpy(self.output_indices),
            self.num_auxiliary_node_per_output)

        ppr_matrix = ppr_matrix[:, self.output_indices]

        logging.info("Getting PPR pairs")
        ppr_pairs = get_pairs(ppr_matrix)

        output_list = prime_orient_merge(
            ppr_pairs,
            self.num_output_nodes_per_batch,
            len(self.output_indices),
        )
        output_list = prime_post_process(
            output_list,
            self.num_output_nodes_per_batch,
        )
        node_wise_out_aux_pairs = []

        if isinstance(neighbors, list):
            neighbors = np.array(neighbors, dtype=object)

        def _union(inputs):
            return np.unique(np.concatenate(inputs))

        for p in output_list:
            node_wise_out_aux_pairs.append(
                (self.output_indices[p],
                 _union(neighbors[p]).astype(np.int64)))

        indices_complete_check(node_wise_out_aux_pairs, self.output_indices)
        self.node_wise_out_aux_pairs = node_wise_out_aux_pairs

        if self.cache_data:
            self.subgraphs = self.prepare_cache(
                graph,
                node_wise_out_aux_pairs,
                adj,
                self.return_edge_index_type,
            )


class IBMBOrderedSampler(torch.utils.data.Sampler[int]):
    r"""A sampler with given order, specially for IBMB loaders.

    Args:
        data_source (np.ndarray, torch.Tensor, List): A :obj:`np.ndarray`,
            :obj:`torch.Tensor`, or :obj:`List` data object. Contains the
            order of the batches.
    """
    def __init__(self, data_source: Union[np.ndarray, torch.Tensor,
                                          List]) -> None:
        self.data_source = data_source
        super().__init__(data_source)

    def __iter__(self) -> Iterator[int]:
        return iter(self.data_source)

    def __len__(self) -> int:
        return len(self.data_source)


class IBMBWeightedSampler(torch.utils.data.Sampler[int]):
    r"""A weighted sampler wrt the pair wise KL divergence.
    The very first batch after initialization is sampled randomly,
    with the next ones being sampled according to the last batch,
    including the first batch in the next round.

    Args:
        batch_kl_div (np.ndarray, torch.Tensor): A :obj:`np.ndarray` or
            :obj:`torch.Tensor`, each element [i, j] contains the pair wise
            KL divergence between batch i and j.
    """
    def __init__(self, batch_kl_div: Union[np.ndarray, torch.Tensor]) -> None:
        data_source = np.arange(batch_kl_div.shape[0])
        self.data_source = data_source
        self.batch_kl_div = batch_kl_div
        self.last_train_batch_id = 0
        super().__init__(data_source)

    def __iter__(self) -> Iterator[int]:
        probs = self.batch_kl_div.copy()

        last = self.last_train_batch_id
        num_batches = probs.shape[0]

        fetch_idx = []

        next_id = 0
        while np.any(probs):
            next_id = np.random.choice(num_batches, size=None, replace=False,
                                       p=probs[last] / probs[last].sum())
            last = next_id
            fetch_idx.append(next_id)
            probs[:, next_id] = 0.

        self.last_train_batch_id = next_id

        return iter(fetch_idx)

    def __len__(self) -> int:
        return len(self.data_source)