# Copyright (c) 2022-2023, NVIDIA CORPORATION.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import torch
from ..utils import maybe_empty, get_tensor_meta_data
from ..graph import CSC
from torch.autograd.function import once_differentiable
from pylibcugraphops.operators import (
    agg_dmpnn_e2e_fwd,
    agg_dmpnn_e2e_bwd,
)


def agg_dmpnn_e2e(
    edge_emb: torch.Tensor,
    graph: CSC,
    rev_edge_ids: torch.Tensor,
    concat_own: bool,
) -> torch.Tensor:
    r"""PyTorch autograd function for the Directed Message Passing NN (agg_dmpnn)
    primitive which performs edge-based (e2e) aggregation taking the
    direction of edges into account.

    Parameters
    ----------
    edge_emb : torch.Tensor
        The input edge features.
        Shape: (n_edges, dim).

    graph : CSC
        The graph used for the operation.

    rev_edge_ids : torch.Tensor
        the position indices of each edge for the corresponding reverse edge
        Shape: (n_indices,).

    concat_own : bool
        flag indicating whether to concat the own edge features

    Returns
    -------
    output : torch.Tensor
        The aggregation output.
        Shape: (n_edges, dim)

    """

    if graph.is_bipartite:
        raise ValueError("Operator does not support bipartite graphs.")

    return _agg_dmpnn_e2e_autograd.apply(edge_emb, graph, rev_edge_ids, concat_own)


class _agg_dmpnn_e2e_autograd(torch.autograd.Function):
    r"""Custom autograd function for the Directed Message Passing Neural Network (DMPNN)
    primitive which performs edge-based (e2e) aggregation taking the direction of edges
    into account while operating on a graph."""

    @staticmethod
    def forward(ctx, edge_emb, graph, rev_edge_ids, concat_own):
        if concat_own:
            out = torch.empty(
                (edge_emb.size(0), 2 * edge_emb.size(1)),
                dtype=edge_emb.dtype,
                device=edge_emb.device,
            )
        else:
            out = torch.empty_like(edge_emb)

        stream = torch.cuda.current_stream().cuda_stream
        edge_emb = edge_emb.detach().contiguous()
        agg_dmpnn_e2e_fwd(
            out,
            edge_emb,
            graph.graph_csc,
            rev_edge_ids,
            concat_own,
            stream_id=stream,
        )

        ctx.save_for_backward(rev_edge_ids)
        ctx.graph = graph
        ctx.concat_own = concat_own
        ctx.edge_emb_meta_data = get_tensor_meta_data(edge_emb)
        return out

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_output):
        (rev_edge_ids,) = ctx.saved_tensors
        needs_grad_emb, _, _, _ = ctx.needs_input_grad
        if not needs_grad_emb:
            return None, None, None, None

        grad_edge_emb = maybe_empty(**ctx.edge_emb_meta_data, create_tensor=True)

        stream = torch.cuda.current_stream().cuda_stream
        grad_output = grad_output.detach().contiguous()
        agg_dmpnn_e2e_bwd(
            grad_edge_emb,
            grad_output,
            ctx.graph.graph_csc,
            rev_edge_ids,
            ctx.concat_own,
            stream_id=stream,
        )
        return grad_edge_emb, None, None, None
