# Copyright (c) 2023-2024, 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 typing import Optional, Tuple, Union
from torch.autograd.function import once_differentiable

from ..utils import (
    maybe_empty_like,
    maybe_empty,
    maybe_detach,
    get_tensor_meta_data,
    maybe_zeros_like,
    maybe_zeros,
)
from ..graph import CSC
from pylibcugraphops import ActivationOp, mha_params
from pylibcugraphops.operators import (
    activation_params,
    mha_gat_v2_n2n_fwd,
    mha_gat_v2_n2n_bwd,
    mha_gat_v2_n2n_efeat_fwd,
    mha_gat_v2_n2n_efeat_bwd,
)

_activation_ops = [x for x in dir(ActivationOp) if not x.startswith("__")]


def mha_gat_v2_n2n(
    feat: Union[torch.Tensor, Tuple[torch.Tensor]],
    attn_weights: torch.Tensor,
    graph: CSC,
    num_heads: int = 1,
    activation: str = "LeakyReLU",
    negative_slope: float = 0.2,
    concat_heads: bool = True,
    edge_feat: Optional[torch.Tensor] = None,
    deterministic_dgrad: Optional[bool] = False,
    deterministic_wgrad: Optional[bool] = False,
) -> torch.Tensor:
    """
    PyTorch autograd function for a multi-head attention layer (GAT-like)
    without using cudnn (mha_gat_v2) with an activation prior to the dot
    product but none afterwards in a node-to-node reduction (n2n).

    Parameters
    ----------
    feat : torch.Tensor | Tuple[torch.Tensor, torch.Tensor]
        The input node features.
        Depending on whether the graph is bipartite or not, this has
        to be a tuple of two tensors or a single tensor correspondingly.
        For bipartite graphs, the rows of each tensors consists of
        concatenated features from different heads after the linear transformation.
        Shapes: ``(n_src_nodes, dim_in), (n_dst_nodes, dim_in)``,
        where ``dim_in = dim_head * num_heads``, with ``dim_head`` being the
        feature dimension per head.
        For non-bipartite graphs, the rows of the single tensor consists of
        concatenated features from different heads after the linear transformation.
        Shape: ``(n_src_nodes, dim_in)``, where ``dim_in = dim_head * num_heads``,
        with ``dim_head`` being the feature dimension per head.

    attn_weights : torch.Tensor
        The attention weight.
        Shape: ``(dim_in,)``.

    graph : CSC
        The graph used in for the operation.

    num_heads : int, default=1
        Number of heads in multi-head attention.

    activation : str, default="LeakyReLU"
        The activation function used in the attention mechanism. Choose from
        ``"ELU"``, ``"LeakyReLU"``, ``"Linear"``, ``"ReLU"``, ``"Scalar"``,
        ``"Sigmoid"`` and ``"Tanh"``.

    negative_slope : float, default=0.2
        LeakyReLU angle of negative slope. Only effective when ``activation`` is
        ``"LeakyReLU"``.

    concat_heads : bool, default=True
        Aggregated embeddings from each head are concatenated if ``True`` or
        averaged if ``False``.

    edge_feat : Optional[torch.Tensor], default=None
        Optional input edge features. Each row consists of concatenated features
        from different heads after the linear transformation.
        Shape: ``(n_edges, dim_in)``.

    deterministic_dgrad : Optional[bool], default=False
        Optional flag indicating whether the gradient w.r.t. the features
        is computed deterministically using a dedicated workspace buffer.

    deterministic_wgrad : Optional[bool], default=False
        Optional flag indicating whether the gradient w.r.t. the weights
        is computed deterministically using a dedicated workspace buffer.

    Returns
    -------
    output : torch.Tensor
        The aggregation output.
        Shape for ``concat_heads=True``: ``(n_dst_nodes, dim_in)``;
        Shape for ``concat_heads=False``: ``(n_dst_nodes, dim_head)``.
    """
    if activation not in _activation_ops:
        raise ValueError(
            f"Supported activation functions: {_activation_ops}, "
            f"but got '{activation}'."
        )
    act_params = activation_params(negative_slope, getattr(ActivationOp, activation))
    params = mha_params(act_params, num_heads, concat_heads)
    if graph.is_bipartite:
        assert isinstance(feat, Tuple) and len(feat) == 2, (
            f"Expected tuple of source and destination node features,"
            f" got {type(feat)} instead."
        )
        return _mha_gat_v2_n2n_bipartite_autograd.apply(
            feat[0],
            feat[1],
            edge_feat,
            attn_weights,
            graph,
            params,
            deterministic_dgrad,
            deterministic_wgrad,
        )

    assert isinstance(
        feat, torch.Tensor
    ), f"Expected torch.Tensor of node features, got {type(feat)} instead."
    return _mha_gat_v2_n2n_autograd.apply(
        feat,
        edge_feat,
        attn_weights,
        graph,
        params,
        deterministic_dgrad,
        deterministic_wgrad,
    )


class _mha_gat_v2_n2n_autograd(torch.autograd.Function):
    """
    Custom autograd function for a multi-head attention layer (GAT-like)
    without using cudnn (mha_gat_v2) with an activation prior to the dot
    product but none afterwards operating on a graph in a node-to-node
    reduction (n2n).
    """

    @staticmethod
    def forward(
        ctx,
        feat,
        edge_feat,
        attn_weights,
        graph,
        params,
        deterministic_dgrad,
        deterministic_wgrad,
    ):
        dim = feat.size(-1)
        feat = feat.detach().contiguous()
        attn_weights = attn_weights.detach().contiguous()
        edge_feat = maybe_detach(edge_feat)
        dim_out = dim if params.concat_heads else (dim // params.num_heads)
        output = torch.empty(
            graph.num_dst_nodes, dim_out, dtype=feat.dtype, device=feat.device
        )
        sm_scores = torch.empty(
            2,
            params.num_heads,
            graph.num_indices,
            dtype=feat.dtype,
            device=feat.device,
        )
        act_scores = torch.empty(
            graph.num_indices, dim, dtype=feat.dtype, device=feat.device
        )

        stream = torch.cuda.current_stream().cuda_stream

        if graph.num_dst_nodes > 0:
            if edge_feat is None:
                mha_gat_v2_n2n_fwd(
                    output,
                    sm_scores,
                    act_scores,
                    feat,
                    attn_weights,
                    graph.graph_csc,
                    params,
                    stream_id=stream,
                )
            else:
                mha_gat_v2_n2n_efeat_fwd(
                    output,
                    sm_scores,
                    act_scores,
                    feat,
                    edge_feat,
                    attn_weights,
                    graph.graph_csc,
                    params,
                    stream_id=stream,
                )

        ctx.graph = graph
        ctx.params = params
        ctx.deterministic_dgrad = deterministic_dgrad
        ctx.deterministic_wgrad = deterministic_wgrad
        ctx.save_for_backward(feat, edge_feat, attn_weights, sm_scores, act_scores)

        return output

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_output):
        feat, edge_feat, attn_weights, sm_scores, act_scores = ctx.saved_tensors
        (
            needs_grad_feat,
            needs_grad_efeat,
            needs_grad_weights,
            _,
            _,
            _,
            _,
        ) = ctx.needs_input_grad
        if not (needs_grad_feat or needs_grad_weights or needs_grad_efeat):
            return None, None, None, None, None, None, None

        grad_output = grad_output.detach().contiguous()

        if ctx.graph.num_dst_nodes > 0:
            grad_feat = maybe_empty_like(feat, needs_grad_feat)
            # needs_grad_efeat is also False if edge_feat is None
            # thus capturing all possible cases
            grad_edge_feat = maybe_empty_like(edge_feat, needs_grad_efeat)
            grad_attn_weights = maybe_empty_like(attn_weights, needs_grad_weights)
            grad_wspace_feat = maybe_empty(
                (ctx.graph.num_indices, feat.size(1)),
                create_tensor=(needs_grad_feat and ctx.deterministic_dgrad),
                dtype=feat.dtype,
                device=feat.device,
            )
            grad_wspace_weight = maybe_empty(
                (ctx.graph.num_dst_nodes, attn_weights.size(0)),
                create_tensor=(needs_grad_weights and ctx.deterministic_wgrad),
                dtype=attn_weights.dtype,
                device=attn_weights.device,
            )
            # grad_sm_scores is basically a workspace, thus always required
            grad_sm_scores = torch.empty_like(sm_scores)

            stream = torch.cuda.current_stream().cuda_stream

            if edge_feat is None:
                mha_gat_v2_n2n_bwd(
                    grad_feat,
                    grad_attn_weights,
                    grad_sm_scores,
                    grad_output,
                    feat,
                    attn_weights,
                    sm_scores,
                    act_scores,
                    ctx.graph.graph_csc,
                    ctx.params,
                    grad_wspace_feat,
                    grad_wspace_weight,
                    stream_id=stream,
                )

            else:
                mha_gat_v2_n2n_efeat_bwd(
                    grad_feat,
                    grad_edge_feat,
                    grad_attn_weights,
                    grad_sm_scores,
                    grad_output,
                    feat,
                    attn_weights,
                    sm_scores,
                    act_scores,
                    ctx.graph.graph_csc,
                    ctx.params,
                    grad_wspace_feat,
                    grad_wspace_weight,
                    stream_id=stream,
                )
        else:
            grad_feat = maybe_zeros_like(feat, needs_grad_feat)
            grad_edge_feat = maybe_zeros_like(edge_feat, needs_grad_efeat)
            grad_attn_weights = maybe_zeros_like(attn_weights, needs_grad_weights)

        return grad_feat, grad_edge_feat, grad_attn_weights, None, None, None, None


class _mha_gat_v2_n2n_bipartite_autograd(torch.autograd.Function):
    """
    Custom autograd function for a multi-head attention layer (GAT-like)
    without using cudnn (mha_gat_v2) with an activation prior to the dot
    product but none afterwards operating on bipartite graphs in a node-to-node
    reduction (n2n).
    """

    @staticmethod
    def forward(
        ctx,
        src_feat,
        dst_feat,
        edge_feat,
        attn_weights,
        graph,
        params,
        deterministic_dgrad,
        deterministic_wgrad,
    ):
        dim = src_feat.size(-1)
        src_feat = src_feat.detach().contiguous()
        dst_feat = dst_feat.detach().contiguous()
        attn_weights = attn_weights.detach().contiguous()
        edge_feat = maybe_detach(edge_feat)
        dim_out = dim if params.concat_heads else (dim // params.num_heads)
        output = torch.empty(
            graph.num_dst_nodes, dim_out, dtype=src_feat.dtype, device=src_feat.device
        )
        sm_scores = torch.empty(
            2,
            params.num_heads,
            graph.num_indices,
            dtype=src_feat.dtype,
            device=src_feat.device,
        )
        act_scores = torch.empty(
            graph.num_indices, dim, dtype=src_feat.dtype, device=src_feat.device
        )

        stream = torch.cuda.current_stream().cuda_stream

        if graph.num_dst_nodes > 0:
            if edge_feat is None:
                mha_gat_v2_n2n_fwd(
                    output,
                    sm_scores,
                    act_scores,
                    src_feat,
                    dst_feat,
                    attn_weights,
                    graph.graph_csc,
                    params,
                    stream_id=stream,
                )
            else:
                mha_gat_v2_n2n_efeat_fwd(
                    output,
                    sm_scores,
                    act_scores,
                    src_feat,
                    dst_feat,
                    edge_feat,
                    attn_weights,
                    graph.graph_csc,
                    params,
                    stream_id=stream,
                )

        ctx.graph = graph
        ctx.params = params
        ctx.deterministic_dgrad = deterministic_dgrad
        ctx.deterministic_wgrad = deterministic_wgrad
        ctx.dst_feat_meta_data = get_tensor_meta_data(dst_feat)
        ctx.save_for_backward(src_feat, edge_feat, attn_weights, sm_scores, act_scores)

        return output

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_output):
        src_feat, edge_feat, attn_weights, sm_scores, act_scores = ctx.saved_tensors
        (
            needs_grad_src_feat,
            needs_grad_dst_feat,
            needs_grad_efeat,
            needs_grad_weights,
            _,
            _,
            _,
            _,
        ) = ctx.needs_input_grad
        if not (
            needs_grad_src_feat
            or needs_grad_dst_feat
            or needs_grad_weights
            or needs_grad_efeat
        ):
            return None, None, None, None, None, None, None, None

        grad_output = grad_output.detach().contiguous()

        if ctx.graph.num_dst_nodes > 0:
            grad_src_feat = maybe_empty_like(src_feat, needs_grad_src_feat)
            grad_dst_feat = maybe_empty(
                **ctx.dst_feat_meta_data, create_tensor=needs_grad_dst_feat
            )
            # needs_grad_efeat is also False if edge_feat is None
            # thus capturing all possible cases
            grad_edge_feat = maybe_empty_like(edge_feat, needs_grad_efeat)
            grad_attn_weights = maybe_empty_like(attn_weights, needs_grad_weights)
            grad_wspace_feat = maybe_empty(
                (ctx.graph.num_indices, src_feat.size(1)),
                create_tensor=(needs_grad_src_feat and ctx.deterministic_dgrad),
                dtype=src_feat.dtype,
                device=src_feat.device,
            )
            grad_wspace_weights = maybe_empty(
                (ctx.graph.num_dst_nodes, attn_weights.size(0)),
                create_tensor=(needs_grad_weights and ctx.deterministic_wgrad),
                dtype=attn_weights.dtype,
                device=attn_weights.device,
            )
            # grad_sm_scores is basically a workspace, thus always required
            grad_sm_scores = torch.empty_like(sm_scores)

            stream = torch.cuda.current_stream().cuda_stream

            if edge_feat is None:
                mha_gat_v2_n2n_bwd(
                    grad_src_feat,
                    grad_dst_feat,
                    grad_attn_weights,
                    grad_sm_scores,
                    grad_output,
                    src_feat,
                    attn_weights,
                    sm_scores,
                    act_scores,
                    ctx.graph.graph_csc,
                    ctx.params,
                    grad_wspace_feat,
                    grad_wspace_weights,
                    stream_id=stream,
                )

            else:
                mha_gat_v2_n2n_efeat_bwd(
                    grad_src_feat,
                    grad_dst_feat,
                    grad_edge_feat,
                    grad_attn_weights,
                    grad_sm_scores,
                    grad_output,
                    src_feat,
                    attn_weights,
                    sm_scores,
                    act_scores,
                    ctx.graph.graph_csc,
                    ctx.params,
                    grad_wspace_feat,
                    grad_wspace_weights,
                    stream_id=stream,
                )
        else:
            grad_src_feat = maybe_zeros_like(src_feat, needs_grad_src_feat)
            grad_dst_feat = maybe_zeros(
                **ctx.dst_feat_meta_data, create_tensor=needs_grad_dst_feat
            )
            grad_edge_feat = maybe_zeros_like(edge_feat, needs_grad_efeat)
            grad_attn_weights = maybe_zeros_like(attn_weights, needs_grad_weights)

        return (
            grad_src_feat,
            grad_dst_feat,
            grad_edge_feat,
            grad_attn_weights,
            None,
            None,
            None,
            None,
        )
