import random
from typing import Optional, Tuple, Union

import numpy as np
import torch
from torch import Tensor

from torch_geometric.utils import coalesce, degree, remove_self_loops

from .num_nodes import maybe_num_nodes


def negative_sampling(edge_index: Tensor,
                      num_nodes: Optional[Union[int, Tuple[int, int]]] = None,
                      num_neg_samples: Optional[int] = None,
                      method: str = "sparse",
                      force_undirected: bool = False) -> Tensor:
    r"""Samples random negative edges of a graph given by :attr:`edge_index`.

    Args:
        edge_index (LongTensor): The edge indices.
        num_nodes (int or Tuple[int, int], optional): The number of nodes,
            *i.e.* :obj:`max_val + 1` of :attr:`edge_index`.
            If given as a tuple, then :obj:`edge_index` is interpreted as a
            bipartite graph with shape :obj:`(num_src_nodes, num_dst_nodes)`.
            (default: :obj:`None`)
        num_neg_samples (int, optional): The (approximate) number of negative
            samples to return.
            If set to :obj:`None`, will try to return a negative edge for every
            positive edge. (default: :obj:`None`)
        method (string, optional): The method to use for negative sampling,
            *i.e.*, :obj:`"sparse"` or :obj:`"dense"`.
            This is a memory/runtime trade-off.
            :obj:`"sparse"` will work on any graph of any size, while
            :obj:`"dense"` can perform faster true-negative checks.
            (default: :obj:`"sparse"`)
        force_undirected (bool, optional): If set to :obj:`True`, sampled
            negative edges will be undirected. (default: :obj:`False`)

    :rtype: LongTensor
    """
    assert method in ['sparse', 'dense']

    size = num_nodes
    bipartite = isinstance(size, (tuple, list))
    size = maybe_num_nodes(edge_index) if size is None else size
    size = (size, size) if not bipartite else size
    force_undirected = False if bipartite else force_undirected

    idx, population = edge_index_to_vector(edge_index, size, bipartite,
                                           force_undirected)

    if idx.numel() >= population:
        return edge_index.new_empty((2, 0))

    if num_neg_samples is None:
        num_neg_samples = edge_index.size(1)
    if force_undirected:
        num_neg_samples = num_neg_samples // 2

    prob = 1. - idx.numel() / population  # Probability to sample a negative.
    sample_size = int(1.1 * num_neg_samples / prob)  # (Over)-sample size.

    neg_idx = None
    if method == 'dense':
        # The dense version creates a mask of shape `population` to check for
        # invalid samples.
        mask = idx.new_ones(population, dtype=torch.bool)
        mask[idx] = False
        for _ in range(3):  # Number of tries to sample negative indices.
            rnd = sample(population, sample_size, idx.device)
            rnd = rnd[mask[rnd]]  # Filter true negatives.
            neg_idx = rnd if neg_idx is None else torch.cat([neg_idx, rnd])
            if neg_idx.numel() >= num_neg_samples:
                neg_idx = neg_idx[:num_neg_samples]
                break
            mask[neg_idx] = False

    else:  # 'sparse'
        # The sparse version checks for invalid samples via `np.isin`.
        idx = idx.to('cpu')
        for _ in range(3):  # Number of tries to sample negative indices.
            rnd = sample(population, sample_size, device='cpu')
            mask = np.isin(rnd, idx)
            if neg_idx is not None:
                mask |= np.isin(rnd, neg_idx.to('cpu'))
            mask = torch.from_numpy(mask).to(torch.bool)
            rnd = rnd[~mask].to(edge_index.device)
            neg_idx = rnd if neg_idx is None else torch.cat([neg_idx, rnd])
            if neg_idx.numel() >= num_neg_samples:
                neg_idx = neg_idx[:num_neg_samples]
                break

    return vector_to_edge_index(neg_idx, size, bipartite, force_undirected)


def batched_negative_sampling(
    edge_index: Tensor,
    batch: Union[Tensor, Tuple[Tensor, Tensor]],
    num_neg_samples: Optional[int] = None,
    method: str = "sparse",
    force_undirected: bool = False,
) -> Tensor:
    r"""Samples random negative edges of multiple graphs given by
    :attr:`edge_index` and :attr:`batch`.

    Args:
        edge_index (LongTensor): The edge indices.
        batch (LongTensor or Tuple[LongTensor, LongTensor]): Batch vector
            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
            node to a specific example.
            If given as a tuple, then :obj:`edge_index` is interpreted as a
            bipartite graph connecting two different node types.
        num_neg_samples (int, optional): The number of negative samples to
            return. If set to :obj:`None`, will try to return a negative edge
            for every positive edge. (default: :obj:`None`)
        method (string, optional): The method to use for negative sampling,
            *i.e.*, :obj:`"sparse"` or :obj:`"dense"`.
            This is a memory/runtime trade-off.
            :obj:`"sparse"` will work on any graph of any size, while
            :obj:`"dense"` can perform faster true-negative checks.
            (default: :obj:`"sparse"`)
        force_undirected (bool, optional): If set to :obj:`True`, sampled
            negative edges will be undirected. (default: :obj:`False`)

    :rtype: LongTensor
    """
    if isinstance(batch, Tensor):
        src_batch, dst_batch = batch, batch
    else:
        src_batch, dst_batch = batch[0], batch[1]

    split = degree(src_batch[edge_index[0]], dtype=torch.long).tolist()
    edge_indices = torch.split(edge_index, split, dim=1)

    num_src = degree(src_batch, dtype=torch.long)
    cum_src = torch.cat([src_batch.new_zeros(1), num_src.cumsum(0)[:-1]])

    if isinstance(batch, Tensor):
        num_nodes = num_src.tolist()
        cumsum = cum_src
    else:
        num_dst = degree(dst_batch, dtype=torch.long)
        cum_dst = torch.cat([dst_batch.new_zeros(1), num_dst.cumsum(0)[:-1]])

        num_nodes = torch.stack([num_src, num_dst], dim=1).tolist()
        cumsum = torch.stack([cum_src, cum_dst], dim=1).unsqueeze(-1)

    neg_edge_indices = []
    for i, edge_index in enumerate(edge_indices):
        edge_index = edge_index - cumsum[i]
        neg_edge_index = negative_sampling(edge_index, num_nodes[i],
                                           num_neg_samples, method,
                                           force_undirected)
        neg_edge_index += cumsum[i]
        neg_edge_indices.append(neg_edge_index)

    return torch.cat(neg_edge_indices, dim=1)


def structured_negative_sampling(edge_index, num_nodes: Optional[int] = None,
                                 contains_neg_self_loops: bool = True):
    r"""Samples a negative edge :obj:`(i,k)` for every positive edge
    :obj:`(i,j)` in the graph given by :attr:`edge_index`, and returns it as a
    tuple of the form :obj:`(i,j,k)`.

    Args:
        edge_index (LongTensor): The edge indices.
        num_nodes (int, optional): The number of nodes, *i.e.*
            :obj:`max_val + 1` of :attr:`edge_index`. (default: :obj:`None`)
        contains_neg_self_loops (bool, optional): If set to
            :obj:`False`, sampled negative edges will not contain self loops.
            (default: :obj:`True`)

    :rtype: (LongTensor, LongTensor, LongTensor)
    """
    num_nodes = maybe_num_nodes(edge_index, num_nodes)

    row, col = edge_index.cpu()
    pos_idx = row * num_nodes + col
    if not contains_neg_self_loops:
        loop_idx = torch.arange(num_nodes) * (num_nodes + 1)
        pos_idx = torch.cat([pos_idx, loop_idx], dim=0)

    rand = torch.randint(num_nodes, (row.size(0), ), dtype=torch.long)
    neg_idx = row * num_nodes + rand

    mask = torch.from_numpy(np.isin(neg_idx, pos_idx)).to(torch.bool)
    rest = mask.nonzero(as_tuple=False).view(-1)
    while rest.numel() > 0:  # pragma: no cover
        tmp = torch.randint(num_nodes, (rest.size(0), ), dtype=torch.long)
        rand[rest] = tmp
        neg_idx = row[rest] * num_nodes + tmp

        mask = torch.from_numpy(np.isin(neg_idx, pos_idx)).to(torch.bool)
        rest = rest[mask]

    return edge_index[0], edge_index[1], rand.to(edge_index.device)


def structured_negative_sampling_feasible(
        edge_index: Tensor, num_nodes: Optional[int] = None,
        contains_neg_self_loops: bool = True) -> bool:
    r"""Returns :obj:`True` if
    :meth:`~torch_geometric.utils.structured_negative_sampling` is feasible
    on the graph given by :obj:`edge_index`.
    :obj:`~torch_geometric.utils.structured_negative_sampling` is infeasible
    if atleast one node is connected to all other nodes.

    Args:
        edge_index (LongTensor): The edge indices.
        num_nodes (int, optional): The number of nodes, *i.e.*
            :obj:`max_val + 1` of :attr:`edge_index`. (default: :obj:`None`)
        contains_neg_self_loops (bool, optional): If set to
            :obj:`False`, sampled negative edges will not contain self loops.
            (default: :obj:`True`)

    :rtype: bool
    """
    num_nodes = maybe_num_nodes(edge_index, num_nodes)
    max_num_neighbors = num_nodes

    edge_index = coalesce(edge_index, num_nodes=num_nodes)

    if not contains_neg_self_loops:
        edge_index, _ = remove_self_loops(edge_index)
        max_num_neighbors -= 1  # Reduce number of valid neighbors

    deg = degree(edge_index[0], num_nodes)
    # True if there exists no node that is connected to all other nodes.
    return bool(torch.all(deg < max_num_neighbors))


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


def sample(population: int, k: int, device=None) -> Tensor:
    if population <= k:
        return torch.arange(population, device=device)
    else:
        return torch.tensor(random.sample(range(population), k), device=device)


def edge_index_to_vector(
    edge_index: Tensor,
    size: Tuple[int, int],
    bipartite: bool,
    force_undirected: bool = False,
) -> Tuple[Tensor, int]:

    row, col = edge_index

    if bipartite:  # No need to account for self-loops.
        idx = (row * size[1]).add_(col)
        population = size[0] * size[1]
        return idx, population

    elif force_undirected:
        assert size[0] == size[1]
        num_nodes = size[0]

        # We only operate on the upper triangular matrix:
        mask = row < col
        row, col = row[mask], col[mask]
        offset = torch.arange(1, num_nodes, device=row.device).cumsum(0)[row]
        idx = row.mul_(num_nodes).add_(col).sub_(offset)
        population = (num_nodes * (num_nodes + 1)) // 2 - num_nodes
        return idx, population

    else:
        assert size[0] == size[1]
        num_nodes = size[0]

        # We remove self-loops as we do not want to take them into account
        # when sampling negative values.
        mask = row != col
        row, col = row[mask], col[mask]
        col[row < col] -= 1
        idx = row.mul_(num_nodes - 1).add_(col)
        population = num_nodes * num_nodes - num_nodes
        return idx, population


def vector_to_edge_index(idx: Tensor, size: Tuple[int, int], bipartite: bool,
                         force_undirected: bool = False) -> Tensor:

    if bipartite:  # No need to account for self-loops.
        row = idx.div(size[1], rounding_mode='floor')
        col = idx % size[1]
        return torch.stack([row, col], dim=0)

    elif force_undirected:
        assert size[0] == size[1]
        num_nodes = size[0]

        offset = torch.arange(1, num_nodes, device=idx.device).cumsum(0)
        end = torch.arange(num_nodes, num_nodes * num_nodes, num_nodes,
                           device=idx.device)
        row = torch.bucketize(idx, end.sub_(offset), right=True)
        col = offset[row].add_(idx) % num_nodes
        return torch.stack([torch.cat([row, col]), torch.cat([col, row])], 0)

    else:
        assert size[0] == size[1]
        num_nodes = size[0]

        row = idx.div(num_nodes - 1, rounding_mode='floor')
        col = idx % (num_nodes - 1)
        col[row <= col] += 1
        return torch.stack([row, col], dim=0)

[((134, 19, 134, 56), 'torch.split', 'torch.split', (), '', False, 'import torch\n'), ((136, 14, 136, 49), 'torch_geometric.utils.degree', 'degree', (), '', False, 'from torch_geometric.utils import coalesce, degree, remove_self_loops\n'), ((158, 11, 158, 45), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((223, 17, 223, 58), 'torch_geometric.utils.coalesce', 'coalesce', (), '', False, 'from torch_geometric.utils import coalesce, degree, remove_self_loops\n'), ((229, 10, 229, 42), 'torch_geometric.utils.degree', 'degree', ({(229, 17, 229, 30): 'edge_index[0]', (229, 32, 229, 41): 'num_nodes'}, {}), '(edge_index[0], num_nodes)', False, 'from torch_geometric.utils import coalesce, degree, remove_self_loops\n'), ((143, 18, 143, 53), 'torch_geometric.utils.degree', 'degree', (), '', False, 'from torch_geometric.utils import coalesce, degree, remove_self_loops\n'), ((183, 18, 183, 55), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((226, 24, 226, 53), 'torch_geometric.utils.remove_self_loops', 'remove_self_loops', ({(226, 42, 226, 52): 'edge_index'}, {}), '(edge_index)', False, 'from torch_geometric.utils import coalesce, degree, remove_self_loops\n'), ((231, 16, 231, 50), 'torch.all', 'torch.all', ({(231, 26, 231, 49): '(deg < max_num_neighbors)'}, {}), '(deg < max_num_neighbors)', False, 'import torch\n'), ((239, 15, 239, 54), 'torch.arange', 'torch.arange', (), '', False, 'import torch\n'), ((290, 15, 290, 45), 'torch.stack', 'torch.stack', (), '', False, 'import torch\n'), ((84, 19, 84, 36), 'numpy.isin', 'np.isin', ({(84, 27, 84, 30): 'rnd', (84, 32, 84, 35): 'idx'}, {}), '(rnd, idx)', True, 'import numpy as np\n'), ((133, 12, 133, 62), 'torch_geometric.utils.degree', 'degree', (), '', False, 'from torch_geometric.utils import coalesce, degree, remove_self_loops\n'), ((182, 19, 182, 42), 'torch.arange', 'torch.arange', ({(182, 32, 182, 41): 'num_nodes'}, {}), '(num_nodes)', False, 'import torch\n'), ((297, 14, 298, 45), 'torch.arange', 'torch.arange', (), '', False, 'import torch\n'), ((310, 15, 310, 45), 'torch.stack', 'torch.stack', (), '', False, 'import torch\n'), ((73, 50, 73, 75), 'torch.cat', 'torch.cat', ({(73, 60, 73, 74): '[neg_idx, rnd]'}, {}), '([neg_idx, rnd])', False, 'import torch\n'), ((89, 50, 89, 75), 'torch.cat', 'torch.cat', ({(89, 60, 89, 74): '[neg_idx, rnd]'}, {}), '([neg_idx, rnd])', False, 'import torch\n'), ((146, 20, 146, 58), 'torch.stack', 'torch.stack', (), '', False, 'import torch\n'), ((147, 17, 147, 55), 'torch.stack', 'torch.stack', (), '', False, 'import torch\n'), ((188, 28, 188, 53), 'numpy.isin', 'np.isin', ({(188, 36, 188, 43): 'neg_idx', (188, 45, 188, 52): 'pos_idx'}, {}), '(neg_idx, pos_idx)', True, 'import numpy as np\n'), ((87, 19, 87, 41), 'torch.from_numpy', 'torch.from_numpy', ({(87, 36, 87, 40): 'mask'}, {}), '(mask)', False, 'import torch\n'), ((195, 32, 195, 57), 'numpy.isin', 'np.isin', ({(195, 40, 195, 47): 'neg_idx', (195, 49, 195, 56): 'pos_idx'}, {}), '(neg_idx, pos_idx)', True, 'import numpy as np\n'), ((296, 17, 296, 62), 'torch.arange', 'torch.arange', (), '', False, 'import torch\n'), ((301, 28, 301, 49), 'torch.cat', 'torch.cat', ({(301, 38, 301, 48): '[row, col]'}, {}), '([row, col])', False, 'import torch\n'), ((301, 51, 301, 72), 'torch.cat', 'torch.cat', ({(301, 61, 301, 71): '[col, row]'}, {}), '([col, row])', False, 'import torch\n'), ((265, 17, 265, 62), 'torch.arange', 'torch.arange', (), '', False, 'import torch\n')]