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

from __future__ import annotations
import warnings
from typing import Tuple, Any

import torch
from torch import Tensor
import pylibcugraphops


class CSC:
    def __init__(
        self,
        offsets: Tensor,
        indices: Tensor,
        num_src_nodes: int,
        map_csc_to_coo: Tensor | None = None,
        rev_offsets: Tensor | None = None,
        rev_indices: Tensor | None = None,
        map_rev_to_coo: Tensor | None = None,
        map_dst_to_src: Tensor | None = None,
        dst_max_in_degree: int = -1,
        reverse_graph_bwd: bool = False,
        is_bipartite: bool = False,
    ):
        self.map_dst_to_src = map_dst_to_src
        self.num_src_nodes = num_src_nodes
        self.num_dst_nodes = offsets.size(0) - 1
        self.num_indices = indices.size(0)

        self.offsets = offsets
        self.indices = indices
        self.map_csc_to_coo = map_csc_to_coo

        self._rev_cache_filled = False
        if reverse_graph_bwd:
            if (
                isinstance(rev_offsets, Tensor)
                and isinstance(rev_indices, Tensor)
                and isinstance(map_rev_to_coo, Tensor)
            ):
                self.rev_offsets = rev_offsets
                self.rev_indices = rev_indices
                self.map_rev_to_coo = map_rev_to_coo
                self._rev_cache_filled = True
            else:
                self._create_rev_buffers()
        else:
            self.rev_offsets = self.rev_indices = self.map_rev_to_coo = None
        self.reverse_graph_bwd = reverse_graph_bwd

        self.is_bipartite = is_bipartite

        if map_dst_to_src is not None:
            raise NotImplementedError("map_dst_to_src is currently not supported")

        # TODO remove warning once we again have support for these
        if dst_max_in_degree > 0:
            warnings.warn("dst_max_in_degree currently has no effect")
        self.dst_max_in_degree = dst_max_in_degree

        self._build_graph()

        if self.reverse_graph_bwd:
            self.add_reverse_graph()

    def _create_rev_buffers(self):
        self.rev_offsets = torch.empty(
            self.num_src_nodes + 1,
            device=self.offsets.device,
            dtype=self.offsets.dtype,
        )
        self.rev_indices = torch.empty_like(self.indices)
        self.map_rev_to_coo = torch.empty_like(self.indices)

    def _build_graph(self):
        if self.is_bipartite:
            make_graph_func = pylibcugraphops.make_bipartite_csc
        else:
            make_graph_func = pylibcugraphops.make_csc

        if not self.reverse_graph_bwd:
            self.graph_csc = make_graph_func(
                self.offsets,
                self.indices,
                self.num_src_nodes,
                self.map_csc_to_coo,
            )
        else:
            self.graph_csc = make_graph_func(
                self.offsets,
                self.indices,
                self.num_src_nodes,
                self.map_csc_to_coo,
                self.rev_offsets,
                self.rev_indices,
                self.map_rev_to_coo,
            )

    def add_reverse_graph(self):
        if self._rev_cache_filled:
            return

        # skip reversing empty graphs to avoid zero-grid kernel launch
        if self.indices.numel() == 0:
            return

        if self.reverse_graph_bwd is False:
            self.reverse_graph_bwd = True
            self._create_rev_buffers()
            self._build_graph()

        stream = torch.cuda.current_stream().cuda_stream
        # call reverse to get necessary workspace size and
        # create necessary workspace buffers
        node_counts = torch.empty(
            self.num_src_nodes,
            dtype=self.offsets.dtype,
            device=self.offsets.device,
        )
        cub_workspace = None
        cub_workspace_size = pylibcugraphops.reverse_graph(
            self.graph_csc, node_counts, cub_workspace, 0, stream
        )
        cub_workspace = torch.empty(
            cub_workspace_size, dtype=torch.int8, device=self.offsets.device
        )
        # actually reverse graph (in-place)
        pylibcugraphops.reverse_graph(
            self.graph_csc, node_counts, cub_workspace, cub_workspace_size, stream
        )

        self._rev_cache_filled = True

    def to(self, *args, **kwargs):
        self.offsets = self.offsets.to(*args, **kwargs)
        self.indices = self.indices.to(*args, **kwargs)
        if self.map_dst_to_src is not None:
            self.map_dst_to_src = self.map_dst_to_src.to(*args, **kwargs)
        if self.map_csc_to_coo is not None:
            self.map_csc_to_coo = self.map_csc_to_coo.to(*args, **kwargs)
        if self.map_rev_to_coo is not None:
            self.map_rev_to_coo = self.map_rev_to_coo.to(*args, **kwargs)
        if self.reverse_graph_bwd:
            self.rev_offsets = self.rev_offsets.to(*args, **kwargs)
            self.rev_indices = self.rev_indices.to(*args, **kwargs)

        self._build_graph()

    def __to__(self, *args, **kwargs):
        self.to(*args, **kwargs)

    @classmethod
    def from_edge_index(
        cls,
        edge_index: Any,
        reverse_graph_bwd: bool = False,
        is_bipartite: bool = False,
    ):
        """Create CSC graph from a torch_geometric.EdgeIndex object."""
        num_src_nodes = edge_index.sparse_size(0)
        edge_index = edge_index.sort_by("col")[0]
        colptr, row = edge_index.get_csc()[0]

        if not reverse_graph_bwd:
            return cls(
                colptr,
                row,
                num_src_nodes,
                reverse_graph_bwd=False,
                is_bipartite=is_bipartite,
            )

        (rowptr, col), perm = edge_index.get_csr()
        perm = perm.to(dtype=edge_index.dtype)
        return cls(
            colptr,
            row,
            num_src_nodes,
            rev_offsets=rowptr,
            rev_indices=col,
            map_rev_to_coo=perm,
            reverse_graph_bwd=True,
            is_bipartite=is_bipartite,
        )


class HeteroCSC(CSC):
    def __init__(
        self,
        offsets: Tensor,
        indices: Tensor,
        edge_types: Tensor,
        num_src_nodes: int,
        num_edge_types: int,
        num_node_types: int = 0,
        node_types: Tensor | Tuple[Tensor, Tensor] | None = None,
        map_csc_to_coo: Tensor | None = None,
        rev_offsets: Tensor | None = None,
        rev_indices: Tensor | None = None,
        map_rev_to_coo: Tensor | None = None,
        map_dst_to_src: Tensor | None = None,
        dst_max_in_degree: int = -1,
        reverse_graph_bwd: bool = False,
        is_bipartite: bool = False,
    ):
        if is_bipartite:
            if num_node_types > 0:
                assert isinstance(node_types, Tuple), (
                    "Expected tuple of node types for a bipartite graph"
                    f", but got {type(node_types)} instead."
                )

        else:
            if num_node_types > 0:
                assert isinstance(node_types, Tensor), (
                    "Expected node_types to be of type Tensor"
                    f", but got {type(node_types)} instead."
                )

        self.edge_types = edge_types.to(torch.int32)
        self.num_node_types = num_node_types
        self.num_edge_types = num_edge_types
        if node_types is None:
            self.node_types = None
        elif isinstance(node_types, Tensor):
            self.node_types = node_types.to(torch.int32)
        else:
            self.node_types = (
                node_types[0].to(torch.int32),
                node_types[1].to(torch.int32),
            )

        super().__init__(
            offsets,
            indices,
            num_src_nodes,
            map_csc_to_coo=map_csc_to_coo,
            rev_offsets=rev_offsets,
            rev_indices=rev_indices,
            map_rev_to_coo=map_rev_to_coo,
            map_dst_to_src=map_dst_to_src,
            dst_max_in_degree=dst_max_in_degree,
            reverse_graph_bwd=reverse_graph_bwd,
            is_bipartite=is_bipartite,
        )

    def _build_graph(self):
        if self.is_bipartite:
            src_node_types, dst_node_types = self.node_types
            self.graph_csc = pylibcugraphops.make_bipartite_csc_hg(
                self.offsets,
                self.indices,
                self.num_src_nodes,
                self.num_node_types,
                self.num_edge_types,
                src_node_types,
                dst_node_types,
                self.edge_types,
                self.map_csc_to_coo,
                self.rev_offsets,
                self.rev_indices,
                self.map_rev_to_coo,
            )

        else:
            self.graph_csc = pylibcugraphops.make_csc_hg(
                self.offsets,
                self.indices,
                self.num_src_nodes,
                self.num_node_types,
                self.num_edge_types,
                self.node_types,
                self.edge_types,
                self.map_csc_to_coo,
                self.rev_offsets,
                self.rev_indices,
                self.map_rev_to_coo,
            )

    def to(self, *args, **kwargs):
        self.edge_types = self.edge_types.to(*args, **kwargs).to(torch.int32)
        if self.node_types is not None:
            if isinstance(self.node_types, Tensor):
                self.node_types = self.node_types.to(*args, **kwargs).to(torch.int32)
            else:
                self.node_types = (
                    self.node_types[0].to(*args, **kwargs).to(torch.int32),
                    self.node_types[1].to(*args, **kwargs).to(torch.int32),
                )
        super().to(*args, **kwargs)
