import torch as th

from ..._sparse_ops import (
    _bwd_segment_cmp,
    _csrmask,
    _csrmm,
    _csrsum,
    _edge_softmax_backward,
    _edge_softmax_forward,
    _gather_mm,
    _gather_mm_scatter,
    _gsddmm,
    _gsddmm_hetero,
    _gspmm,
    _gspmm_hetero,
    _scatter_add,
    _segment_mm,
    _segment_mm_backward_B,
    _segment_reduce,
    _update_grad_minmax_hetero,
)

from ...base import ALL, is_all
from ...heterograph_index import create_unitgraph_from_csr

__all__ = [
    "gspmm",
    "gsddmm",
    "gspmm_hetero",
    "gsddmm_hetero",
    "edge_softmax",
    "edge_softmax_hetero",
    "segment_reduce",
    "scatter_add",
    "csrmm",
    "csrsum",
    "csrmask",
    "gather_mm",
    "segment_mm",
]


def _reduce_grad(grad, shape):
    """Reduce gradient on the broadcast dimension
    If there is broadcast in forward pass, gradients need to be reduced on
    broadcast dimension. This function checks the input tensor shape and
    gradient shape and perform the reduction.

    Parameters
    ----------
    grad: Tensor
        Gradient tensor
    shape: tuple
        Shape of input tensor

    Returns
    -------
    Tensor
    """
    grad_shape = grad.shape[1:]
    in_shape = shape[1:]
    if in_shape == grad_shape:
        # no need to reduce
        return grad
    num_to_squeeze = len(grad_shape) - len(in_shape)
    # pad inshape
    in_shape = (1,) * num_to_squeeze + in_shape
    reduce_idx = th.nonzero(
        th.tensor(grad_shape) - th.tensor(in_shape), as_tuple=False
    )
    reduce_idx += 1  # skip batch dim
    if len(reduce_idx) > 0:
        grad = grad.sum(dim=tuple(reduce_idx), keepdim=True)
    return grad.view(-1, *shape[1:])


def _need_reduce_last_dim(ufeat, efeat):
    """Indicates whether to reduce the last dimension on edges
    in the backward pass of spmm,
    if so, use dot instead of mul."""
    if ufeat is None or efeat is None:
        return False
    ushp = ufeat.shape
    eshp = efeat.shape
    return ushp[1:-1] == eshp[1:-1] and eshp[-1] == 1 and ushp[-1] > 1


def _expand(x, shape):
    return x.expand(-1, *shape)


def spmm_cache_X(binary_op, reduce_op, req_grad_X, req_grad_Y):
    """Rules to identify whether to cache X in SpMM forward stage."""
    if binary_op != "copy_lhs" and req_grad_Y:
        if reduce_op == "sum":
            return True
        else:
            if binary_op == "mul":
                return True
    return False


def spmm_cache_Y(binary_op, reduce_op, req_grad_X, req_grad_Y):
    """Rules to identify whether to cache Y in SpMM forward stage."""
    if binary_op != "copy_rhs" and req_grad_X:
        if reduce_op == "sum":
            if binary_op in ["mul", "add"]:
                return True
        else:
            if binary_op == "mul":
                return True
    return False


def spmm_cache_argX(binary_op, reduce_op, req_grad_X, req_grad_Y):
    """Rules to identify whether to cache argX in SpMM forward stage."""
    if req_grad_X or req_grad_Y:
        if reduce_op in ["min", "max"]:
            return True
    return False


def spmm_cache_argY(binary_op, reduce_op, req_grad_X, req_grad_Y):
    """Rules to identify whether to cache argY in SpMM forward stage."""
    if req_grad_X or req_grad_Y:
        if reduce_op in ["min", "max"]:
            return True
    return False


class empty_context:
    """Empty context that does nothing"""

    def __init__(self, *args, **kargs):
        return

    def __enter__(self, *args, **kargs):
        return self

    def __exit__(self, *args, **kargs):
        return


# Disable CUDA autocast since we have casted args manually,
# and do it only in a nested autocast context.
def _disable_autocast_if_enabled():
    if th.is_autocast_enabled():
        return th.cuda.amp.autocast(enabled=False)
    else:
        return empty_context()


def _cast_if_autocast_enabled(*args):
    if not th.is_autocast_enabled():
        return args
    else:
        return th.cuda.amp.autocast_mode._cast(
            args, th.get_autocast_gpu_dtype()
        )


class GSpMM(th.autograd.Function):
    @staticmethod
    def forward(ctx, gidx, op, reduce_op, X, Y):
        out, (argX, argY) = _gspmm(gidx, op, reduce_op, X, Y)
        reduce_last = _need_reduce_last_dim(X, Y)
        X_shape = X.shape if X is not None else None
        Y_shape = Y.shape if Y is not None else None
        dtype = X.dtype if X is not None else Y.dtype
        device = X.device if X is not None else Y.device
        ctx.backward_cache = (
            gidx,
            op,
            reduce_op,
            X_shape,
            Y_shape,
            dtype,
            device,
            reduce_last,
        )
        req_grad_X = X.requires_grad if X is not None else False
        req_grad_Y = Y.requires_grad if Y is not None else False
        if not spmm_cache_X(op, reduce_op, req_grad_X, req_grad_Y):
            X = None
        if not spmm_cache_Y(op, reduce_op, req_grad_X, req_grad_Y):
            Y = None
        if not spmm_cache_argX(op, reduce_op, req_grad_X, req_grad_Y):
            argX = None
        if not spmm_cache_argY(op, reduce_op, req_grad_X, req_grad_Y):
            argY = None
        ctx.save_for_backward(X, Y, argX, argY)
        return out

    @staticmethod
    def backward(ctx, dZ):
        (
            gidx,
            op,
            reduce_op,
            X_shape,
            Y_shape,
            dtype,
            device,
            reduce_last,
        ) = ctx.backward_cache
        X, Y, argX, argY = ctx.saved_tensors
        if op != "copy_rhs" and ctx.needs_input_grad[3]:
            g_rev = gidx.reverse()
            if reduce_op == "sum":
                if op == "mul":
                    dX = gspmm(g_rev, "mul", "sum", dZ, Y)
                elif op == "add":
                    dX = gspmm(g_rev, "copy_lhs", "sum", dZ, Y)
                elif op == "copy_lhs":
                    dX = gspmm(g_rev, "copy_lhs", "sum", dZ, None)
            else:  # max/min
                dX = th.zeros(
                    (X_shape[0],) + dZ.shape[1:], dtype=dtype, device=device
                )
                if op == "mul":
                    grad = _expand(Y, dZ.shape[1:]).gather(0, argY.long()) * dZ
                    dX.scatter_add_(0, argX.long(), grad)
                elif op in ["add", "copy_lhs"]:
                    dX.scatter_add_(0, argX.long(), dZ)
            dX = _reduce_grad(dX, X_shape)
        else:  # X has not gradient
            dX = None
        if op != "copy_lhs" and ctx.needs_input_grad[4]:
            if reduce_op == "sum":
                if op == "mul" and reduce_last:
                    dY = gsddmm(gidx, "dot", X, dZ)
                elif op == "mul":
                    dY = gsddmm(gidx, "mul", X, dZ)
                elif op in ["add", "copy_rhs"]:
                    dY = gsddmm(gidx, "copy_rhs", X, dZ)
            else:  # max/min
                dY = th.zeros(
                    (Y_shape[0],) + dZ.shape[1:], dtype=dtype, device=device
                )
                if op == "mul":
                    grad = _expand(X, dZ.shape[1:]).gather(0, argX.long()) * dZ
                    dY.scatter_add_(0, argY.long(), grad)
                elif op in ["add", "copy_rhs"]:
                    dY.scatter_add_(0, argY.long(), dZ)
            dY = _reduce_grad(dY, Y_shape)
        else:  # Y has no gradient
            dY = None
        return None, None, None, dX, dY


class GSpMM_hetero(th.autograd.Function):
    @staticmethod
    def forward(
        ctx, gidx, op, reduce_op, X_len, *feats
    ):  # feats = lhs_data + rhs_data
        out, (argX, argY, argX_ntype, argY_etype) = _gspmm_hetero(
            gidx, op, reduce_op, X_len, feats
        )
        X, Y = feats[:X_len], feats[X_len:]
        # TODO (Israt): check target to decide src_id/dst_id?
        src_id, dst_id = gidx.metagraph.find_edge(0)
        reduce_last = _need_reduce_last_dim(X[src_id], Y[dst_id])
        X_shape = tuple(
            [X[i].shape if X[i] is not None else None for i in range(X_len)]
        )
        Y_shape = tuple(
            [Y[i].shape if Y[i] is not None else None for i in range(len(Y))]
        )
        dtype = X[src_id].dtype if X[src_id] is not None else Y[dst_id].dtype
        device = X[src_id].device if X[src_id] is not None else Y[dst_id].device
        ctx.backward_cache = (
            gidx,
            op,
            reduce_op,
            X_shape,
            Y_shape,
            dtype,
            device,
            reduce_last,
            X_len,
        )
        req_grad_X = tuple(
            [
                X[i].requires_grad if X[i] is not None else False
                for i in range(X_len)
            ]
        )
        req_grad_Y = tuple(
            [
                Y[i].requires_grad if Y[i] is not None else False
                for i in range(len(Y))
            ]
        )

        # checking the first relation to decide for all the relations
        if not spmm_cache_argX(
            op, reduce_op, req_grad_X[src_id], req_grad_Y[dst_id]
        ):
            argX = tuple([None] * len(X))
        if not spmm_cache_argY(
            op, reduce_op, req_grad_X[src_id], req_grad_Y[dst_id]
        ):
            argY = tuple([None] * len(X))

        ctx.save_for_backward(*feats, *argX, *argX_ntype, *argY, *argY_etype)
        return out

    @staticmethod
    def backward(ctx, *dZ):
        (
            gidx,
            op,
            reduce_op,
            X_shape,
            Y_shape,
            dtype,
            device,
            reduce_last,
            X_len,
        ) = ctx.backward_cache
        num_ntypes = gidx.number_of_ntypes()
        feats = ctx.saved_tensors[: -(4 * num_ntypes)]
        argX = ctx.saved_tensors[-(4 * num_ntypes) : -(3 * num_ntypes)]
        argX_ntype = ctx.saved_tensors[-(3 * num_ntypes) : -(2 * num_ntypes)]
        argY = ctx.saved_tensors[-(2 * num_ntypes) : -num_ntypes]
        argY_etype = ctx.saved_tensors[-num_ntypes:]
        X, Y = feats[:X_len], feats[X_len:]

        if op != "copy_rhs" and any([x is not None for x in X]):
            g_rev = gidx.reverse()
            if reduce_op == "sum":
                if op == "mul":
                    dX = gspmm_hetero(
                        g_rev, "mul", "sum", len(X), *tuple(dZ + Y)
                    )
                elif op == "add":
                    dX = gspmm_hetero(
                        g_rev, "copy_lhs", "sum", len(X), *tuple(dZ + Y)
                    )
                elif op == "copy_lhs":
                    tpl_None = tuple([None] * len(Y))
                    dX = gspmm_hetero(
                        g_rev, "copy_lhs", "sum", len(X), *tuple(dZ + tpl_None)
                    )
            else:  # max/min
                # Assuming that the features are of the same dimension (enforced by the forward function)
                src_id, dst_id = gidx.metagraph.find_edge(0)
                dX = tuple(
                    [
                        th.zeros(
                            (X_shape[i][0],) + dZ[dst_id].shape[1:],
                            dtype=dtype,
                            device=device,
                        )
                        if X[i] is not None
                        else None
                        for i in range(len(X))
                    ]
                )
                if op == "mul":
                    grad = _expand(Y, dZ.shape[1:]).gather(0, argY.long()) * dZ
                    dX.scatter_add_(0, argX.long(), grad)
                elif op in ["add", "copy_lhs"]:
                    dX = _update_grad_minmax_hetero(
                        g_rev, op, dZ, argX, argX_ntype, dX
                    )
            dX = tuple(
                [
                    _reduce_grad(dX[i], X_shape[i])
                    if X[i] is not None
                    else None
                    for i in range(len(X))
                ]
            )
        else:  # X has not gradient
            dX = tuple([None] * len(X))
        if op != "copy_lhs" and any([y is not None for y in Y]):
            # TODO(Israt): implement other combinations of reduce functions
            if reduce_op == "sum":
                tpl_dZ = tuple(
                    [
                        dZ[i] if dZ[i] is not None else None
                        for i in range(len(dZ))
                    ]
                )
                tpl_X_dZ = tuple(X + tpl_dZ)
                if op == "mul" and reduce_last:
                    dY = gsddmm_hetero(gidx, "dot", X_len, "u", "v", *tpl_X_dZ)
                elif op == "mul":
                    dY = gsddmm_hetero(gidx, "mul", X_len, "u", "v", *tpl_X_dZ)
                elif op in ["add", "copy_rhs"]:
                    dY = gsddmm_hetero(
                        gidx, "copy_rhs", X_len, "u", "v", *tpl_X_dZ
                    )
            else:  # max/min
                src_id, dst_id = gidx.metagraph.find_edge(0)
                dY = tuple(
                    [
                        th.zeros(
                            (Y_shape[i][0],) + dZ[dst_id].shape[1:],
                            dtype=dtype,
                            device=device,
                        )
                        if Y[i] is not None
                        else None
                        for i in range(len(Y))
                    ]
                )
                if op == "mul":
                    grad = _expand(X, dZ.shape[1:]).gather(0, argX.long()) * dZ
                    dY.scatter_add_(0, argY.long(), grad)
                elif op in ["add", "copy_rhs"]:
                    dY = _update_grad_minmax_hetero(
                        gidx.reverse(), op, dZ, argY, argY_etype, dY
                    )
            dY = tuple(
                [
                    _reduce_grad(dY[i], Y_shape[i])
                    if dY[i] is not None
                    else None
                    for i in range(len(dY))
                ]
            )
        else:  # Y has no gradient
            dY = tuple([None] * len(Y))
        return (None, None, None, None) + dX + dY


def sddmm_cache_X(op, req_grad_X, req_grad_Y):
    """Rules to identify whether to cache X in SDDMM forward stage."""
    if op in ["mul", "dot"] and req_grad_Y:
        return True
    return False


def sddmm_cache_Y(op, req_grad_X, req_grad_Y):
    """Rules to identify whether to cache Y in SDDMM forward stage."""
    if op in ["mul", "dot"] and req_grad_X:
        return True
    return False


class GSDDMM(th.autograd.Function):
    @staticmethod
    def forward(ctx, gidx, op, X, Y, lhs_target, rhs_target):
        out = _gsddmm(gidx, op, X, Y, lhs_target, rhs_target)
        X_shape = X.shape if X is not None else None
        Y_shape = Y.shape if Y is not None else None
        ctx.backward_cache = gidx, op, lhs_target, rhs_target, X_shape, Y_shape
        req_grad_X = X.requires_grad if X is not None else False
        req_grad_Y = Y.requires_grad if Y is not None else False
        if not sddmm_cache_X(op, req_grad_X, req_grad_Y):
            X = None
        if not sddmm_cache_Y(op, req_grad_X, req_grad_Y):
            Y = None
        ctx.save_for_backward(X, Y)
        return out

    @staticmethod
    def backward(ctx, dZ):
        gidx, op, lhs_target, rhs_target, X_shape, Y_shape = ctx.backward_cache
        X, Y = ctx.saved_tensors
        if op != "copy_rhs" and ctx.needs_input_grad[2]:
            if lhs_target in ["u", "v"]:
                _gidx = gidx if lhs_target == "v" else gidx.reverse()
                if op in ["add", "copy_lhs"]:
                    dX = gspmm(_gidx, "copy_rhs", "sum", None, dZ)
                else:  # mul, dot
                    if rhs_target == lhs_target:
                        dX = gspmm(_gidx, "copy_rhs", "sum", None, dZ) * Y
                    elif rhs_target == "e":
                        dX = gspmm(_gidx, "copy_rhs", "sum", None, dZ * Y)
                    else:  # rhs_target = !lhs_target
                        dX = gspmm(_gidx, "mul", "sum", Y, dZ)
            else:  # lhs_target == 'e'
                if op in ["add", "copy_lhs"]:
                    dX = dZ
                else:  # mul, dot
                    dX = gsddmm(gidx, "mul", dZ, Y, "e", rhs_target)
            dX = _reduce_grad(dX, X_shape)
        else:
            dX = None
        if op != "copy_lhs" and ctx.needs_input_grad[3]:
            if rhs_target in ["u", "v"]:
                _gidx = gidx if rhs_target == "v" else gidx.reverse()
                if op in ["add", "copy_rhs"]:
                    dY = gspmm(_gidx, "copy_rhs", "sum", None, dZ)
                else:  # mul, dot
                    if lhs_target == rhs_target:
                        dY = gspmm(_gidx, "copy_rhs", "sum", None, dZ) * X
                    elif lhs_target == "e":
                        dY = gspmm(_gidx, "copy_rhs", "sum", None, dZ * X)
                    else:  # rhs_target = !lhs_target
                        dY = gspmm(_gidx, "mul", "sum", X, dZ)
            else:
                if op in ["add", "copy_rhs"]:
                    dY = dZ
                else:  # mul, dot
                    dY = gsddmm(gidx, "mul", dZ, X, "e", lhs_target)
            dY = _reduce_grad(dY, Y_shape)
        else:
            dY = None
        return None, None, dX, dY, None, None


class GSDDMM_hetero(th.autograd.Function):
    @staticmethod
    def forward(
        ctx, gidx, op, X_len, lhs_target, rhs_target, *feats
    ):  # feats = X+Y
        out = _gsddmm_hetero(gidx, op, X_len, lhs_target, rhs_target, feats)
        X, Y = feats[:X_len], feats[X_len:]
        X_shape = tuple(
            [X[i].shape if X[i] is not None else None for i in range(len(X))]
        )
        Y_shape = tuple(
            [Y[i].shape if Y[i] is not None else None for i in range(len(Y))]
        )
        ctx.backward_cache = (
            gidx,
            op,
            lhs_target,
            rhs_target,
            X_shape,
            Y_shape,
            X_len,
        )
        req_grad_X = tuple(
            [
                X[i].requires_grad if X[i] is not None else False
                for i in range(len(X))
            ]
        )
        req_grad_Y = tuple(
            [
                Y[i].requires_grad if Y[i] is not None else False
                for i in range(len(Y))
            ]
        )
        ctx.save_for_backward(*feats)
        return out

    @staticmethod
    # TODO(Israt): Implement the complete backward operator
    def backward(ctx, *dZ):
        (
            gidx,
            op,
            lhs_target,
            rhs_target,
            X_shape,
            Y_shape,
            X_len,
        ) = ctx.backward_cache
        feats = ctx.saved_tensors
        X, Y = feats[:X_len], feats[X_len:]
        if op != "copy_rhs" and any([x is not None for x in X]):
            if lhs_target in ["u", "v"]:
                _gidx = gidx if lhs_target == "v" else gidx.reverse()
                tpl_of_None = tuple([None] * len(X))
                if op in ["add", "copy_lhs"]:
                    dX = gspmm_hetero(
                        _gidx,
                        "copy_rhs",
                        "sum",
                        len(X),
                        *(tuple(tpl_of_None + dZ))
                    )
                else:  # mul, dot
                    if rhs_target == lhs_target:
                        dX = (
                            gspmm_hetero(
                                _gidx,
                                "copy_rhs",
                                "sum",
                                len(X),
                                *(tuple(tpl_of_None + dZ))
                            )
                            * Y
                        )
                    elif rhs_target == "e":
                        dZ_mul_Y = tuple(
                            [
                                dZ[i] * Y[i] if dZ[i] is not None else None
                                for i in range(len(Y))
                            ]
                        )
                        dX = gspmm_hetero(
                            _gidx,
                            "copy_rhs",
                            "sum",
                            len(X),
                            *(tuple(tpl_of_None + dZ_mul_Y))
                        )
                    else:  # rhs_target = !lhs_target
                        dX = gspmm_hetero(
                            _gidx, "mul", "sum", len(X), *tuple(Y + dZ)
                        )
            else:  # lhs_target == 'e'
                if op in ["add", "copy_lhs"]:
                    dX = dZ
                else:  # mul, dot
                    num_etype = gidx.number_of_etypes()
                    dX = gsddmm_hetero(
                        gidx, "mul", num_etype, "e", rhs_target, *tuple(dZ + Y)
                    )
            dX = tuple(
                [
                    _reduce_grad(dX[i], X_shape[i])
                    if X[i] is not None
                    else None
                    for i in range(len(X))
                ]
            )
        else:
            dX = tuple([None] * len(X))
        if op != "copy_lhs" and any([y is not None for y in Y]):
            if rhs_target in ["u", "v"]:
                _gidx = gidx if rhs_target == "v" else gidx.reverse()
                tpl_of_None = tuple([None] * len(X))
                if op in ["add", "copy_rhs"]:
                    dY = gspmm_hetero(
                        _gidx,
                        "copy_rhs",
                        "sum",
                        len(X),
                        *(tuple(tpl_of_None + dZ))
                    )
                else:  # mul, dot
                    if lhs_target == rhs_target:
                        dY = (
                            gspmm_hetero(
                                _gidx,
                                "copy_rhs",
                                "sum",
                                len(X),
                                *(tuple(tpl_of_None + dZ))
                            )
                            * X
                        )
                    elif lhs_target == "e":
                        dZ_mul_X = tuple(
                            [
                                dZ[i] * X[i] if dZ[i] is not None else None
                                for i in range(len(X))
                            ]
                        )
                        dY = gspmm_hetero(
                            _gidx,
                            "copy_rhs",
                            "sum",
                            len(X),
                            *(tuple(tpl_of_None + dZ_mul_X))
                        )
                    else:  # rhs_target = !lhs_target
                        dY = gspmm_hetero(
                            _gidx, "mul", "sum", len(X), *tuple(X + dZ)
                        )
            else:
                if op in ["add", "copy_rhs"]:
                    dY = tuple(
                        [
                            dZ[i] if dZ[i] is not None else None
                            for i in range(len(dZ))
                        ]
                    )
                else:  # mul, dot
                    num_etype = gidx.number_of_etypes()
                    dY = gsddmm_hetero(
                        gidx, "mul", num_etype, "e", lhs_target, *tuple(dZ + X)
                    )
            dY = tuple(
                [
                    _reduce_grad(dY[i], Y_shape[i])
                    if Y[i] is not None
                    else None
                    for i in range(len(Y))
                ]
            )
        else:
            dY = tuple([None] * len(Y))
        return (None, None, None, None, None) + dX + dY


class EdgeSoftmax(th.autograd.Function):
    @staticmethod
    def forward(ctx, gidx, score, eids, norm_by):
        """Forward function.

        Pseudo-code:

        .. code:: python

            score = dgl.EData(g, score)
            score_max = score.dst_max()  # of type dgl.NData
            score = score - score_max  # edge_sub_dst, ret dgl.EData
            score_sum = score.dst_sum()  # of type dgl.NData
            out = score / score_sum    # edge_div_dst, ret dgl.EData
            return out.data
        """
        # remember to save the graph to backward cache before making it
        # a local variable
        if not is_all(eids):
            gidx = gidx.edge_subgraph([eids], True).graph
        if norm_by == "src":
            gidx = gidx.reverse()
        # Note: Now _edge_softmax_forward op only supports CPU
        # TODO(Zhejiang): We will support GPU in the future
        if score.is_cuda:
            score_max = _gspmm(gidx, "copy_rhs", "max", None, score)[0]
            score = th.exp(_gsddmm(gidx, "sub", score, score_max, "e", "v"))
            score_sum = _gspmm(gidx, "copy_rhs", "sum", None, score)[0]
            out = _gsddmm(gidx, "div", score, score_sum, "e", "v")
        else:
            out = _edge_softmax_forward(gidx, score, "copy_rhs")
        ctx.backward_cache = gidx
        ctx.save_for_backward(out)
        return out

    @staticmethod
    def backward(ctx, grad_out):
        """Backward function.

        Pseudo-code:

        .. code:: python

            g, out = ctx.backward_cache
            grad_out = dgl.EData(g, grad_out)
            out = dgl.EData(g, out)
            sds = out * grad_out  # type dgl.EData
            sds_sum = sds.dst_sum()  # type dgl.NData
            grad_score = sds - out * sds_sum  # multiple expressions
            return grad_score.data
        """
        gidx = ctx.backward_cache
        (out,) = ctx.saved_tensors
        sds = out * grad_out
        # Note: Now _edge_softmax_backward op only supports CPU
        # TODO(Zhejiang): We will support GPU in the future
        if out.is_cuda:
            accum = gspmm(gidx, "copy_rhs", "sum", None, sds)

            grad_score = sds - gsddmm(gidx, "mul", out, accum, "e", "v")
        else:
            grad_score = _edge_softmax_backward(gidx, out, sds)
        return None, grad_score, None, None


class EdgeSoftmax_hetero(th.autograd.Function):
    @staticmethod
    def forward(ctx, gidx, eids, norm_by, *score):
        """Forward function.

        Pseudo-code:

        .. code:: python

            score = dgl.EData(g, score)
            score_max = score.dst_max()  # of type dgl.NData
            score = score - score_max  # edge_sub_dst, ret dgl.EData
            score_sum = score.dst_sum()  # of type dgl.NData
            out = score / score_sum    # edge_div_dst, ret dgl.EData
            return out.data
        """
        # remember to save the graph to backward cache before making it
        # a local variable
        if not is_all(eids):
            gidx = gidx.edge_subgraph([eids], True).graph
        if norm_by == "src":
            gidx = gidx.reverse()
        u_len = gidx.number_of_ntypes()
        e_len = gidx.number_of_etypes()
        lhs = [None] * u_len
        feats = tuple(lhs + list(score))
        score_max = _gspmm_hetero(gidx, "copy_rhs", "max", u_len, feats)[0]
        out_tmp = _gsddmm_hetero(
            gidx, "sub", e_len, "e", "v", tuple(list(score) + list(score_max))
        )
        score = tuple(
            [
                th.exp(out_tmp[i]) if out_tmp[i] is not None else None
                for i in range(len(out_tmp))
            ]
        )
        score_sum = _gspmm_hetero(
            gidx, "copy_rhs", "sum", u_len, tuple(lhs + list(score))
        )[0]
        out = _gsddmm_hetero(
            gidx, "div", e_len, "e", "v", tuple(list(score) + list(score_sum))
        )
        ctx.backward_cache = gidx
        ctx.save_for_backward(*out)
        return out

    @staticmethod
    def backward(ctx, *grad_out):
        """Backward function.

        Pseudo-code:

        .. code:: python

            g, out = ctx.backward_cache
            grad_out = dgl.EData(g, grad_out)
            out = dgl.EData(g, out)
            sds = out * grad_out  # type dgl.EData
            sds_sum = sds.dst_sum()  # type dgl.NData
            grad_score = sds - out * sds_sum  # multiple expressions
            return grad_score.data
        """
        gidx = ctx.backward_cache
        u_len = gidx.number_of_ntypes()
        e_len = gidx.number_of_etypes()
        lhs = [None] * u_len
        out = ctx.saved_tensors
        sds = tuple([out[i] * grad_out[i] for i in range(len(out))])
        accum = _gspmm_hetero(
            gidx, "copy_rhs", "sum", u_len, tuple(lhs + list(sds))
        )[0]
        out_sddmm = _gsddmm_hetero(
            gidx, "mul", e_len, "e", "v", tuple(list(out) + list(accum))
        )
        grad_score = tuple([sds[i] - out_sddmm[i] for i in range(len(sds))])
        return (None, None, None) + grad_score


class SegmentReduce(th.autograd.Function):
    @staticmethod
    def forward(ctx, op, x, offsets):
        y, arg = _segment_reduce(op, x, offsets)
        ctx.save_for_backward(arg, offsets)
        ctx.backward_cache = op
        return y

    @staticmethod
    def backward(ctx, dy):
        op = ctx.backward_cache
        arg, offsets = ctx.saved_tensors
        m = offsets[-1].item()
        if op == "sum":
            offsets = offsets[1:]
            # To address the issue of trailing zeros, related issue:
            # https://github.com/dmlc/dgl/pull/2610
            indices = th.zeros(
                (m + 1,), device=offsets.device, dtype=offsets.dtype
            )
            indices.scatter_add_(0, offsets, th.ones_like(offsets))
            indices = th.cumsum(indices, -1)[:-1]
            dx = dy[indices]
        else:
            dx = _bwd_segment_cmp(dy, arg, m)
        return None, dx, None


class ScatterAdd(th.autograd.Function):
    @staticmethod
    def forward(ctx, x, idx, m):
        y = _scatter_add(x, idx, m)
        ctx.save_for_backward(idx)
        return y

    @staticmethod
    def backward(ctx, dy):
        idx = ctx.saved_tensors
        return dy[idx], None, None


class CSRMM(th.autograd.Function):
    @staticmethod
    def forward(ctx, gidxA, A_weights, gidxB, B_weights, num_vtypes):
        gidxC, C_weights = _csrmm(
            gidxA, A_weights, gidxB, B_weights, num_vtypes
        )
        (
            nrows,
            ncols,
            C_indptr,
            C_indices,
            C_eids,
        ) = gidxC.adjacency_matrix_tensors(0, False, "csr")
        # Note: the returned C_indptr, C_indices and C_eids tensors MUST be the same
        # as the underlying tensors of the created graph gidxC.
        ctx.backward_cache = gidxA, gidxB, gidxC
        ctx.save_for_backward(A_weights, B_weights)
        return (
            th.tensor(nrows),
            th.tensor(ncols),
            C_indptr,
            C_indices,
            C_eids,
            C_weights,
        )

    @staticmethod
    def backward(
        ctx, dnrows, dncols, dC_indptr, dC_indices, dC_eids, dC_weights
    ):
        # Only the last argument is meaningful.
        gidxA, gidxB, gidxC = ctx.backward_cache
        A_weights, B_weights = ctx.saved_tensors
        dgidxA, dA_weights = csrmm(
            gidxC,
            dC_weights,
            gidxB.reverse(),
            B_weights,
            gidxA.number_of_ntypes(),
        )
        dgidxB, dB_weights = csrmm(
            gidxA.reverse(),
            A_weights,
            gidxC,
            dC_weights,
            gidxB.number_of_ntypes(),
        )
        dA_weights = csrmask(dgidxA, dA_weights, gidxA)
        dB_weights = csrmask(dgidxB, dB_weights, gidxB)
        return None, dA_weights, None, dB_weights, None


class CSRSum(th.autograd.Function):
    @staticmethod
    def forward(ctx, gidxs, *weights):
        # PyTorch tensors must be explicit arguments of the forward function
        gidxC, C_weights = _csrsum(gidxs, weights)
        (
            nrows,
            ncols,
            C_indptr,
            C_indices,
            C_eids,
        ) = gidxC.adjacency_matrix_tensors(0, False, "csr")
        # Note: the returned C_indptr, C_indices and C_eids tensors MUST be the same
        # as the underlying tensors of the created graph gidxC.
        ctx.backward_cache = gidxs, gidxC
        return (
            th.tensor(nrows),
            th.tensor(ncols),
            C_indptr,
            C_indices,
            C_eids,
            C_weights,
        )

    @staticmethod
    def backward(
        ctx, dnrows, dncols, dC_indptr, dC_indices, dC_eids, dC_weights
    ):
        # Only the last argument is meaningful.
        gidxs, gidxC = ctx.backward_cache
        return (None,) + tuple(
            csrmask(gidxC, dC_weights, gidx) for gidx in gidxs
        )


class CSRMask(th.autograd.Function):
    @staticmethod
    def forward(ctx, gidxA, A_weights, gidxB):
        ctx.backward_cache = gidxA, gidxB
        return _csrmask(gidxA, A_weights, gidxB)

    @staticmethod
    def backward(ctx, dB_weights):
        gidxA, gidxB = ctx.backward_cache
        return None, csrmask(gidxB, dB_weights, gidxA), None


class SEGMENTMM(th.autograd.Function):
    @staticmethod
    def forward(ctx, A, B, seglen_A):
        if B.dim() != 3:
            raise ValueError("segment_mm expects B to be a 3D tensor.")
        C = th.empty((A.shape[0], B.shape[2]), device=A.device, dtype=A.dtype)
        C = _segment_mm(A, B, C, seglen_A)
        ctx.backward_cache = A, B, seglen_A
        return C

    @staticmethod
    def backward(ctx, dZ):
        A, B, seglen_A = ctx.backward_cache
        A_grad = B_grad = None
        if ctx.needs_input_grad[0]:
            #  Compute A_grad = Out_grad * B^T
            A_grad = th.empty(A.shape, device=A.device, dtype=A.dtype)
            A_grad = _segment_mm(dZ, B, A_grad, seglen_A, b_trans=True)
        if ctx.needs_input_grad[1]:
            #  Compute B_grad = A^T * Out_grad
            B_grad = th.empty(B.shape, device=B.device, dtype=B.dtype)
            B_grad = _segment_mm_backward_B(A, dZ, B_grad, seglen_A)
        return A_grad, B_grad, None


class GATHERMM(th.autograd.Function):
    @staticmethod
    def forward(ctx, A, B, idx_a, idx_b):
        if B.dim() != 3:
            raise ValueError(
                "Expected dimension of B is 3. Got " + str(B.dim())
            )
        N = len(idx_b) if idx_a is None else len(idx_a)
        C = th.zeros((N, B.shape[2]), device=A.device, dtype=A.dtype)
        C = _gather_mm(A, B, C, idx_a, idx_b)
        ctx.backward_cache = A, B, idx_a, idx_b
        return C

    @staticmethod
    def backward(ctx, dZ):
        A, B, idx_a, idx_b = ctx.backward_cache
        A_grad = B_grad = None
        if ctx.needs_input_grad[0]:
            #  Compute A_grad = Out_grad * B^T
            A_grad = th.zeros(A.shape, device=A.device, dtype=A.dtype)
            A_grad = _gather_mm_scatter(
                dZ, B.transpose(1, 2), A_grad, idx_b=idx_b, idx_c=idx_a
            )
        if ctx.needs_input_grad[1]:
            #  Compute B_grad = A^T * Out_grad
            B_grad = th.zeros(B.shape, device=B.device, dtype=B.dtype)
            B_grad = _gather_mm_scatter(A, dZ, B_grad, idx_a=idx_a, idx_c=idx_b)
        return A_grad, B_grad, None, None


def gspmm(gidx, op, reduce_op, lhs_data, rhs_data):
    if op == "sub":
        op = "add"
        rhs_data = -rhs_data
    if op == "div":
        op = "mul"
        rhs_data = 1.0 / rhs_data
    args = _cast_if_autocast_enabled(gidx, op, reduce_op, lhs_data, rhs_data)
    with _disable_autocast_if_enabled():
        return GSpMM.apply(*args)


def gsddmm(gidx, op, lhs_data, rhs_data, lhs_target="u", rhs_target="v"):
    if op == "sub":
        op = "add"
        rhs_data = -rhs_data
    if op == "div":
        op = "mul"
        rhs_data = 1.0 / rhs_data
    args = _cast_if_autocast_enabled(
        gidx, op, lhs_data, rhs_data, lhs_target, rhs_target
    )
    with _disable_autocast_if_enabled():
        return GSDDMM.apply(*args)


def gspmm_hetero(g, op, reduce_op, lhs_len, *lhs_and_rhs_tuple):
    lhs_tuple, rhs_tuple = (
        lhs_and_rhs_tuple[:lhs_len],
        lhs_and_rhs_tuple[lhs_len:],
    )
    if op == "sub":
        op = "add"
        rhs_tuple = tuple(
            [
                -rhs_tuple[i] if rhs_tuple[i] is not None else None
                for i in range(len(rhs_tuple))
            ]
        )
    if op == "div":
        op = "mul"
        rhs_tuple = tuple(
            [
                (1.0 / rhs_tuple[i]) if rhs_tuple[i] is not None else None
                for i in range(len(rhs_tuple))
            ]
        )
    if op in ["add", "mul"]:
        lhs_and_rhs_tuple = tuple(list(lhs_tuple) + list(rhs_tuple))

    args = _cast_if_autocast_enabled(
        g, op, reduce_op, lhs_len, *lhs_and_rhs_tuple
    )
    with _disable_autocast_if_enabled():
        return GSpMM_hetero.apply(*args)


def gsddmm_hetero(
    g, op, lhs_len, lhs_target="u", rhs_target="v", *lhs_and_rhs_tuple
):
    lhs_tuple, rhs_tuple = (
        lhs_and_rhs_tuple[:lhs_len],
        lhs_and_rhs_tuple[lhs_len:],
    )
    if op == "sub":
        op = "add"
        rhs_tuple = tuple(
            [
                -rhs_tuple[i] if rhs_tuple[i] is not None else None
                for i in range(len(rhs_tuple))
            ]
        )
    if op == "div":
        op = "mul"
        rhs_tuple = tuple(
            [
                (1.0 / rhs_tuple[i]) if rhs_tuple[i] is not None else None
                for i in range(len(rhs_tuple))
            ]
        )
    if op in ["add", "mul"]:
        lhs_and_rhs_tuple = tuple(list(lhs_tuple) + list(rhs_tuple))

    args = _cast_if_autocast_enabled(
        g, op, lhs_len, lhs_target, rhs_target, *lhs_and_rhs_tuple
    )
    with _disable_autocast_if_enabled():
        return GSDDMM_hetero.apply(*args)


def edge_softmax(gidx, logits, eids=ALL, norm_by="dst"):
    args = _cast_if_autocast_enabled(gidx, logits, eids, norm_by)
    with _disable_autocast_if_enabled():
        return EdgeSoftmax.apply(*args)


def edge_softmax_hetero(gidx, eids=ALL, norm_by="dst", *logits):
    args = _cast_if_autocast_enabled(gidx, eids, norm_by, *logits)
    with _disable_autocast_if_enabled():
        return EdgeSoftmax_hetero.apply(*args)


def segment_reduce(op, x, offsets):
    args = _cast_if_autocast_enabled(op, x, offsets)
    with _disable_autocast_if_enabled():
        return SegmentReduce.apply(*args)


def scatter_add(x, idx, m):
    args = _cast_if_autocast_enabled(x, idx, m)
    with _disable_autocast_if_enabled():
        return ScatterAdd.apply(*args)


def csrmm(gidxA, A_weights, gidxB, B_weights, num_vtypes):
    nrows, ncols, C_indptr, C_indices, C_eids, C_weights = CSRMM.apply(
        gidxA, A_weights, gidxB, B_weights, num_vtypes
    )
    gidxC = create_unitgraph_from_csr(
        num_vtypes,
        nrows.item(),
        ncols.item(),
        C_indptr,
        C_indices,
        C_eids,
        ["coo", "csr", "csc"],
    )
    return gidxC, C_weights


def csrsum(gidxs, weights):
    nrows, ncols, C_indptr, C_indices, C_eids, C_weights = CSRSum.apply(
        gidxs, *weights
    )
    gidxC = create_unitgraph_from_csr(
        gidxs[0].number_of_ntypes(),
        nrows.item(),
        ncols.item(),
        C_indptr,
        C_indices,
        C_eids,
        ["coo", "csr", "csc"],
    )
    return gidxC, C_weights


def csrmask(gidxA, A_weights, gidxB):
    return CSRMask.apply(gidxA, A_weights, gidxB)


def segment_mm(A, B, seglen_A):
    if A.device.type == "cpu":
        C = []
        off = 0
        for i in range(B.shape[0]):
            C.append(A[off : off + seglen_A[i]] @ B[i])
            off += seglen_A[i]
        return th.cat(C)
    else:
        args = _cast_if_autocast_enabled(A, B, seglen_A)
        with _disable_autocast_if_enabled():
            return SEGMENTMM.apply(*args)


def gather_mm(A, B, idx_A=None, idx_B=None):
    if A.device.type == "cpu":
        A = A[idx_A] if idx_A is not None else A
        B = B[idx_B] if idx_B is not None else B
        return th.bmm(A.unsqueeze(1), B).squeeze(1)
    else:
        args = _cast_if_autocast_enabled(A, B, idx_A, idx_B)
        with _disable_autocast_if_enabled():
            return GATHERMM.apply(*args)