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

import torch

from ...typing import CuGraphAnyHgCsr
from ...binding.pylibcugraphops_ext.operators import (
    agg_hg_basis_fg_n2n_post_fwd as fg_n2n_post_fwd,
    agg_hg_basis_fg_n2n_post_bwd as fg_n2n_post_bwd,
    agg_hg_basis_mfg_n2n_post_fwd as mfg_n2n_post_fwd,
    agg_hg_basis_mfg_n2n_post_bwd as mfg_n2n_post_bwd,
)


def agg_hg_basis_n2n_post(
    feat: torch.Tensor,
    weights_comb: Optional[torch.Tensor],
    graph: CuGraphAnyHgCsr,
    concat_own: bool = False,
    norm_by_out_degree: bool = False,
) -> torch.Tensor:
    r"""PyTorch autograd function for node-to-node RGCN-like basis regularized
    aggregation, with features being transformed after (post) this aggregation.

    Parameters
    ----------
    feat : torch.Tensor
        The input node features.
        Shape: (n_in_nodes, in_dim).

    weights_comb : torch.Tensor | None
        The combination weights. If `None`, no weights are used and features of
        different edge types are aggregated into separate output dimensions.
        Shape: (n_edge_types, n_bases) if provided.

    graph : fg_csr_hg_int[32|64] | mfg_csr_hg_int[32|64]
        The graph used for the operation.

    concat_own : bool, default=False
        Concatenate output node embeddings in the aggregation.

    norm_by_out_degree : bool, default=False
        If set, output embeddings are normed by the degree of the output node
        per edge type.

    Returns
    -------
    output : torch.Tensor
        The aggregation output.
        Shape: (n_out_nodes, out_dim) if `concat_own` is `True`;
               (n_out_nodes, out_dim + in_dim) otherwise;
               where out_dim = in_dim * n_bases if `weights_comb` is provided,
               otherwise, out_dim = in_dim * n_edge_types.

    Notes
    -----
    Replace n_in_nodes and n_out_nodes with n_nodes when given a full graph.
    """
    use_mfg = hasattr(graph, "sample_size")
    options = (concat_own, norm_by_out_degree)
    if use_mfg:
        return _agg_hg_basis_mfg_n2n_post_autograd.apply(
            feat, weights_comb, graph, options
        )

    return _agg_hg_basis_fg_n2n_post_autograd.apply(feat, weights_comb, graph, options)


class _agg_hg_basis_fg_n2n_post_autograd(torch.autograd.Function):
    r"""Custom autograd function for node-to-node full-graph RGCN-like basis
    regularized aggregation, with features being transformed after (post)
    this aggregation."""

    @staticmethod
    def forward(ctx, feat, weights_comb, fg, options):
        feat = feat.contiguous()
        concat_own, norm_by_out_degree = options
        dim = feat.shape[-1]
        n_bases = fg.n_edge_types if weights_comb is None else weights_comb.shape[-1]
        out_dim = dim * (n_bases + (1 if concat_own else 0))

        agg_output = torch.zeros(
            fg.n_nodes, out_dim, dtype=feat.dtype, device=feat.device
        )

        stream = torch.cuda.current_stream().cuda_stream
        if weights_comb is None:
            fg_n2n_post_fwd(
                agg_output,
                feat.detach(),
                None,
                fg,
                concat_own,
                norm_by_out_degree,
                stream_id=stream,
            )
        else:
            weights_comb = weights_comb.contiguous()
            fg_n2n_post_fwd(
                agg_output,
                feat.detach(),
                weights_comb.detach(),
                fg,
                concat_own,
                norm_by_out_degree,
                stream_id=stream,
            )

        ctx.fg = fg
        ctx.concat_own = concat_own
        ctx.norm_by_out_degree = norm_by_out_degree
        ctx.save_for_backward(feat, weights_comb)

        return agg_output

    @staticmethod
    def backward(ctx, grad_output):
        grad_output = grad_output.contiguous()
        feat, weights_comb = ctx.saved_tensors
        grad_feat = torch.zeros_like(feat)

        stream = torch.cuda.current_stream().cuda_stream
        if weights_comb is None:
            grad_weights_comb = None
            fg_n2n_post_bwd(
                grad_feat,
                None,
                grad_output,
                feat.detach(),
                None,
                ctx.fg,
                ctx.concat_own,
                ctx.norm_by_out_degree,
                stream_id=stream,
            )
        else:
            grad_weights_comb = torch.zeros_like(weights_comb)
            fg_n2n_post_bwd(
                grad_feat,
                grad_weights_comb,
                grad_output,
                feat.detach(),
                weights_comb.detach(),
                ctx.fg,
                ctx.concat_own,
                ctx.norm_by_out_degree,
                stream_id=stream,
            )

        return grad_feat, grad_weights_comb, None, None


class _agg_hg_basis_mfg_n2n_post_autograd(torch.autograd.Function):
    r"""Custom autograd function for node-to-node message-flow-graph RGCN-like
    basis regularized aggregation, with features being transformed after (post)
    this aggregation."""

    @staticmethod
    def forward(ctx, feat, weights_comb, mfg, options):
        feat = feat.contiguous()
        concat_own, norm_by_out_degree = options
        dim = feat.shape[-1]
        n_bases = mfg.n_edge_types if weights_comb is None else weights_comb.shape[-1]
        out_dim = dim * (n_bases + (1 if concat_own else 0))

        agg_output = torch.zeros(
            mfg.n_out_nodes, out_dim, dtype=feat.dtype, device=feat.device
        )

        stream = torch.cuda.current_stream().cuda_stream
        if weights_comb is None:
            mfg_n2n_post_fwd(
                agg_output,
                feat.detach(),
                None,
                mfg,
                concat_own,
                norm_by_out_degree,
                stream_id=stream,
            )
        else:
            weights_comb = weights_comb.contiguous()
            mfg_n2n_post_fwd(
                agg_output,
                feat.detach(),
                weights_comb.detach(),
                mfg,
                concat_own,
                norm_by_out_degree,
                stream_id=stream,
            )

        ctx.mfg = mfg
        ctx.concat_own = concat_own
        ctx.norm_by_out_degree = norm_by_out_degree
        ctx.save_for_backward(feat, weights_comb)

        return agg_output

    @staticmethod
    def backward(ctx, grad_output):
        grad_output = grad_output.contiguous()
        feat, weights_comb = ctx.saved_tensors

        grad_feat = torch.zeros_like(feat)
        stream = torch.cuda.current_stream().cuda_stream
        if weights_comb is None:
            grad_weights_comb = None
            mfg_n2n_post_bwd(
                grad_feat,
                None,
                grad_output,
                feat.detach(),
                None,
                ctx.mfg,
                ctx.concat_own,
                ctx.norm_by_out_degree,
                stream_id=stream,
            )
        else:
            grad_weights_comb = torch.zeros_like(weights_comb)
            mfg_n2n_post_bwd(
                grad_feat,
                grad_weights_comb,
                grad_output,
                feat.detach(),
                weights_comb.detach(),
                ctx.mfg,
                ctx.concat_own,
                ctx.norm_by_out_degree,
                stream_id=stream,
            )

        return grad_feat, grad_weights_comb, None, None
