# Copyright (c) 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.

from typing import Optional, Tuple, Union

import torch
from torch.autograd.function import once_differentiable

from ..utils import maybe_size, maybe_empty, maybe_detach, get_tensor_meta_data
from ..graph import CSC
from pylibcugraphops.operators import (
    update_efeat_e2e_sum_fwd,
    update_efeat_e2e_sum_bwd,
    update_efeat_e2e_concat_fwd,
    update_efeat_e2e_concat_bwd,
)


def update_efeat_e2e(
    edge_feat: Optional[torch.Tensor],
    node_feat: Union[torch.Tensor, Tuple[Optional[torch.Tensor]]],
    graph: CSC,
    use_source_emb: bool,
    use_target_emb: bool,
    mode: str = "concat",
) -> torch.Tensor:
    """
    PyTorch autograd function for creating new edge features (update_efeat)
    based on either concatenating or summing edge features and the features of
    the corresponding source and destination node of each edge in an edge-to-edge
    fashion (e2e).

    Parameters
    ----------
    edge_feat : torch.Tensor | None
        The input edge features - ignored if passed as ``None``.
        Shape: ``(num_indices, dim_edge)``.

    node_feat : torch.Tensor | Tuple[torch.Tensor] | None
        The input source features. For bipartite graphs, a tuple of
        source and destionations node features with shapes
        ``(num_src_nodes, dim_src)`` and ``(num_dst_nodes, dim_dst)``
        respectively is expected. For non-bipartite graph, a single
        tensor of node features of shape ``(num_dst_nodes, dim_node)``
        is expected with ``dim_src = dim_dst = dim_node``.

    use_source_emb : bool
        whether to use node embeddings indexed by indices of source nodes

    use_target_emb : bool
        whether to use node embeddings indexed by indices of source nodes

    graph : CSC
        The graph used for the operation.

    mode : str, default="concat"
        Mode of update, either ``"concat"`` or ``"sum"``.

    Returns
    -------
    output : torch.Tensor
        The output after either concatenation or summation.
        Shape for ``mode=concat``: ``(num_indices, dim_edge + dim_src + dim_dst)``.
        Shape for ``mode=sum``: ``(num_indices, dim_edge)``.
    """
    modes = ["sum", "concat"]
    if mode.lower() not in modes:
        raise ValueError(f"Supported operations: {modes}, " f"but got '{mode}'.")

    if isinstance(node_feat, Tuple) and len(node_feat) == 2:
        assert (
            graph.is_bipartite
        ), "Expected bipartite graph, but got graph.is_bipartite == False."
        return _update_efeat_e2e_bipartite_autograd.apply(
            edge_feat,
            node_feat[0],
            node_feat[1],
            (graph, mode, use_source_emb, use_target_emb),
        )

    else:
        assert isinstance(node_feat, torch.Tensor) or (node_feat is None)
        assert (
            not graph.is_bipartite
        ), "Expected non-bipartite graph, but got graph.is_bipartite == True."
        return _update_efeat_e2e_autograd.apply(
            edge_feat, node_feat, (graph, mode, use_source_emb, use_target_emb)
        )


class _update_efeat_e2e_autograd(torch.autograd.Function):
    """
    Custom autograd function for an update of the edge features based on
    concatenating or summing edge features, and source node and destination
    node features
    """

    @staticmethod
    def forward(ctx, edge_feat, node_feat, args):
        graph, mode, use_source_emb, use_target_emb = args
        ctx.mode = mode
        ctx.graph = graph

        edge_feat = maybe_detach(edge_feat)
        node_feat = maybe_detach(node_feat)

        dim_edge = maybe_size(edge_feat)
        dim_node = maybe_size(node_feat)

        ctx.dim_edge = dim_edge
        ctx.dim_node = dim_node
        ctx.use_source_emb = use_source_emb
        ctx.use_target_emb = use_target_emb

        if dim_edge > 0:
            dtype = edge_feat.dtype
            device = edge_feat.device

        elif dim_node > 0:
            dtype = node_feat.dtype
            device = node_feat.device

        else:
            raise ValueError("forward called with all the inputs being None")

        stream = torch.cuda.current_stream().cuda_stream

        if mode.lower() == "concat":
            dim_src = dim_node if use_source_emb else 0
            dim_dst = dim_node if use_target_emb else 0
            output = torch.empty(
                graph.num_indices,
                dim_edge + dim_dst + dim_src,
                dtype=dtype,
                device=device,
            )
            update_efeat_e2e_concat_fwd(
                output,
                edge_feat,
                node_feat,
                graph.graph_csc,
                use_source_emb,
                use_target_emb,
                stream,
            )

        else:
            # all should be equal, so find the one which is not 0
            # check for equality is done in binding
            dim_edge = max(dim_edge, dim_node)
            output = torch.empty(
                graph.num_indices, dim_edge, dtype=dtype, device=device
            )
            update_efeat_e2e_sum_fwd(
                output,
                edge_feat,
                node_feat,
                graph.graph_csc,
                use_source_emb,
                use_target_emb,
                stream,
            )

        ctx.node_feat_meta_data = get_tensor_meta_data(node_feat)
        ctx.edge_feat_meta_data = get_tensor_meta_data(edge_feat)

        return output

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_output):
        grad_output = grad_output.detach().contiguous()
        needs_grad_edge, needs_grad_node, _ = ctx.needs_input_grad

        grad_edge_feat = maybe_empty(
            **ctx.edge_feat_meta_data,
            create_tensor=needs_grad_edge,
        )
        grad_node_feat = maybe_empty(
            **ctx.node_feat_meta_data,
            create_tensor=needs_grad_node,
        )

        stream = torch.cuda.current_stream().cuda_stream
        grad_output = grad_output.detach().contiguous()

        if ctx.mode.lower() == "concat":
            update_efeat_e2e_concat_bwd(
                grad_edge_feat,
                grad_node_feat,
                grad_output,
                ctx.graph.graph_csc,
                ctx.dim_edge,
                ctx.dim_node,
                ctx.use_source_emb,
                ctx.use_target_emb,
                stream,
            )

        else:
            update_efeat_e2e_sum_bwd(
                grad_edge_feat,
                grad_node_feat,
                grad_output,
                ctx.graph.graph_csc,
                ctx.use_source_emb,
                ctx.use_target_emb,
                stream,
            )

        return grad_edge_feat, grad_node_feat, None


class _update_efeat_e2e_bipartite_autograd(torch.autograd.Function):
    """
    Custom autograd function for an update of the edge features based on
    concatenating or summing edge features, and source node and destination
    node features
    """

    @staticmethod
    def forward(ctx, edge_feat, src_node_feat, dst_node_feat, args):
        graph, mode, use_source_emb, use_target_emb = args
        assert (src_node_feat is not None) == use_source_emb
        assert (dst_node_feat is not None) == use_target_emb
        ctx.graph = graph
        ctx.mode = mode

        edge_feat = maybe_detach(edge_feat)
        src_node_feat = maybe_detach(src_node_feat)
        dst_node_feat = maybe_detach(dst_node_feat)

        dim_edge = maybe_size(edge_feat)
        dim_src = maybe_size(src_node_feat)
        dim_dst = maybe_size(dst_node_feat)

        ctx.dim_edge = dim_edge
        ctx.dim_src = dim_src
        ctx.dim_dst = dim_dst

        if dim_edge > 0:
            dtype = edge_feat.dtype
            device = edge_feat.device

        elif dim_src > 0:
            dtype = src_node_feat.dtype
            device = src_node_feat.device

        elif dim_dst > 0:
            dtype = dst_node_feat.dtype
            device = dst_node_feat.device
        else:
            raise ValueError("forward called with all the inputs being None")

        stream = torch.cuda.current_stream().cuda_stream

        if mode.lower() == "concat":
            output = torch.empty(
                graph.num_indices,
                dim_edge + dim_dst + dim_src,
                dtype=dtype,
                device=device,
            )
            update_efeat_e2e_concat_fwd(
                output,
                edge_feat,
                src_node_feat,
                dst_node_feat,
                graph.graph_csc,
                stream,
            )

        else:
            # all should be equal, so find the one which is not 0
            # check for equality is done in binding
            dim_edge = max(dim_edge, dim_src, dim_dst)
            output = torch.empty(
                graph.num_indices, dim_edge, dtype=dtype, device=device
            )
            update_efeat_e2e_sum_fwd(
                output,
                edge_feat,
                src_node_feat,
                dst_node_feat,
                graph.graph_csc,
                stream,
            )

        ctx.src_node_feat_meta_data = get_tensor_meta_data(src_node_feat)
        ctx.dst_node_feat_meta_data = get_tensor_meta_data(dst_node_feat)
        ctx.edge_feat_meta_data = get_tensor_meta_data(edge_feat)

        return output

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_output):
        grad_output = grad_output.detach().contiguous()
        needs_grad_edge, needs_grad_src, needs_grad_dst, _ = ctx.needs_input_grad

        grad_edge_feat = maybe_empty(
            **ctx.edge_feat_meta_data,
            create_tensor=needs_grad_edge,
        )
        grad_src_feat = maybe_empty(
            **ctx.src_node_feat_meta_data,
            create_tensor=needs_grad_src,
        )
        grad_dst_feat = maybe_empty(
            **ctx.dst_node_feat_meta_data,
            create_tensor=needs_grad_dst,
        )

        stream = torch.cuda.current_stream().cuda_stream
        grad_output = grad_output.detach().contiguous()

        if ctx.mode.lower() == "concat":
            update_efeat_e2e_concat_bwd(
                grad_edge_feat,
                grad_src_feat,
                grad_dst_feat,
                grad_output,
                ctx.graph.graph_csc,
                ctx.dim_edge,
                ctx.dim_src,
                ctx.dim_dst,
                stream,
            )

        else:
            update_efeat_e2e_sum_bwd(
                grad_edge_feat,
                grad_src_feat,
                grad_dst_feat,
                grad_output,
                ctx.graph.graph_csc,
                stream,
            )

        return grad_edge_feat, grad_src_feat, grad_dst_feat, None
