# 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,
    get_operator_from_module,
)
from ..graph import CSC
from pylibcugraphops import ActivationOp, mha_params
from pylibcugraphops.operators import activation_params
import pylibcugraphops.operators as ops

_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,
    deterministic_dgrad: bool = False,
    deterministic_wgrad: bool = False,
    high_precision_dgrad: bool = False,
    high_precision_wgrad: bool = False,
) -> 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.

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

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

    high_precision_dgrad : Optional[bool], default=False
        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.

    high_precision_wgrad : Optional[bool], default=False
        Optional flag indicating whether gradients for weights 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,
            deterministic_dgrad,
            deterministic_wgrad,
            high_precision_dgrad,
            high_precision_wgrad,
        )

    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,
        deterministic_dgrad,
        deterministic_wgrad,
        high_precision_dgrad,
        high_precision_wgrad,
    )


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,
        deterministic_dgrad,
        deterministic_wgrad,
        high_precision_dgrad,
        high_precision_wgrad,
    ):
        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:
                fwd_fun = get_operator_from_module(
                    ops,
                    "mha_gat_n2n_fwd",
                    (
                        output.dtype,
                        sm_scores.dtype,
                    ),
                )
                fwd_fun(
                    output,
                    sm_scores,
                    feat,
                    attn_weights,
                    graph.graph_csc,
                    params,
                    stream_id=stream,
                )

            else:
                fwd_fun = get_operator_from_module(
                    ops,
                    "mha_gat_n2n_efeat_fwd",
                    (
                        output.dtype,
                        sm_scores.dtype,
                    ),
                )

                fwd_fun(
                    output,
                    sm_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.high_precision_dgrad = high_precision_dgrad
        ctx.high_precision_wgrad = high_precision_wgrad
        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, None, None

        if ctx.graph.num_dst_nodes > 0:
            dgrad_dtype = (
                torch.float32 if ctx.high_precision_dgrad else grad_output.dtype
            )
            wgrad_dtype = (
                torch.float32 if ctx.high_precision_dgrad else grad_output.dtype
            )
            grad_feat = maybe_empty_like(feat, needs_grad_feat, dtype=dgrad_dtype)
            grad_edge_feat = maybe_empty_like(
                edge_feat, needs_grad_efeat, dtype=dgrad_dtype
            )
            grad_attn_weights = maybe_empty_like(
                attn_weights, needs_grad_weights, dtype=wgrad_dtype
            )
            grad_wspace_feat = maybe_empty(
                (ctx.graph.num_indices, feat.size(1)),
                create_tensor=(needs_grad_feat and ctx.deterministic_dgrad),
                dtype=dgrad_dtype,
                device=attn_weights.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=wgrad_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:
                bwd_fun = get_operator_from_module(
                    ops,
                    "mha_gat_n2n_bwd",
                    (
                        grad_output.dtype,
                        dgrad_dtype,
                        wgrad_dtype,
                        sm_scores.dtype,
                    ),
                )
                bwd_fun(
                    grad_feat,
                    grad_attn_weights if grad_wspace_weight is None else None,
                    grad_sm_scores,
                    grad_output,
                    feat,
                    attn_weights,
                    sm_scores,
                    ctx.graph.graph_csc,
                    ctx.params,
                    grad_wspace_feat,
                    grad_wspace_weight,
                    stream_id=stream,
                )

                if (grad_wspace_weight is not None) and (grad_attn_weights is not None):
                    grad_attn_weights = grad_wspace_weight.sum(dim=0)

            else:
                bwd_fun = get_operator_from_module(
                    ops,
                    "mha_gat_n2n_efeat_bwd",
                    (
                        grad_output.dtype,
                        dgrad_dtype,
                        wgrad_dtype,
                        sm_scores.dtype,
                    ),
                )
                bwd_fun(
                    grad_feat,
                    grad_edge_feat,
                    grad_attn_weights if grad_wspace_weight is None else None,
                    grad_sm_scores,
                    grad_output,
                    feat,
                    edge_feat,
                    attn_weights,
                    sm_scores,
                    ctx.graph.graph_csc,
                    ctx.params,
                    grad_wspace_feat,
                    grad_wspace_weight,
                    stream_id=stream,
                )

                if (grad_wspace_weight is not None) and (grad_attn_weights is not None):
                    grad_attn_weights = grad_wspace_weight.sum(dim=0)

            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,
            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,
        deterministic_dgrad,
        deterministic_wgrad,
        high_precision_dgrad,
        high_precision_wgrad,
    ):
        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:
                fwd_fun = get_operator_from_module(
                    ops,
                    "mha_gat_n2n_fwd",
                    (
                        output.dtype,
                        sm_scores.dtype,
                    ),
                )
                fwd_fun(
                    output,
                    sm_scores,
                    src_feat,
                    dst_feat,
                    attn_weights,
                    graph.graph_csc,
                    params,
                    stream_id=stream,
                )

            else:
                fwd_fun = get_operator_from_module(
                    ops,
                    "mha_gat_n2n_efeat_fwd",
                    (
                        output.dtype,
                        sm_scores.dtype,
                    ),
                )
                fwd_fun(
                    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.deterministic_dgrad = deterministic_dgrad
        ctx.deterministic_wgrad = deterministic_wgrad
        ctx.high_precision_dgrad = high_precision_dgrad
        ctx.high_precision_wgrad = high_precision_wgrad
        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, None, None

        if ctx.graph.num_dst_nodes > 0:
            dgrad_dtype = (
                torch.float32 if ctx.high_precision_dgrad else grad_output.dtype
            )
            wgrad_dtype = (
                torch.float32 if ctx.high_precision_wgrad else grad_output.dtype
            )
            grad_src_feat = maybe_empty_like(
                src_feat, needs_grad_src_feat, dtype=dgrad_dtype
            )
            grad_dst_feat = maybe_empty_like(
                dst_feat, needs_grad_dst_feat, dtype=dgrad_dtype
            )
            grad_edge_feat = maybe_empty_like(
                edge_feat, needs_grad_efeat, dtype=dgrad_dtype
            )
            grad_attn_weights = maybe_empty_like(
                attn_weights, needs_grad_weights, dtype=wgrad_dtype
            )
            grad_wspace_feat = maybe_empty(
                (ctx.graph.num_indices, src_feat.size(1)),
                create_tensor=(needs_grad_dst_feat and ctx.deterministic_dgrad),
                dtype=dgrad_dtype,
                device=attn_weights.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=wgrad_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:
                bwd_fun = get_operator_from_module(
                    ops,
                    "mha_gat_n2n_bwd",
                    (
                        grad_output.dtype,
                        dgrad_dtype,
                        wgrad_dtype,
                        sm_scores.dtype,
                    ),
                )
                bwd_fun(
                    grad_src_feat,
                    grad_dst_feat,
                    grad_attn_weights if grad_wspace_weight is None else None,
                    grad_sm_scores,
                    grad_output,
                    src_feat,
                    dst_feat,
                    attn_weights,
                    sm_scores,
                    ctx.graph.graph_csc,
                    ctx.params,
                    grad_wspace_feat,
                    grad_wspace_weight,
                    stream_id=stream,
                )

                if (grad_wspace_weight is not None) and (grad_attn_weights is not None):
                    grad_attn_weights = grad_wspace_weight.sum(dim=0)

            else:
                bwd_fun = get_operator_from_module(
                    ops,
                    "mha_gat_n2n_efeat_bwd",
                    (
                        grad_output.dtype,
                        dgrad_dtype,
                        wgrad_dtype,
                        sm_scores.dtype,
                    ),
                )
                bwd_fun(
                    grad_src_feat,
                    grad_dst_feat,
                    grad_edge_feat,
                    grad_attn_weights if grad_wspace_weight is None else None,
                    grad_sm_scores,
                    grad_output,
                    src_feat,
                    dst_feat,
                    edge_feat,
                    attn_weights,
                    sm_scores,
                    ctx.graph.graph_csc,
                    ctx.params,
                    grad_wspace_feat,
                    grad_wspace_weight,
                    stream_id=stream,
                )

                if (grad_wspace_weight is not None) and (grad_attn_weights is not None):
                    grad_attn_weights = grad_wspace_weight.sum(dim=0)

            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,
            None,
            None,
        )
