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

import torch
import pylibcugraphops

from typing import Optional, Union, Any, List


class BipartiteCSC:
    def __init__(
        self,
        offsets: torch.Tensor,
        indices: torch.Tensor,
        num_src_nodes: int,
        ef_indices: Optional[torch.Tensor] = None,
        reverse_graph_bwd: bool = False,
    ):
        self.offsets = offsets
        self.indices = indices
        self.num_src_nodes = num_src_nodes
        self.num_dst_nodes = offsets.size(0) - 1
        self.num_edges = indices.size(0)
        self.ef_indices = ef_indices
        self.rev_offsets = None
        self.rev_indices = None
        self.rev_edge_ids = None

        self._fwd_graph = pylibcugraphops.make_bipartite_csc(
            offsets, indices, num_src_nodes, ef_indices
        )

        if reverse_graph_bwd:
            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.rev_edge_ids = torch.empty_like(self.indices)
            self._bwd_graph = pylibcugraphops.make_bipartite_csc_csr(
                self.offsets,
                self.indices,
                self.rev_offsets,
                self.rev_indices,
                self.num_src_nodes,
                self.ef_indices,
                self.rev_edge_ids,
            )
            stream = torch.cuda.current_stream().cuda_stream
            node_counts = torch.empty(
                self.num_src_nodes + 1,
                dtype=self.offsets.dtype,
                device=self.offsets.device,
            )
            cub_workspace = None
            cub_workspace_size = pylibcugraphops.reverse_graph(
                self._bwd_graph, node_counts, cub_workspace, 0, stream
            )
            cub_workspace = torch.empty(
                cub_workspace_size, dtype=torch.int8, device=self.offsets.device
            )
            pylibcugraphops.reverse_graph(
                self._bwd_graph, node_counts, cub_workspace, cub_workspace_size, stream
            )
        else:
            self._bwd_graph = self._fwd_graph


class SampledCSC:
    def __init__(
        self,
        offsets: torch.Tensor,
        indices: torch.Tensor,
        max_num_neighbors: int,
        num_src_nodes: int,
    ):
        self.offsets = offsets
        self.indices = indices

        self.num_src_nodes = num_src_nodes
        self.num_dst_nodes = offsets.size(0) - 1
        self.num_edges = indices.size(0)
        self.max_num_neighbors = max_num_neighbors
        self.out_node_ids = torch.arange(
            self.num_dst_nodes, device=offsets.device, dtype=offsets.dtype
        )

        self._fwd_graph = pylibcugraphops.make_mfg_csr(
            self.out_node_ids, offsets, indices, max_num_neighbors, num_src_nodes
        )
        self._bwd_graph = self._fwd_graph


class SampledHeteroCSC:
    def __init__(
        self,
        offsets: torch.Tensor,
        indices: torch.Tensor,
        edge_types: torch.Tensor,
        max_num_neighbors: int,
        num_src_nodes: int,
        num_edge_types: int,
    ):
        self.offsets = offsets
        self.indices = indices

        self.num_src_nodes = num_src_nodes
        self.num_dst_nodes = offsets.size(0) - 1
        self.num_edges = indices.size(0)
        self.max_num_neighbors = max_num_neighbors
        self.out_node_ids = torch.arange(
            self.num_dst_nodes, device=offsets.device, dtype=offsets.dtype
        )

        self.num_edge_types = num_edge_types
        self.edge_types = edge_types

        self._fwd_graph = pylibcugraphops.make_mfg_csr_hg(
            self.out_node_ids,
            self.offsets,
            self.indices,
            self.max_num_neighbors,
            self.num_src_nodes,
            0,
            self.num_edge_types,
            None,
            None,
            self.edge_types,
        )
        self._bwd_graph = self._fwd_graph


class StaticCSC:
    def __init__(
        self,
        offsets: torch.Tensor,
        indices: torch.Tensor,
        ef_indices: Optional[torch.Tensor] = None,
        rev_edge_ids: Optional[torch.Tensor] = None,
        reverse_graph_bwd: bool = False,
    ):
        self.offsets = offsets
        self.indices = indices
        self.num_nodes = offsets.size(0) - 1
        self.num_edges = indices.size(0)

        self.ef_indices = ef_indices
        self.rev_offsets = None
        self.rev_indices = None
        self.rev_edge_ids = rev_edge_ids

        self._fwd_graph = pylibcugraphops.make_fg_csr(
            self.offsets, self.indices, self.ef_indices, self.rev_edge_ids
        )

        if reverse_graph_bwd:
            self.rev_offsets = torch.empty_like(self.offsets)
            self.rev_indices = torch.empty_like(self.indices)
            self.rev_edge_ids = torch.empty_like(self.indices)
            self._bwd_graph = pylibcugraphops.make_fg_csr_rev(
                self.offsets,
                self.indices,
                self.rev_offsets,
                self.rev_indices,
                self.ef_indices,
                self.rev_edge_ids,
            )
            stream = torch.cuda.current_stream().cuda_stream
            node_counts = torch.empty_like(self.offsets)
            cub_workspace = None
            cub_workspace_size = pylibcugraphops.reverse_graph(
                self._bwd_graph, node_counts, cub_workspace, 0, stream
            )
            cub_workspace = torch.empty(
                cub_workspace_size, dtype=torch.int8, device=self.offsets.device
            )
            pylibcugraphops.reverse_graph(
                self._bwd_graph, node_counts, cub_workspace, cub_workspace_size, stream
            )

        else:
            self._bwd_graph = self._fwd_graph


class StaticHeteroCSC:
    def __init__(
        self,
        offsets: torch.Tensor,
        indices: torch.Tensor,
        edge_types: torch.Tensor,
        num_edge_types: int,
        ef_indices: Optional[torch.Tensor] = None,
        rev_edge_ids: Optional[torch.Tensor] = None,
    ):
        self.offsets = offsets
        self.indices = indices
        self.num_nodes = offsets.size(0) - 1
        self.num_edges = indices.size(0)

        self.num_edge_types = num_edge_types
        self.edge_types = edge_types

        self.ef_indices = ef_indices
        self.rev_edge_ids = rev_edge_ids

        self._fwd_graph = pylibcugraphops.make_fg_csr_hg(
            self.offsets,
            self.indices,
            0,
            self.num_edge_types,
            None,
            self.edge_types,
            self.ef_indices,
            self.rev_edge_ids,
        )
        self._bwd_graph = self._fwd_graph


class UnsupportedGraphError(ValueError):
    def __init__(self, supported_types: Union[Any, List], graph_instance: Any):
        if isinstance(supported_types, list):
            supported_types = [s.__name__ for s in supported_types]
            supported_types = ", or ".join(supported_types)
            msg = (f"expected graph of type {supported_types}, "
                   f"got instance of type {graph_instance.__class__.__name__}.")
        else:
            msg = (f"expected graph of type {supported_types.__name__}, "
                   f"got instance of type {graph_instance.__class__.__name__}.")
        super().__init__(msg)
