# 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.

import torch

from ...typing import CuGraphAnyCsr
from ...binding.pylibcugraphops_ext.operators import (
    AggOp,
    agg_concat_fg_n2n_fwd as fg_n2n_fwd,
    agg_concat_fg_n2n_bwd as fg_n2n_bwd,
    agg_concat_mfg_n2n_fwd as mfg_n2n_fwd,
    agg_concat_mfg_n2n_bwd as mfg_n2n_bwd,
)

_aggr_ops = ["max", "mean", "min", "sum"]


def agg_concat_n2n(
    feat: torch.Tensor, graph: CuGraphAnyCsr, aggr: str = "sum"
) -> torch.Tensor:
    r"""PyTorch autograd function for simple aggregation using node features in
    an node-to-node reduction (n2n) while concatenating the original features of
    output at the end (agg_concat).

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

    graph : fg_csr_int[32|64] | mfg_csr_int[32|64]
        The graph used for the operation.

    aggr : str, default="sum"
        Aggregation operation. Choose from 'max', 'mean', 'min', 'sum'.

    Returns
    -------
    output : torch.Tensor
        The aggregation output.
        Shape: (n_out_nodes, 2 * dim_in).

    Notes
    -----
    Replace n_in_nodes and n_out_nodes with n_nodes when given a full graph.
    """
    use_mfg = hasattr(graph, "sample_size")
    if aggr not in _aggr_ops:
        raise ValueError(
            f"Supported aggregation operations: {_aggr_ops}, " f"but got '{aggr}'."
        )
    aggr = getattr(AggOp, aggr.capitalize())

    if use_mfg:
        return _agg_concat_mfg_n2n_autograd.apply(feat, graph, aggr)
    return _agg_concat_fg_n2n_autograd.apply(feat, graph, aggr)


class _agg_concat_fg_n2n_autograd(torch.autograd.Function):
    r"""Custom autograd function for simple aggregation using node features in
    an node-to-node reduction (n2n) on a full graph (fg) while concatenating the
    original features of output at the end (agg_concat)."""

    @staticmethod
    def forward(ctx, feat, fg, aggr):
        feat = feat.contiguous()
        dim_in = feat.shape[-1]
        agg_output = torch.zeros(
            fg.n_nodes, 2 * dim_in, dtype=feat.dtype, device=feat.device
        )

        if aggr in (AggOp.Max, AggOp.Min):
            idx_type = fg.__class__.__name__.split("_")[-1]
            out_pos = torch.zeros(
                fg.n_nodes, dim_in, dtype=getattr(torch, idx_type), device=feat.device
            )
        else:
            out_pos = None

        stream = torch.cuda.current_stream().cuda_stream
        fg_n2n_fwd(agg_output, feat.detach(), fg, aggr, out_pos, stream_id=stream)

        ctx.fg = fg
        ctx.aggr = aggr
        ctx.save_for_backward(feat, out_pos)

        return agg_output

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

        stream = torch.cuda.current_stream().cuda_stream
        fg_n2n_bwd(grad_feat, grad_output, ctx.fg, ctx.aggr, out_pos, stream_id=stream)

        return grad_feat, None, None


class _agg_concat_mfg_n2n_autograd(torch.autograd.Function):
    r"""Custom autograd function for simple aggregation using node features in
    an node-to-node reduction (n2n) on a message-flow-graph (mfg) while
    concatenating the original features of output at the end (agg_concat)."""

    @staticmethod
    def forward(ctx, feat, mfg, aggr):
        feat = feat.contiguous()
        dim_in = feat.shape[-1]
        agg_output = torch.zeros(
            mfg.n_out_nodes, 2 * dim_in, dtype=feat.dtype, device=feat.device
        )

        if aggr in (AggOp.Max, AggOp.Min):
            idx_type = mfg.__class__.__name__.split("_")[-1]
            out_pos = torch.zeros(
                mfg.n_out_nodes,
                dim_in,
                dtype=getattr(torch, idx_type),
                device=feat.device,
            )
        else:
            out_pos = None

        stream = torch.cuda.current_stream().cuda_stream
        mfg_n2n_fwd(agg_output, feat.detach(), mfg, aggr, out_pos, stream_id=stream)

        ctx.mfg = mfg
        ctx.aggr = aggr
        ctx.save_for_backward(feat, out_pos)

        return agg_output

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

        stream = torch.cuda.current_stream().cuda_stream
        mfg_n2n_bwd(
            grad_feat, grad_output, ctx.mfg, ctx.aggr, out_pos, stream_id=stream
        )

        return grad_feat, None, None
