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

from ..utils import maybe_empty_like, maybe_empty, get_tensor_meta_data
from ..graph import CSC
from pylibcugraphops.operators import (
    AggOp,
    agg_simple_n2n_fwd,
    agg_simple_n2n_bwd,
    agg_simple_e2n_fwd,
    agg_simple_e2n_bwd,
    agg_simple_n2n_e2n_fwd,
    agg_simple_n2n_e2n_bwd,
    agg_weighted_n2n_fwd,
    agg_weighted_n2n_bwd,
)

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


def agg_simple_n2n(
    feat: torch.Tensor,
    graph: CSC,
    aggr: str = "sum",
    edge_weight: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    r"""PyTorch autograd function for simple aggregation (agg_simple)
    using node features in an node-to-node reduction (n2n).

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

    graph : CSC
        The graph used for the operation.

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

    edge_weight : Optional[torch.Tensor], default=None
        When passing additional edge_weights, the aggregation is done
        in a weighted fashion. Note that for min/max the result after
        weighting features with those edge weights is returned.
        Edge weights are indexed in similar to indices in `graph`.
        Shape: (n_edges, ).

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

    if aggr not in _aggr_ops:
        raise ValueError(
            f"Supported aggregation operations: {_aggr_ops}, " f"but got '{aggr}'."
        )
    aggr = getattr(AggOp, aggr.capitalize())
    return _agg_simple_n2n_autograd.apply(feat, edge_weight, graph, aggr)


def agg_simple_e2n(
    edge_feat: torch.Tensor,
    graph: CSC,
    aggr: str = "sum",
) -> torch.Tensor:
    r"""PyTorch autograd function for simple aggregation (agg_simple)
    using edge features in an edge-to-node reduction (e2n).

    Parameters
    ----------
    edge_feat : torch.Tensor
        The input edge features.
        Shape: (n_edges, dim_edge).

    graph : CSC
        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_dst_nodes, dim_node + dim_edge).
    """

    if aggr not in _aggr_ops:
        raise ValueError(
            f"Supported aggregation operations: {_aggr_ops}, " f"but got '{aggr}'."
        )
    aggr = getattr(AggOp, aggr.capitalize())
    return _agg_simple_e2n_autograd.apply(edge_feat, graph, aggr)


def agg_simple_n2n_e2n(
    node_feat: torch.Tensor,
    edge_feat: torch.Tensor,
    graph: CSC,
    aggr: str = "sum",
) -> torch.Tensor:
    r"""PyTorch autograd function for simple aggregation (agg_simple)
    using node features in an node-to-node reduction (n2n) and
    edge features in and edge-to-node reduction (e2n) an edge-to-node reduction (n2n).

    Parameters
    ----------
    node_feat : torch.Tensor
        The input node features.
        Shape: (n_dst_nodes, dim_node)

    edge_feat : torch.Tensor
        The input edge features.
        Shape: (n_edges, dim_edge).

    graph : CSC
        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_dst_nodes, dim_node + dim_edge).
    """

    if aggr not in _aggr_ops:
        raise ValueError(
            f"Supported aggregation operations: {_aggr_ops}, " f"but got '{aggr}'."
        )
    aggr = getattr(AggOp, aggr.capitalize())
    return _agg_simple_n2n_e2n_autograd.apply(node_feat, edge_feat, graph, aggr)


class _agg_simple_n2n_autograd(torch.autograd.Function):
    r"""Custom autograd function for simple aggregation (agg_simple)
    using node features in an node-to-node reduction (n2n) on a graph."""

    @staticmethod
    def forward(ctx, feat, edge_weight, graph, aggr):
        dim_in = feat.size(-1)
        agg_output = torch.empty(
            (graph.num_dst_nodes, dim_in), dtype=feat.dtype, device=feat.device
        )
        node_deg = None

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

        stream = torch.cuda.current_stream().cuda_stream

        feat = feat.detach().contiguous()

        if edge_weight is None:
            agg_simple_n2n_fwd(
                agg_output,
                feat,
                graph.graph_csc,
                aggr,
                out_pos,
                stream_id=stream,
            )
        else:
            node_deg = torch.empty(
                graph.num_dst_nodes, dtype=edge_weight.dtype, device=edge_weight.device
            )
            edge_weight = edge_weight.detach().contiguous()
            agg_weighted_n2n_fwd(
                agg_output,
                feat,
                edge_weight,
                graph.graph_csc,
                aggr,
                out_pos,
                node_deg,
                stream_id=stream,
            )

        ctx.graph = graph
        ctx.aggr = aggr
        ctx.feat_meta_data = get_tensor_meta_data(feat)
        saved_feat = feat if edge_weight is not None else None
        saved_agg_output = (
            agg_output.detach().contiguous()
            if (edge_weight is not None) and (aggr == AggOp.Mean)
            else None
        )
        ctx.save_for_backward(
            saved_feat, edge_weight, node_deg, out_pos, saved_agg_output
        )

        return agg_output

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_output):
        grad_output = grad_output.detach().contiguous()
        feat, edge_weight, node_deg, out_pos, agg_output = ctx.saved_tensors
        needs_grad_feat, needs_grad_weight, _, _ = ctx.needs_input_grad
        if not (needs_grad_feat or needs_grad_weight):
            return None, None, None, None

        grad_output = grad_output.detach().contiguous()
        grad_feat = maybe_empty(**ctx.feat_meta_data, create_tensor=needs_grad_feat)
        grad_edge_weight = maybe_empty_like(edge_weight, edge_weight is not None)

        stream = torch.cuda.current_stream().cuda_stream

        if edge_weight is None:
            agg_simple_n2n_bwd(
                grad_feat,
                grad_output,
                ctx.graph.graph_csc,
                ctx.aggr,
                out_pos,
                stream_id=stream,
            )

        else:
            agg_weighted_n2n_bwd(
                grad_feat,
                grad_edge_weight,
                grad_output,
                edge_weight,
                feat,
                ctx.graph.graph_csc,
                aggregation_operation=ctx.aggr,
                output_embedding=agg_output,
                out_pos=out_pos,
                node_degree=node_deg,
                stream_id=stream,
            )

        return grad_feat, grad_edge_weight, None, None


class _agg_simple_e2n_autograd(torch.autograd.Function):
    r"""Custom autograd function for simple aggregation (agg_simple)
    using edge features in an edge-to-node reduction (e2n) on a graph."""

    @staticmethod
    def forward(ctx, edge_feat, graph, aggr):
        dim_edge = edge_feat.size(-1)
        agg_output = torch.empty(
            graph.num_dst_nodes,
            dim_edge,
            dtype=edge_feat.dtype,
            device=edge_feat.device,
        )

        if aggr in (AggOp.Max, AggOp.Min):
            idx_type = graph.graph_csc.__class__.__name__.split("_")[-1]
            out_pos = torch.empty(
                graph.num_dst_nodes,
                dim_edge,
                dtype=getattr(torch, idx_type),
                device=edge_feat.device,
            )
        else:
            out_pos = None

        stream = torch.cuda.current_stream().cuda_stream

        agg_simple_e2n_fwd(
            agg_output,
            edge_feat.detach(),
            graph.graph_csc,
            aggr,
            out_pos,
            stream_id=stream,
        )

        ctx.graph = graph
        ctx.aggr = aggr
        ctx.edge_feat_meta_data = get_tensor_meta_data(edge_feat)
        ctx.save_for_backward(out_pos)

        return agg_output

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_output):
        grad_output = grad_output.detach().contiguous()
        (out_pos,) = ctx.saved_tensors
        needs_grad_edge_feat, _, _ = ctx.needs_input_grad
        if not needs_grad_edge_feat:
            return None, None, None

        grad_output = grad_output.detach().contiguous()
        grad_edge_feat = maybe_empty(**ctx.edge_feat_meta_data, create_tensor=True)
        stream = torch.cuda.current_stream().cuda_stream

        agg_simple_e2n_bwd(
            grad_edge_feat,
            grad_output,
            ctx.graph.graph_csc,
            ctx.aggr,
            out_pos,
            stream_id=stream,
        )

        return grad_edge_feat, None, None


class _agg_simple_n2n_e2n_autograd(torch.autograd.Function):
    r"""Custom autograd function for simple aggregation (agg_simple)
    using node features in an node-to-node reduction (n2n)
    and edge features in an edge-to-node reduction (e2n)."""

    @staticmethod
    def forward(ctx, node_feat, edge_feat, graph, aggr):
        dim_node = node_feat.size(-1)
        dim_edge = edge_feat.size(-1)
        agg_output = torch.empty(
            graph.num_dst_nodes,
            dim_node + dim_edge,
            dtype=edge_feat.dtype,
            device=edge_feat.device,
        )

        if aggr in (AggOp.Max, AggOp.Min):
            idx_type = graph.graph_csc.__class__.__name__.split("_")[-1]
            out_pos = torch.empty(
                graph.num_dst_nodes,
                dim_node + dim_edge,
                dtype=getattr(torch, idx_type),
                device=edge_feat.device,
            )
        else:
            out_pos = None

        stream = torch.cuda.current_stream().cuda_stream

        agg_simple_n2n_e2n_fwd(
            agg_output,
            node_feat.detach(),
            edge_feat.detach(),
            graph.graph_csc,
            aggr,
            out_pos,
            stream_id=stream,
        )

        ctx.graph = graph
        ctx.aggr = aggr
        ctx.dim_node = dim_node
        ctx.dim_edge = dim_edge
        ctx.node_feat_meta_data = get_tensor_meta_data(node_feat)
        ctx.edge_feat_meta_data = get_tensor_meta_data(edge_feat)
        ctx.save_for_backward(out_pos)

        return agg_output

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_output):
        grad_output = grad_output.detach().contiguous()
        (out_pos,) = ctx.saved_tensors
        needs_grad_node_feat, needs_grad_edge_feat, _, _ = ctx.needs_input_grad
        if not (needs_grad_node_feat or needs_grad_edge_feat):
            return None, None, None, None

        grad_output = grad_output.detach().contiguous()
        grad_node_feat = maybe_empty(
            **ctx.node_feat_meta_data, create_tensor=needs_grad_node_feat
        )
        grad_edge_feat = maybe_empty(
            **ctx.edge_feat_meta_data, create_tensor=needs_grad_edge_feat
        )

        stream = torch.cuda.current_stream().cuda_stream

        agg_simple_n2n_e2n_bwd(
            grad_node_feat,
            grad_edge_feat,
            grad_output,
            ctx.graph.graph_csc,
            ctx.dim_node,
            ctx.dim_edge,
            ctx.aggr,
            out_pos,
            stream_id=stream,
        )

        return grad_node_feat, grad_edge_feat, None, None
