import torch


def preprocess_adj(adj, device):
    """
    Perform the processing of the adjacency matrix proposed by Kipf et al. 2017.

    Parameters
    ----------
    adj: sp.spmatrix
        Input adjacency matrix.

    Returns
    -------
    The matrix (D+1)^(-0.5) (adj + I) (D+1)^(-0.5)

    """
    adj_ = adj + torch.eye(adj.shape[0], device=device)
    rowsum = torch.sum(adj_, dim=0)
    degree_mat_inv_sqrt = torch.diag(torch.pow(rowsum, -0.5))
    adj_norm = torch.mm(torch.mm(adj_, degree_mat_inv_sqrt.transpose(1, 0)),
                        degree_mat_inv_sqrt)
    return adj_norm