# 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, Union, Tuple
from torch.autograd.function import once_differentiable

from ..utils import maybe_empty_like, maybe_detach, maybe_empty, maybe_zeros_like
from ..graph import CSC
from pylibcugraphops import ActivationOp, mha_params
from pylibcugraphops.operators import (
    activation_params,
    mha_gat_n2n_fwd,
    mha_gat_n2n_bwd,
    mha_gat_n2n_efeat_fwd,
    mha_gat_n2n_efeat_bwd,
)

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


def mha_gat_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,
    is_deterministic: bool = False,
    high_precision_gradient: bool = True,
) -> torch.Tensor:
    """
    PyTorch autograd function for a multi-head attention layer (GAT-like)
    without using cudnn (mha_gat) 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 graph, 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 weights with ``dim_w = 2 * dim_in + dim_in_edges`` where
        ``dim_in_edges`` is assumed to be 0 if ``edge_feat`` is None.
        Shape: ``(dim_w,)``.

    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_edge)`` where
        ``dim_in_edge = dim_head_edge * num_heads``, with ``dim_head_edge``
        being the feature dimension per head.

    is_deterministic : Optional[bool], default=False
        Optional flag indicating whether the gradients
        are computed deterministically using a dedicated workspace buffer.

    high_precision_gradient : Optional[bool], default=True
        Optional flag indicating whether gradients for inputs in half precision
        are kept in single precision as long as possible and only casted to
        the corresponding input type at the very end.

    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_n2n_bipartite_autograd.apply(
            feat[0],
            feat[1],
            edge_feat,
            attn_weights,
            graph,
            params,
            is_deterministic,
            high_precision_gradient,
        )

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


class _mha_gat_n2n_autograd(torch.autograd.Function):
    """
    Custom autograd function for a multi-head attention layer (GAT-like)
    without using cudnn (mha_gat) operating on a graph in a node-to-node
    reduction (n2n).
    """

    @staticmethod
    def forward(
        ctx,
        feat,
        edge_feat,
        attn_weights,
        graph,
        params,
        is_deterministic,
        high_precision_gradient,
    ):
        dim = feat.size(-1)
        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
        )

        feat = feat.detach().contiguous()
        attn_weights = attn_weights.detach().contiguous()
        edge_feat = maybe_detach(edge_feat)

        sm_scores = torch.empty(
            2,
            params.num_heads,
            graph.num_indices,
            dtype=torch.float32,
            device=feat.device,
        )

        stream = torch.cuda.current_stream().cuda_stream

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

            else:
                mha_gat_n2n_efeat_fwd(
                    output,
                    sm_scores,
                    feat,
                    edge_feat,
                    attn_weights,
                    graph.graph_csc,
                    params,
                    stream_id=stream,
                )

        ctx.graph = graph
        ctx.params = params
        ctx.is_deterministic = is_deterministic
        ctx.high_precision_gradient = high_precision_gradient
        ctx.save_for_backward(feat, edge_feat, attn_weights, sm_scores)

        return output

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_output):
        feat, edge_feat, attn_weights, sm_scores = ctx.saved_tensors
        grad_output = grad_output.detach().contiguous()
        (
            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

        if ctx.graph.num_dst_nodes > 0:
            grad_dtype = (
                torch.float32 if ctx.high_precision_gradient else grad_output.dtype
            )
            grad_feat = maybe_empty_like(feat, needs_grad_feat, dtype=grad_dtype)
            grad_edge_feat = maybe_empty_like(
                edge_feat, needs_grad_efeat, dtype=grad_dtype
            )
            grad_attn_weights = maybe_empty_like(
                attn_weights, needs_grad_weights, dtype=grad_dtype
            )
            workspace_size = ctx.graph.num_indices * feat.size(1)
            workspace_size += ctx.graph.num_dst_nodes * attn_weights.size(0)
            grad_workspace = maybe_empty(
                workspace_size,
                create_tensor=ctx.is_deterministic,
                dtype=grad_dtype,
                device=attn_weights.device,
            )
            grad_sm_scores = torch.empty_like(sm_scores)

            stream = torch.cuda.current_stream().cuda_stream

            if edge_feat is None:
                mha_gat_n2n_bwd(
                    grad_feat,
                    grad_attn_weights,
                    grad_sm_scores,
                    grad_output,
                    feat,
                    attn_weights,
                    sm_scores,
                    ctx.graph.graph_csc,
                    ctx.params,
                    grad_workspace,
                    stream_id=stream,
                )

            else:
                mha_gat_n2n_efeat_bwd(
                    grad_feat,
                    grad_edge_feat,
                    grad_attn_weights,
                    grad_sm_scores,
                    grad_output,
                    feat,
                    edge_feat,
                    attn_weights,
                    sm_scores,
                    ctx.graph.graph_csc,
                    ctx.params,
                    grad_workspace,
                    stream_id=stream,
                )

            if (grad_feat is not None) and (grad_feat.dtype != feat.dtype):
                grad_feat = grad_feat.to(dtype=feat.dtype)

            if (grad_edge_feat is not None) and (
                grad_edge_feat.dtype != edge_feat.dtype
            ):
                grad_edge_feat = grad_edge_feat.to(dtype=edge_feat.dtype)

            if (grad_attn_weights is not None) and (
                grad_attn_weights.dtype != attn_weights.dtype
            ):
                grad_attn_weights = grad_attn_weights.to(dtype=attn_weights.dtype)

        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_n2n_bipartite_autograd(torch.autograd.Function):
    """
    Custom autograd function for a multi-head attention layer (GAT-like)
    without using cudnn (mha_gat) 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,
        is_deterministic,
        high_precision_gradient,
    ):
        dim = src_feat.size(-1)
        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
        )
        src_feat = src_feat.detach().contiguous()
        dst_feat = dst_feat.detach().contiguous()
        attn_weights = attn_weights.detach().contiguous()
        edge_feat = maybe_detach(edge_feat)

        sm_scores = torch.empty(
            2,
            params.num_heads,
            graph.num_indices,
            dtype=torch.float32,
            device=src_feat.device,
        )

        stream = torch.cuda.current_stream().cuda_stream

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

            else:
                mha_gat_n2n_efeat_fwd(
                    output,
                    sm_scores,
                    src_feat,
                    dst_feat,
                    edge_feat,
                    attn_weights,
                    graph.graph_csc,
                    params,
                    stream_id=stream,
                )

        ctx.graph = graph
        ctx.params = params
        ctx.is_deterministic = is_deterministic
        ctx.high_precision_gradient = high_precision_gradient
        ctx.save_for_backward(src_feat, dst_feat, edge_feat, attn_weights, sm_scores)

        return output

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_output):
        src_feat, dst_feat, edge_feat, attn_weights, sm_scores = ctx.saved_tensors
        grad_output = grad_output.detach().contiguous()
        (
            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

        if ctx.graph.num_dst_nodes > 0:
            grad_dtype = (
                torch.float32 if ctx.high_precision_gradient else grad_output.dtype
            )
            grad_src_feat = maybe_empty_like(
                src_feat, needs_grad_src_feat, dtype=grad_dtype
            )
            grad_dst_feat = maybe_empty_like(
                dst_feat, needs_grad_dst_feat, dtype=grad_dtype
            )
            grad_edge_feat = maybe_empty_like(
                edge_feat, needs_grad_efeat, dtype=grad_dtype
            )
            grad_attn_weights = maybe_empty_like(
                attn_weights, needs_grad_weights, dtype=grad_dtype
            )
            workspace_size = ctx.graph.num_indices * src_feat.size(1)
            workspace_size += ctx.graph.num_dst_nodes * attn_weights.size(0)
            grad_workspace = maybe_empty(
                workspace_size,
                create_tensor=ctx.is_deterministic,
                dtype=grad_dtype,
                device=attn_weights.device,
            )
            grad_sm_scores = torch.empty_like(sm_scores)

            stream = torch.cuda.current_stream().cuda_stream

            if edge_feat is None:
                mha_gat_n2n_bwd(
                    grad_src_feat,
                    grad_dst_feat,
                    grad_attn_weights,
                    grad_sm_scores,
                    grad_output,
                    src_feat,
                    dst_feat,
                    attn_weights,
                    sm_scores,
                    ctx.graph.graph_csc,
                    ctx.params,
                    grad_workspace,
                    stream_id=stream,
                )

            else:
                mha_gat_n2n_efeat_bwd(
                    grad_src_feat,
                    grad_dst_feat,
                    grad_edge_feat,
                    grad_attn_weights,
                    grad_sm_scores,
                    grad_output,
                    src_feat,
                    dst_feat,
                    edge_feat,
                    attn_weights,
                    sm_scores,
                    ctx.graph.graph_csc,
                    ctx.params,
                    grad_workspace,
                    stream_id=stream,
                )

            if (grad_src_feat is not None) and (grad_src_feat.dtype != src_feat.dtype):
                grad_src_feat = grad_src_feat.to(dtype=src_feat.dtype)

            if (grad_dst_feat is not None) and (grad_dst_feat.dtype != dst_feat.dtype):
                grad_dst_feat = grad_dst_feat.to(dtype=dst_feat.dtype)

            if (grad_edge_feat is not None) and (
                grad_edge_feat.dtype != edge_feat.dtype
            ):
                grad_edge_feat = grad_edge_feat.to(dtype=edge_feat.dtype)

            if (grad_attn_weights is not None) and (
                grad_attn_weights.dtype != attn_weights.dtype
            ):
                grad_attn_weights = grad_attn_weights.to(dtype=attn_weights.dtype)

        else:
            grad_src_feat = maybe_zeros_like(src_feat, needs_grad_src_feat)
            grad_dst_feat = maybe_zeros_like(dst_feat, 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,
        )
