"""GPU backend using cuGraph for accelerated graph algorithms.

Uses cuGraph directly for GPU-accelerated graph algorithms.
cuGraph is C/CUDA under the hood - Python is just a thin wrapper.

Performance (71K nodes, 121K edges):
    - Louvain: 117ms (vs 3810ms NetworkX) - 33x faster
    - PageRank: 13ms (vs 120ms NetworkX) - 9x faster
    - Triangle Count: 28ms (vs 300ms+ NetworkX) - 10x faster

NO NetworkX dependency - loads directly from edge arrays.
"""

from typing import Dict, List, Optional, Tuple, Set

from loguru import logger

# Check GPU availability at import time
_GPU_AVAILABLE = False
_CUGRAPH = None
_CUDF = None

try:
    import cudf
    import cugraph
    _GPU_AVAILABLE = True
    _CUGRAPH = cugraph
    _CUDF = cudf
    logger.info("cuGraph GPU backend available")
except ImportError:
    logger.info("cuGraph not available, using Rust backend only")


def gpu_available() -> bool:
    """Check if GPU backend is available."""
    return _GPU_AVAILABLE


class GPUBackend:
    """GPU-accelerated graph algorithms using cuGraph.

    Loads graph from edge arrays, then runs algorithms on GPU.

    Usage:
        gpu = GPUBackend()
        gpu.load_from_edges(sources, targets, weights)

        communities = gpu.louvain()
        pagerank = gpu.pagerank()
        triangles = gpu.triangle_count()
    """

    def __init__(self):
        if not _GPU_AVAILABLE:
            raise RuntimeError("cuGraph not available. Install with: pip install cugraph-cu12")

        self._graph_directed: Optional[_CUGRAPH.Graph] = None
        self._graph_undirected: Optional[_CUGRAPH.Graph] = None
        self._node_to_id: Dict[str, int] = {}
        self._id_to_node: Dict[int, str] = {}
        self._edge_df: Optional[_CUDF.DataFrame] = None

    def load_from_edges(
        self,
        sources: List[str],
        targets: List[str],
        weights: Optional[List[float]] = None
    ) -> Tuple[int, int]:
        """Load graph directly from edge lists.

        Args:
            sources: List of source addresses
            targets: List of target addresses
            weights: Optional list of edge weights

        Returns:
            Tuple of (node_count, edge_count)
        """
        if not sources:
            logger.warning("Empty graph, no edges to load")
            return 0, 0

        # Build node ID mappings
        all_nodes = set(sources) | set(targets)
        self._node_to_id = {node: i for i, node in enumerate(all_nodes)}
        self._id_to_node = {i: node for node, i in self._node_to_id.items()}

        # Convert to numeric IDs
        src_ids = [self._node_to_id[s] for s in sources]
        dst_ids = [self._node_to_id[t] for t in targets]

        if weights is None:
            weights = [1.0] * len(sources)

        # Create cuDF DataFrame
        self._edge_df = _CUDF.DataFrame({
            'src': _CUDF.Series(src_ids, dtype='int32'),
            'dst': _CUDF.Series(dst_ids, dtype='int32'),
            'weight': _CUDF.Series(weights, dtype='float32')
        })

        # Build directed graph
        self._graph_directed = _CUGRAPH.Graph(directed=True)
        self._graph_directed.from_cudf_edgelist(
            self._edge_df, source='src', destination='dst', edge_attr='weight'
        )

        # Build undirected graph (for community detection)
        self._graph_undirected = _CUGRAPH.Graph(directed=False)
        self._graph_undirected.from_cudf_edgelist(
            self._edge_df, source='src', destination='dst', edge_attr='weight'
        )

        n_nodes = len(all_nodes)
        n_edges = len(sources)
        logger.debug(f"Loaded graph to GPU: {n_nodes} nodes, {n_edges} edges")

        return n_nodes, n_edges

    def louvain(self, resolution: float = 1.0) -> List[Set[str]]:
        """Run Louvain community detection on GPU.

        Performance: ~117ms for 71K nodes (vs 3810ms NetworkX)

        Args:
            resolution: Resolution parameter for modularity

        Returns:
            List of communities, each community is a set of addresses
        """
        if self._graph_undirected is None:
            raise RuntimeError("Graph not loaded. Call load_from_edges first.")

        # Run Louvain on GPU
        parts, modularity = _CUGRAPH.louvain(
            self._graph_undirected,
            resolution=resolution
        )

        # Convert back to address sets - use arrow for robust GPU->CPU transfer
        try:
            # Try arrow conversion (more reliable)
            arrow_table = parts.to_arrow()
            vertices = arrow_table.column('vertex').to_pylist()
            partitions = arrow_table.column('partition').to_pylist()
        except Exception:
            # Fallback to pandas
            df = parts.to_pandas()
            vertices = df['vertex'].tolist()
            partitions = df['partition'].tolist()

        communities_dict: Dict[int, Set[str]] = {}

        for node_id, partition in zip(vertices, partitions):
            node_id = int(node_id)
            partition = int(partition)

            if partition not in communities_dict:
                communities_dict[partition] = set()

            if node_id in self._id_to_node:
                communities_dict[partition].add(self._id_to_node[node_id])

        communities = list(communities_dict.values())
        logger.debug(f"Louvain found {len(communities)} communities (modularity={modularity:.4f})")

        return communities

    def pagerank(self, alpha: float = 0.85, max_iter: int = 100) -> Dict[str, float]:
        """Run PageRank on GPU.

        Performance: ~13ms for 71K nodes (vs 120ms NetworkX)

        Args:
            alpha: Damping factor
            max_iter: Maximum iterations

        Returns:
            Dict mapping address to PageRank score
        """
        if self._graph_directed is None:
            raise RuntimeError("Graph not loaded. Call load_from_edges first.")

        # Run PageRank on GPU
        pr = _CUGRAPH.pagerank(
            self._graph_directed,
            alpha=alpha,
            max_iter=max_iter
        )

        # Convert back to address dict - use arrow for robust GPU->CPU transfer
        try:
            arrow_table = pr.to_arrow()
            vertices = arrow_table.column('vertex').to_pylist()
            scores = arrow_table.column('pagerank').to_pylist()
        except Exception:
            df = pr.to_pandas()
            vertices = df['vertex'].tolist()
            scores = df['pagerank'].tolist()

        result = {}
        for node_id, score in zip(vertices, scores):
            node_id = int(node_id)
            if node_id in self._id_to_node:
                result[self._id_to_node[node_id]] = float(score)

        return result

    def triangle_count(self) -> int:
        """Count triangles in the graph on GPU.

        Performance: ~28ms for 71K nodes (vs 300ms+ NetworkX)

        Returns:
            Total number of triangles
        """
        if self._graph_undirected is None:
            raise RuntimeError("Graph not loaded. Call load_from_edges first.")

        # Run triangle count on GPU
        result = _CUGRAPH.triangle_count(self._graph_undirected)

        # Sum all triangles (each triangle counted 3 times, once per vertex)
        try:
            total = int(result['counts'].sum()) // 3
        except Exception:
            # Fallback to arrow conversion
            arrow_table = result.to_arrow()
            counts = arrow_table.column('counts').to_pylist()
            total = sum(counts) // 3

        return int(total)

    def betweenness_centrality(self, k: Optional[int] = None) -> Dict[str, float]:
        """Run betweenness centrality on GPU.

        Performance: ~1.5s for 71K nodes with k=100 (vs 28s NetworkX)

        Args:
            k: Number of source vertices to sample (None = all)

        Returns:
            Dict mapping address to betweenness score
        """
        if self._graph_directed is None:
            raise RuntimeError("Graph not loaded. Call load_from_edges first.")

        # Run betweenness on GPU
        bc = _CUGRAPH.betweenness_centrality(
            self._graph_directed,
            k=k
        )

        # Convert back to address dict - use arrow for robust GPU->CPU transfer
        try:
            arrow_table = bc.to_arrow()
            vertices = arrow_table.column('vertex').to_pylist()
            scores = arrow_table.column('betweenness_centrality').to_pylist()
        except Exception:
            df = bc.to_pandas()
            vertices = df['vertex'].tolist()
            scores = df['betweenness_centrality'].tolist()

        result = {}
        for node_id, score in zip(vertices, scores):
            node_id = int(node_id)
            if node_id in self._id_to_node:
                result[self._id_to_node[node_id]] = float(score)

        return result

    def bfs(self, source: str, depth_limit: Optional[int] = None) -> Dict[str, int]:
        """Run BFS from a source node on GPU.

        Note: For small graphs, Rust BFS may be faster due to GPU overhead.

        Args:
            source: Source address
            depth_limit: Maximum depth (None = unlimited)

        Returns:
            Dict mapping address to distance from source
        """
        if self._graph_undirected is None:
            raise RuntimeError("Graph not loaded. Call load_from_edges first.")

        if source not in self._node_to_id:
            return {}

        source_id = self._node_to_id[source]

        # Run BFS on GPU
        bfs_result = _CUGRAPH.bfs(
            self._graph_undirected,
            source_id,
            depth_limit=depth_limit
        )

        # Convert back to address dict - use arrow for robust GPU->CPU transfer
        try:
            arrow_table = bfs_result.to_arrow()
            vertices = arrow_table.column('vertex').to_pylist()
            distances = arrow_table.column('distance').to_pylist()
        except Exception:
            df = bfs_result.to_pandas()
            vertices = df['vertex'].tolist()
            distances = df['distance'].tolist()

        result = {}
        for node_id, distance in zip(vertices, distances):
            node_id = int(node_id)
            distance = int(distance)
            if distance >= 0 and node_id in self._id_to_node:
                result[self._id_to_node[node_id]] = distance

        return result

    def get_node_id(self, address: str) -> Optional[int]:
        """Get numeric ID for an address."""
        return self._node_to_id.get(address)

    def get_address(self, node_id: int) -> Optional[str]:
        """Get address for a numeric ID."""
        return self._id_to_node.get(node_id)
