# 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 Any, Optional
from .utils import maybe_size, maybe_empty
import torch

from ...binding.pylibcugraphops_ext.operators import (
    update_efeat_fg_e2e_sum_fwd as sum_efeat_e2e_fwd,
    update_efeat_fg_e2e_sum_bwd as sum_efeat_e2e_bwd,
    update_efeat_fg_e2e_concat_fwd as concat_efeat_e2e_fwd,
    update_efeat_fg_e2e_concat_bwd as concat_efeat_e2e_bwd,
)


def update_efeat_e2e(
    edge_feat: Optional[torch.Tensor],
    src_node_feat: Optional[torch.Tensor],
    dst_node_feat: Optional[torch.Tensor],
    graph: Any,
    mode: str = "concat",
) -> torch.Tensor:
    r"""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: (n_edges, dim_edge).

    src_node_feat : torch.Tensor | None
        The input source node features - ignored if passed as None
        Shape: (n_in_nodes, dim_src)

    dst_node_feat : torch.Tensor | None
        The input destination node features - ignored if passed as None
        Shape: (n_out_nodes, dim_dst)

    graph : fg_csr_int[32|64] | mfg_csr_int[32|64]
        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 "concat": (n_edges, dim_edge + dim_src + dim_dst).
        Shape for "sum": (n_edges, dim_edge).
    """

    modes = ["sum", "concat"]

    use_mfg = hasattr(graph, "sample_size")
    if mode.lower() not in modes:
        raise ValueError(f"Supported operations: {modes}, " f"but got '{mode}'.")

    if use_mfg:
        raise NotImplementedError("MFG version not available.")
    return _fg_e2e_autograd.apply(
        edge_feat, src_node_feat, dst_node_feat, (graph, mode)
    )


class _fg_e2e_autograd(torch.autograd.Function):
    r"""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):
        fg, mode = args
        ctx.fg = fg
        ctx.mode = mode

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

        if dim_edge > 0:
            edge_feat = edge_feat.contiguous()
            dtype = edge_feat.dtype
            device = edge_feat.device

        elif dim_src > 0:
            src_node_feat = src_node_feat.contiguous()
            dtype = src_node_feat.dtype
            device = src_node_feat.device

        elif dim_dst > 0:
            dst_node_feat = dst_node_feat.contiguous()
            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(
                fg.n_indices, dim_edge + dim_dst + dim_src, dtype=dtype, device=device
            )
            concat_efeat_e2e_fwd(
                output,
                edge_feat.detach(),
                src_node_feat.detach(),
                dst_node_feat.detach(),
                fg,
                stream,
            )
            ctx.dim_edge = dim_edge
            ctx.dim_src = dim_src
            ctx.dim_dst = dim_dst

        else:
            output = torch.empty(fg.n_indices, dim_edge, dtype=dtype, device=device)
            sum_efeat_e2e_fwd(
                output,
                edge_feat.detach(),
                src_node_feat.detach(),
                dst_node_feat.detach(),
                fg,
                stream,
            )
            ctx.dim_edge = dim_edge
            ctx.dim_src = dim_edge
            ctx.dim_dst = dim_edge

        return output

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

        grad_edge_feat = maybe_empty(
            (ctx.fg.n_indices, ctx.dim_edge),
            dtype=grad_output.dtype,
            device=grad_output.device,
            create_tensor=needs_grad_edge,
        )
        grad_src_feat = maybe_empty(
            (ctx.fg.n_nodes, ctx.dim_src),
            dtype=grad_output.dtype,
            device=grad_output.device,
            create_tensor=needs_grad_src,
        )
        grad_dst_feat = maybe_empty(
            (ctx.fg.n_nodes, ctx.dim_dst),
            dtype=grad_output.dtype,
            device=grad_output.device,
            create_tensor=needs_grad_dst,
        )

        stream = torch.cuda.current_stream().cuda_stream

        if ctx.mode.lower() == "concat":
            concat_efeat_e2e_bwd(
                grad_edge_feat,
                grad_src_feat,
                grad_dst_feat,
                grad_output.detach(),
                ctx.fg,
                ctx.dim_edge,
                ctx.dim_src,
                ctx.dim_dst,
                stream,
            )

        else:
            sum_efeat_e2e_bwd(
                grad_edge_feat,
                grad_src_feat,
                grad_dst_feat,
                grad_output.detach(),
                ctx.fg,
                ctx.dim_edge,
                stream,
            )

        return grad_edge_feat, grad_src_feat, grad_dst_feat, None
