//! Rust graph algorithms and data I/O using petgraph with PyO3 bindings.
//!
//! Implements: SCC, PageRank, Triangle Count, Louvain, BFS, Parallel Path Finding
//! Data I/O: Parquet reading, timestamp extraction

mod data_io;

use hashbrown::{HashMap as HBMap, HashSet as HBSet};
use petgraph::algo::kosaraju_scc;
use petgraph::graph::{DiGraph, UnGraph, NodeIndex};
use petgraph::visit::EdgeRef;
use pyo3::prelude::*;
use pyo3::types::PyModule;
use pyo3::Bound;
use rayon::prelude::*;
use std::collections::{BTreeSet, HashMap, VecDeque};
use std::sync::atomic::{AtomicUsize, Ordering};

/// Graph wrapper for Python interop
#[pyclass]
pub struct RustGraph {
    directed: DiGraph<String, f64>,
    node_map: HBMap<String, NodeIndex>,
    reverse_map: HBMap<NodeIndex, String>,
}

#[pymethods]
impl RustGraph {
    #[new]
    pub fn new() -> Self {
        RustGraph {
            directed: DiGraph::new(),
            node_map: HBMap::new(),
            reverse_map: HBMap::new(),
        }
    }

    /// Add edges from source/target/weight lists
    pub fn add_edges(&mut self, sources: Vec<String>, targets: Vec<String>, weights: Vec<f64>) {
        for ((src, tgt), weight) in sources.into_iter().zip(targets).zip(weights) {
            let src_idx = *self.node_map.entry(src.clone()).or_insert_with(|| {
                let idx = self.directed.add_node(src.clone());
                self.reverse_map.insert(idx, src);
                idx
            });

            let tgt_idx = *self.node_map.entry(tgt.clone()).or_insert_with(|| {
                let idx = self.directed.add_node(tgt.clone());
                self.reverse_map.insert(idx, tgt);
                idx
            });

            self.directed.add_edge(src_idx, tgt_idx, weight);
        }
    }

    /// Get node and edge counts
    pub fn stats(&self) -> (usize, usize) {
        (self.directed.node_count(), self.directed.edge_count())
    }

    /// Get node count
    pub fn node_count(&self) -> usize {
        self.directed.node_count()
    }

    /// Get edge count
    pub fn edge_count(&self) -> usize {
        self.directed.edge_count()
    }

    /// Get all node names
    pub fn nodes(&self) -> Vec<String> {
        self.directed
            .node_indices()
            .filter_map(|idx| self.reverse_map.get(&idx).cloned())
            .collect()
    }

    /// Get all edges as (source, target, weight) tuples
    pub fn edges(&self) -> Vec<(String, String, f64)> {
        self.directed
            .edge_references()
            .filter_map(|edge| {
                let src = self.reverse_map.get(&edge.source())?;
                let tgt = self.reverse_map.get(&edge.target())?;
                Some((src.clone(), tgt.clone(), *edge.weight()))
            })
            .collect()
    }

    /// Get in-degree for a node
    pub fn in_degree(&self, node: String) -> usize {
        self.node_map
            .get(&node)
            .map(|&idx| self.directed.edges_directed(idx, petgraph::Direction::Incoming).count())
            .unwrap_or(0)
    }

    /// Get out-degree for a node
    pub fn out_degree(&self, node: String) -> usize {
        self.node_map
            .get(&node)
            .map(|&idx| self.directed.edges(idx).count())
            .unwrap_or(0)
    }

    /// Get in-edges for a node as (source, weight) tuples
    pub fn in_edges(&self, node: String) -> Vec<(String, f64)> {
        self.node_map
            .get(&node)
            .map(|&idx| {
                self.directed
                    .edges_directed(idx, petgraph::Direction::Incoming)
                    .filter_map(|edge| {
                        self.reverse_map
                            .get(&edge.source())
                            .map(|src| (src.clone(), *edge.weight()))
                    })
                    .collect()
            })
            .unwrap_or_default()
    }

    /// Get out-edges for a node as (target, weight) tuples
    pub fn out_edges(&self, node: String) -> Vec<(String, f64)> {
        self.node_map
            .get(&node)
            .map(|&idx| {
                self.directed
                    .edges(idx)
                    .filter_map(|edge| {
                        self.reverse_map
                            .get(&edge.target())
                            .map(|tgt| (tgt.clone(), *edge.weight()))
                    })
                    .collect()
            })
            .unwrap_or_default()
    }

    /// Get all neighbors (undirected view) for a node
    pub fn neighbors(&self, node: String) -> Vec<String> {
        let Some(&idx) = self.node_map.get(&node) else {
            return vec![];
        };

        let mut neighbors: HBSet<NodeIndex> = HBSet::new();

        // Outgoing edges
        for edge in self.directed.edges(idx) {
            neighbors.insert(edge.target());
        }

        // Incoming edges
        for edge in self.directed.edges_directed(idx, petgraph::Direction::Incoming) {
            neighbors.insert(edge.source());
        }

        neighbors
            .into_iter()
            .filter_map(|n| self.reverse_map.get(&n).cloned())
            .collect()
    }

    /// Check if node exists in graph
    pub fn has_node(&self, node: String) -> bool {
        self.node_map.contains_key(&node)
    }

    /// Get edge weight between two nodes (None if no edge)
    pub fn get_edge_weight(&self, source: String, target: String) -> Option<f64> {
        let src_idx = self.node_map.get(&source)?;
        let tgt_idx = self.node_map.get(&target)?;
        self.directed.find_edge(*src_idx, *tgt_idx)
            .map(|edge_idx| *self.directed.edge_weight(edge_idx).unwrap_or(&0.0))
    }

    /// Create a subgraph containing only the specified nodes
    pub fn subgraph(&self, nodes: Vec<String>) -> RustGraph {
        let node_set: HBSet<&String> = nodes.iter().collect();

        let mut new_graph = RustGraph::new();
        let mut sources = Vec::new();
        let mut targets = Vec::new();
        let mut weights = Vec::new();

        for edge in self.directed.edge_references() {
            let src = self.reverse_map.get(&edge.source());
            let tgt = self.reverse_map.get(&edge.target());

            if let (Some(src_name), Some(tgt_name)) = (src, tgt) {
                if node_set.contains(src_name) && node_set.contains(tgt_name) {
                    sources.push(src_name.clone());
                    targets.push(tgt_name.clone());
                    weights.push(*edge.weight());
                }
            }
        }

        // Also ensure isolated nodes are added
        for node in &nodes {
            if self.node_map.contains_key(node) && !new_graph.node_map.contains_key(node) {
                let idx = new_graph.directed.add_node(node.clone());
                new_graph.node_map.insert(node.clone(), idx);
                new_graph.reverse_map.insert(idx, node.clone());
            }
        }

        if !sources.is_empty() {
            new_graph.add_edges(sources, targets, weights);
        }

        new_graph
    }

    /// Calculate graph density
    pub fn density(&self) -> f64 {
        let n = self.directed.node_count();
        let m = self.directed.edge_count();

        if n < 2 {
            return 0.0;
        }

        // For directed graph: density = m / (n * (n - 1))
        m as f64 / (n * (n - 1)) as f64
    }

    /// Calculate clustering coefficient for a single node (local clustering coefficient)
    pub fn node_clustering(&self, node: String) -> f64 {
        let Some(&idx) = self.node_map.get(&node) else {
            return 0.0;
        };

        // Get neighbors (undirected view)
        let mut neighbor_set: HBSet<NodeIndex> = HBSet::new();
        for edge in self.directed.edges(idx) {
            neighbor_set.insert(edge.target());
        }
        for edge in self.directed.edges_directed(idx, petgraph::Direction::Incoming) {
            neighbor_set.insert(edge.source());
        }

        let neighbors: Vec<NodeIndex> = neighbor_set.into_iter().collect();
        let k = neighbors.len();

        if k < 2 {
            return 0.0;
        }

        // Count edges between neighbors
        let mut edges_between = 0usize;
        for i in 0..neighbors.len() {
            for j in (i + 1)..neighbors.len() {
                let ni = neighbors[i];
                let nj = neighbors[j];

                // Check for edge in either direction
                if self.directed.find_edge(ni, nj).is_some()
                    || self.directed.find_edge(nj, ni).is_some()
                {
                    edges_between += 1;
                }
            }
        }

        // Clustering coefficient = 2 * edges_between / (k * (k - 1))
        2.0 * edges_between as f64 / (k * (k - 1)) as f64
    }

    /// Calculate clustering coefficient for all nodes
    pub fn clustering_coefficient(&self) -> HashMap<String, f64> {
        self.directed
            .node_indices()
            .filter_map(|idx| {
                let name = self.reverse_map.get(&idx)?;
                let cc = self.node_clustering(name.clone());
                Some((name.clone(), cc))
            })
            .collect()
    }

    /// Calculate average clustering coefficient
    pub fn average_clustering(&self) -> f64 {
        let n = self.directed.node_count();
        if n == 0 {
            return 0.0;
        }

        let total: f64 = self
            .directed
            .node_indices()
            .filter_map(|idx| self.reverse_map.get(&idx))
            .map(|name| self.node_clustering(name.clone()))
            .sum();

        total / n as f64
    }

    /// Calculate closeness centrality for all nodes (using BFS distances)
    /// Parallel implementation using rayon
    pub fn closeness_centrality(&self) -> HashMap<String, f64> {
        let n = self.directed.node_count();
        if n < 2 {
            return HashMap::new();
        }

        // Build undirected adjacency once
        let adj: HBMap<NodeIndex, Vec<NodeIndex>> = {
            let mut adj: HBMap<NodeIndex, Vec<NodeIndex>> = HBMap::new();
            for edge in self.directed.edge_references() {
                let s = edge.source();
                let t = edge.target();
                adj.entry(s).or_default().push(t);
                adj.entry(t).or_default().push(s);
            }
            adj
        };

        // Parallel closeness calculation
        let nodes: Vec<NodeIndex> = self.directed.node_indices().collect();
        let results: Vec<(NodeIndex, f64)> = nodes
            .par_iter()
            .map(|&source_idx| {
                // BFS from this node
                let mut distances: HBMap<NodeIndex, usize> = HBMap::new();
                let mut queue: VecDeque<(NodeIndex, usize)> = VecDeque::new();

                distances.insert(source_idx, 0);
                queue.push_back((source_idx, 0));

                while let Some((node, dist)) = queue.pop_front() {
                    if let Some(neighbors) = adj.get(&node) {
                        for &neighbor in neighbors {
                            if !distances.contains_key(&neighbor) {
                                distances.insert(neighbor, dist + 1);
                                queue.push_back((neighbor, dist + 1));
                            }
                        }
                    }
                }

                // Closeness = (reachable - 1) / sum(distances)
                let reachable = distances.len();
                if reachable < 2 {
                    return (source_idx, 0.0);
                }

                let total_dist: usize = distances.values().sum();
                let closeness = (reachable - 1) as f64 / total_dist as f64;

                (source_idx, closeness)
            })
            .collect();

        results
            .into_iter()
            .filter_map(|(idx, cc)| {
                self.reverse_map.get(&idx).map(|name| (name.clone(), cc))
            })
            .collect()
    }

    /// Calculate k-core number for each node
    /// Parallel degeneracy algorithm
    pub fn core_number(&self) -> HashMap<String, usize> {
        let n = self.directed.node_count();
        if n == 0 {
            return HashMap::new();
        }

        // Work with undirected degree
        let mut degree: HBMap<NodeIndex, usize> = HBMap::new();
        for idx in self.directed.node_indices() {
            let mut neighbors: HBSet<NodeIndex> = HBSet::new();
            for edge in self.directed.edges(idx) {
                neighbors.insert(edge.target());
            }
            for edge in self.directed.edges_directed(idx, petgraph::Direction::Incoming) {
                neighbors.insert(edge.source());
            }
            degree.insert(idx, neighbors.len());
        }

        // Build adjacency
        let mut adj: HBMap<NodeIndex, HBSet<NodeIndex>> = HBMap::new();
        for edge in self.directed.edge_references() {
            let s = edge.source();
            let t = edge.target();
            adj.entry(s).or_default().insert(t);
            adj.entry(t).or_default().insert(s);
        }

        let mut core: HBMap<NodeIndex, usize> = HBMap::new();
        let mut remaining: HBSet<NodeIndex> = self.directed.node_indices().collect();

        // Iteratively remove minimum degree nodes
        while !remaining.is_empty() {
            // Find minimum degree among remaining nodes
            let min_deg = remaining
                .iter()
                .map(|&idx| degree.get(&idx).copied().unwrap_or(0))
                .min()
                .unwrap_or(0);

            // Find all nodes with minimum degree
            let min_nodes: Vec<NodeIndex> = remaining
                .iter()
                .filter(|&&idx| degree.get(&idx).copied().unwrap_or(0) == min_deg)
                .copied()
                .collect();

            for node in min_nodes {
                core.insert(node, min_deg);
                remaining.remove(&node);

                // Update degrees of neighbors
                if let Some(neighbors) = adj.get(&node) {
                    for &neighbor in neighbors {
                        if remaining.contains(&neighbor) {
                            if let Some(d) = degree.get_mut(&neighbor) {
                                *d = d.saturating_sub(1);
                            }
                        }
                    }
                }
            }
        }

        core.into_iter()
            .filter_map(|(idx, k)| {
                self.reverse_map.get(&idx).map(|name| (name.clone(), k))
            })
            .collect()
    }

    /// Detect fan-in motifs: nodes receiving from many sources
    /// Returns: Vec<(center_node, sources, total_volume, edge_count)>
    #[pyo3(signature = (min_sources=5, min_volume=0.0))]
    pub fn detect_fan_in_motifs(
        &self,
        min_sources: usize,
        min_volume: f64,
    ) -> Vec<(String, Vec<String>, f64, usize)> {
        self.directed
            .node_indices()
            .filter_map(|idx| {
                let in_edges: Vec<_> = self
                    .directed
                    .edges_directed(idx, petgraph::Direction::Incoming)
                    .collect();

                if in_edges.len() < min_sources {
                    return None;
                }

                let total_volume: f64 = in_edges.iter().map(|e| e.weight()).sum();
                if total_volume < min_volume {
                    return None;
                }

                let center = self.reverse_map.get(&idx)?.clone();
                let sources: Vec<String> = in_edges
                    .iter()
                    .filter_map(|e| self.reverse_map.get(&e.source()).cloned())
                    .collect();

                Some((center, sources, total_volume, in_edges.len()))
            })
            .collect()
    }

    /// Detect fan-out motifs: nodes sending to many targets
    /// Returns: Vec<(center_node, targets, total_volume, edge_count)>
    #[pyo3(signature = (min_targets=5, min_volume=0.0))]
    pub fn detect_fan_out_motifs(
        &self,
        min_targets: usize,
        min_volume: f64,
    ) -> Vec<(String, Vec<String>, f64, usize)> {
        self.directed
            .node_indices()
            .filter_map(|idx| {
                let out_edges: Vec<_> = self.directed.edges(idx).collect();

                if out_edges.len() < min_targets {
                    return None;
                }

                let total_volume: f64 = out_edges.iter().map(|e| e.weight()).sum();
                if total_volume < min_volume {
                    return None;
                }

                let center = self.reverse_map.get(&idx)?.clone();
                let targets: Vec<String> = out_edges
                    .iter()
                    .filter_map(|e| self.reverse_map.get(&e.target()).cloned())
                    .collect();

                Some((center, targets, total_volume, out_edges.len()))
            })
            .collect()
    }

    /// Detect threshold evasion patterns
    /// Finds edges with amounts just below a threshold
    /// Returns: Vec<(source, target, amount, distance_to_threshold)>
    #[pyo3(signature = (threshold, tolerance_pct=0.1))]
    pub fn detect_threshold_evasion(
        &self,
        threshold: f64,
        tolerance_pct: f64,
    ) -> Vec<(String, String, f64, f64)> {
        let tolerance = threshold * tolerance_pct;
        let lower_bound = threshold - tolerance;

        self.directed
            .edge_references()
            .filter_map(|edge| {
                let amount = *edge.weight();
                if amount >= lower_bound && amount < threshold {
                    let src = self.reverse_map.get(&edge.source())?;
                    let tgt = self.reverse_map.get(&edge.target())?;
                    let distance = threshold - amount;
                    Some((src.clone(), tgt.clone(), amount, distance))
                } else {
                    None
                }
            })
            .collect()
    }

    /// Get all edges between two sets of nodes (for smurfing detection)
    /// Returns edges where source is in sources and target is in targets
    pub fn edges_between(
        &self,
        sources: Vec<String>,
        targets: Vec<String>,
    ) -> Vec<(String, String, f64)> {
        let source_set: HBSet<NodeIndex> = sources
            .iter()
            .filter_map(|s| self.node_map.get(s).copied())
            .collect();
        let target_set: HBSet<NodeIndex> = targets
            .iter()
            .filter_map(|t| self.node_map.get(t).copied())
            .collect();

        self.directed
            .edge_references()
            .filter_map(|edge| {
                if source_set.contains(&edge.source()) && target_set.contains(&edge.target()) {
                    let src = self.reverse_map.get(&edge.source())?;
                    let tgt = self.reverse_map.get(&edge.target())?;
                    Some((src.clone(), tgt.clone(), *edge.weight()))
                } else {
                    None
                }
            })
            .collect()
    }

    /// Get total in-volume for a node
    pub fn in_volume(&self, node: String) -> f64 {
        self.node_map
            .get(&node)
            .map(|&idx| {
                self.directed
                    .edges_directed(idx, petgraph::Direction::Incoming)
                    .map(|e| e.weight())
                    .sum()
            })
            .unwrap_or(0.0)
    }

    /// Get total out-volume for a node
    pub fn out_volume(&self, node: String) -> f64 {
        self.node_map
            .get(&node)
            .map(|&idx| self.directed.edges(idx).map(|e| e.weight()).sum())
            .unwrap_or(0.0)
    }

    /// Get in-degree and out-degree for all nodes
    /// Returns: HashMap<node_name, (in_degree, out_degree)>
    pub fn degree_map(&self) -> HashMap<String, (usize, usize)> {
        self.directed
            .node_indices()
            .filter_map(|idx| {
                let name = self.reverse_map.get(&idx)?;
                let in_deg = self
                    .directed
                    .edges_directed(idx, petgraph::Direction::Incoming)
                    .count();
                let out_deg = self.directed.edges(idx).count();
                Some((name.clone(), (in_deg, out_deg)))
            })
            .collect()
    }

    /// Get volume statistics for all nodes
    /// Returns: HashMap<node_name, (in_volume, out_volume)>
    pub fn volume_map(&self) -> HashMap<String, (f64, f64)> {
        self.directed
            .node_indices()
            .filter_map(|idx| {
                let name = self.reverse_map.get(&idx)?;
                let in_vol: f64 = self
                    .directed
                    .edges_directed(idx, petgraph::Direction::Incoming)
                    .map(|e| e.weight())
                    .sum();
                let out_vol: f64 = self.directed.edges(idx).map(|e| e.weight()).sum();
                Some((name.clone(), (in_vol, out_vol)))
            })
            .collect()
    }

    /// Strongly Connected Components using Kosaraju's algorithm
    pub fn strongly_connected_components(&self) -> Vec<Vec<String>> {
        let sccs = kosaraju_scc(&self.directed);
        sccs.into_iter()
            .map(|component| {
                component
                    .into_iter()
                    .filter_map(|idx| self.reverse_map.get(&idx).cloned())
                    .collect()
            })
            .collect()
    }

    /// BFS shortest path distances from source (treats graph as undirected)
    /// Returns HashMap of node_name -> distance from source
    #[pyo3(signature = (source, cutoff=None))]
    pub fn bfs(&self, source: String, cutoff: Option<usize>) -> HashMap<String, usize> {
        let Some(&source_idx) = self.node_map.get(&source) else {
            return HashMap::new();
        };

        let max_dist = cutoff.unwrap_or(usize::MAX);
        let mut distances: HBMap<NodeIndex, usize> = HBMap::new();
        let mut queue: VecDeque<(NodeIndex, usize)> = VecDeque::new();

        distances.insert(source_idx, 0);
        queue.push_back((source_idx, 0));

        // Build undirected adjacency for BFS
        let mut adj: HBMap<NodeIndex, Vec<NodeIndex>> = HBMap::new();
        for edge in self.directed.edge_references() {
            let s = edge.source();
            let t = edge.target();
            adj.entry(s).or_default().push(t);
            adj.entry(t).or_default().push(s);
        }

        while let Some((node, dist)) = queue.pop_front() {
            if dist >= max_dist {
                continue;
            }

            if let Some(neighbors) = adj.get(&node) {
                for &neighbor in neighbors {
                    if !distances.contains_key(&neighbor) {
                        let new_dist = dist + 1;
                        distances.insert(neighbor, new_dist);
                        if new_dist < max_dist {
                            queue.push_back((neighbor, new_dist));
                        }
                    }
                }
            }
        }

        // Convert to string keys
        distances
            .into_iter()
            .filter_map(|(idx, dist)| {
                self.reverse_map.get(&idx).map(|name| (name.clone(), dist))
            })
            .collect()
    }

    /// Multi-source BFS - PARALLEL proximity search from risk addresses
    /// Returns: (source, target, distance) tuples
    ///
    /// Key for AML: Find all addresses within N hops of known bad actors
    /// Uses rayon for parallel BFS from each source
    #[pyo3(signature = (sources, cutoff=None))]
    pub fn multi_source_bfs(&self, sources: Vec<String>, cutoff: Option<usize>) -> Vec<(String, String, usize)> {
        let max_dist = cutoff.unwrap_or(3);  // Default 3 hops for risk proximity

        // Build undirected adjacency once (shared across all BFS)
        let adj: HBMap<NodeIndex, Vec<NodeIndex>> = {
            let mut adj: HBMap<NodeIndex, Vec<NodeIndex>> = HBMap::new();
            for edge in self.directed.edge_references() {
                let s = edge.source();
                let t = edge.target();
                adj.entry(s).or_default().push(t);
                adj.entry(t).or_default().push(s);
            }
            adj
        };

        // Convert sources to indices
        let source_indices: Vec<(String, NodeIndex)> = sources
            .into_iter()
            .filter_map(|s| self.node_map.get(&s).map(|&idx| (s, idx)))
            .collect();

        // PARALLEL BFS from each source using rayon
        let results: Vec<Vec<(String, String, usize)>> = source_indices
            .par_iter()
            .map(|(source_name, source_idx)| {
                let mut local_results = Vec::new();
                let mut distances: HBMap<NodeIndex, usize> = HBMap::new();
                let mut queue: VecDeque<(NodeIndex, usize)> = VecDeque::new();

                distances.insert(*source_idx, 0);
                queue.push_back((*source_idx, 0));

                while let Some((node, dist)) = queue.pop_front() {
                    if dist >= max_dist {
                        continue;
                    }

                    if let Some(neighbors) = adj.get(&node) {
                        for &neighbor in neighbors {
                            if !distances.contains_key(&neighbor) {
                                let new_dist = dist + 1;
                                distances.insert(neighbor, new_dist);
                                if new_dist < max_dist {
                                    queue.push_back((neighbor, new_dist));
                                }
                            }
                        }
                    }
                }

                // Collect results for this source
                for (idx, dist) in distances {
                    if dist > 0 {
                        if let Some(target) = self.reverse_map.get(&idx) {
                            local_results.push((source_name.clone(), target.clone(), dist));
                        }
                    }
                }
                local_results
            })
            .collect();

        results.into_iter().flatten().collect()
    }

    /// PageRank algorithm
    pub fn pagerank(&self, alpha: f64, max_iter: usize, tol: f64) -> HashMap<String, f64> {
        let n = self.directed.node_count();
        if n == 0 {
            return HashMap::new();
        }

        let n_f64 = n as f64;
        let initial = 1.0 / n_f64;

        // Initialize scores
        let mut scores: HBMap<NodeIndex, f64> = self.directed
            .node_indices()
            .map(|idx| (idx, initial))
            .collect();

        // Precompute out-degrees
        let out_degrees: HBMap<NodeIndex, usize> = self.directed
            .node_indices()
            .map(|idx| (idx, self.directed.edges(idx).count()))
            .collect();

        // Iterative PageRank
        for _ in 0..max_iter {
            let mut new_scores: HBMap<NodeIndex, f64> = HBMap::with_capacity(n);
            let mut diff = 0.0f64;

            for node in self.directed.node_indices() {
                let mut sum = 0.0;

                // Sum contributions from incoming edges
                for edge in self.directed.edges_directed(node, petgraph::Direction::Incoming) {
                    let source = edge.source();
                    let out_deg = *out_degrees.get(&source).unwrap_or(&1);
                    if out_deg > 0 {
                        sum += scores[&source] / out_deg as f64;
                    }
                }

                let new_score = (1.0 - alpha) / n_f64 + alpha * sum;
                diff += (new_score - scores[&node]).abs();
                new_scores.insert(node, new_score);
            }

            scores = new_scores;

            if diff < tol {
                break;
            }
        }

        // Convert to string keys (use std HashMap for PyO3)
        scores
            .into_iter()
            .filter_map(|(idx, score)| {
                self.reverse_map.get(&idx).map(|name| (name.clone(), score))
            })
            .collect()
    }

    /// Count triangles in the graph (treating as undirected)
    pub fn triangle_count(&self) -> usize {
        // Build undirected graph
        let mut undirected: UnGraph<(), ()> = UnGraph::new_undirected();
        let mut und_node_map: HBMap<NodeIndex, petgraph::graph::NodeIndex<u32>> = HBMap::new();

        // Add nodes
        for node in self.directed.node_indices() {
            let und_idx = undirected.add_node(());
            und_node_map.insert(node, und_idx);
        }

        // Add edges (deduplicated)
        let mut seen_edges: HBSet<(usize, usize)> = HBSet::new();
        for edge in self.directed.edge_references() {
            let s = edge.source().index();
            let t = edge.target().index();
            let key = if s < t { (s, t) } else { (t, s) };
            if !seen_edges.contains(&key) {
                seen_edges.insert(key);
                let und_s = und_node_map[&edge.source()];
                let und_t = und_node_map[&edge.target()];
                undirected.add_edge(und_s, und_t, ());
            }
        }

        // Build adjacency sets for fast lookup
        let adj: HBMap<petgraph::graph::NodeIndex<u32>, BTreeSet<petgraph::graph::NodeIndex<u32>>> =
            undirected
                .node_indices()
                .map(|n| {
                    let neighbors: BTreeSet<_> = undirected.neighbors(n).collect();
                    (n, neighbors)
                })
                .collect();

        // Count triangles: for each edge (u,v) where u < v, count common neighbors w > v
        let mut count = 0usize;
        for edge in undirected.edge_references() {
            let u = edge.source();
            let v = edge.target();
            let (smaller, larger) = if u.index() < v.index() { (u, v) } else { (v, u) };

            if let (Some(adj_s), Some(adj_l)) = (adj.get(&smaller), adj.get(&larger)) {
                // Count common neighbors greater than larger
                for &w in adj_s.iter() {
                    if w.index() > larger.index() && adj_l.contains(&w) {
                        count += 1;
                    }
                }
            }
        }

        count
    }

    /// Louvain community detection algorithm
    pub fn louvain(&self, resolution: f64) -> Vec<Vec<String>> {
        let n = self.directed.node_count();
        if n == 0 {
            return vec![];
        }

        // Build undirected weighted graph
        let mut edge_weights: HBMap<(usize, usize), f64> = HBMap::new();
        for edge in self.directed.edge_references() {
            let s = edge.source().index();
            let t = edge.target().index();
            let key = if s < t { (s, t) } else { (t, s) };
            *edge_weights.entry(key).or_insert(0.0) += edge.weight().abs();
        }

        let total_weight: f64 = edge_weights.values().sum();
        if total_weight == 0.0 {
            // Each node is its own community
            return self.directed
                .node_indices()
                .filter_map(|idx| self.reverse_map.get(&idx).map(|s| vec![s.clone()]))
                .collect();
        }

        // Node indices as Vec for indexing
        let nodes: Vec<NodeIndex> = self.directed.node_indices().collect();
        let _node_to_idx: HBMap<NodeIndex, usize> = nodes
            .iter()
            .enumerate()
            .map(|(i, &n)| (n, i))
            .collect();

        // Build adjacency with weights
        let mut adj: Vec<Vec<(usize, f64)>> = vec![vec![]; n];
        for (&(s, t), &w) in &edge_weights {
            adj[s].push((t, w));
            adj[t].push((s, w));
        }

        // Calculate node strengths (weighted degrees)
        let strengths: Vec<f64> = adj
            .iter()
            .map(|neighbors| neighbors.iter().map(|(_, w)| w).sum())
            .collect();

        // Initialize: each node in its own community
        let mut community: Vec<usize> = (0..n).collect();
        let mut community_strength: Vec<f64> = strengths.clone();
        let mut community_internal: Vec<f64> = vec![0.0; n];

        // Calculate initial internal weights
        for (&(s, t), &w) in &edge_weights {
            if community[s] == community[t] {
                community_internal[community[s]] += 2.0 * w;
            }
        }

        let m2 = 2.0 * total_weight;
        let mut improved = true;
        let mut iterations = 0;
        let max_iterations = 100;

        while improved && iterations < max_iterations {
            improved = false;
            iterations += 1;

            for i in 0..n {
                let current_comm = community[i];
                let ki = strengths[i];

                // Calculate delta Q for moving to each neighbor's community
                let mut best_comm = current_comm;
                let mut best_delta = 0.0f64;

                // Compute weight to current community
                let mut weight_to_current = 0.0f64;
                for &(j, w) in &adj[i] {
                    if community[j] == current_comm {
                        weight_to_current += w;
                    }
                }

                // Remove node from current community temporarily
                let sigma_current = community_strength[current_comm] - ki;

                // Try each neighboring community
                let mut tried_comms: HBSet<usize> = HBSet::new();
                tried_comms.insert(current_comm);

                for &(j, _) in &adj[i] {
                    let new_comm = community[j];
                    if tried_comms.contains(&new_comm) {
                        continue;
                    }
                    tried_comms.insert(new_comm);

                    // Weight to new community
                    let mut weight_to_new = 0.0f64;
                    for &(k, w) in &adj[i] {
                        if community[k] == new_comm {
                            weight_to_new += w;
                        }
                    }

                    let sigma_new = community_strength[new_comm];

                    // Delta Q = (weight_to_new - weight_to_current) / m
                    //         - resolution * ki * (sigma_new - sigma_current) / (2 * m^2)
                    let delta = (weight_to_new - weight_to_current) / total_weight
                        - resolution * ki * (sigma_new - sigma_current) / (m2 * total_weight);

                    if delta > best_delta {
                        best_delta = delta;
                        best_comm = new_comm;
                    }
                }

                // Move to best community if improvement
                if best_comm != current_comm && best_delta > 1e-10 {
                    // Update community assignments
                    community_strength[current_comm] -= ki;
                    community_strength[best_comm] += ki;
                    community[i] = best_comm;
                    improved = true;
                }
            }
        }

        // Collect communities
        let mut comm_members: HBMap<usize, Vec<String>> = HBMap::new();
        for (i, &c) in community.iter().enumerate() {
            if let Some(name) = self.reverse_map.get(&nodes[i]) {
                comm_members.entry(c).or_default().push(name.clone());
            }
        }

        comm_members.into_values().collect()
    }

    /// Find layering paths in parallel using DFS
    /// Returns: Vec<(path_nodes, edge_weights, total_volume)>
    ///
    /// This is optimized for the LayeringDetector use case:
    /// - Parallel search from multiple high-volume sources
    /// - Early termination when enough paths found
    /// - Returns edge weights for CV calculation
    pub fn find_layering_paths(
        &self,
        sources: Vec<String>,
        targets: Vec<String>,
        min_length: usize,
        max_length: usize,
        max_paths_per_source: usize,
        min_volume: f64,
        max_cv: f64,
    ) -> Vec<(Vec<String>, Vec<f64>, f64)> {
        // Convert targets to a set of NodeIndex for fast lookup
        let target_set: HBSet<NodeIndex> = targets
            .iter()
            .filter_map(|t| self.node_map.get(t).copied())
            .collect();

        if target_set.is_empty() {
            return vec![];
        }

        // Convert sources to NodeIndex
        let source_indices: Vec<(String, NodeIndex)> = sources
            .into_iter()
            .filter_map(|s| self.node_map.get(&s).map(|&idx| (s, idx)))
            .collect();

        // Global counter for total paths found
        let total_paths = AtomicUsize::new(0);
        let max_total = max_paths_per_source * source_indices.len().min(50);

        // Parallel search from each source
        let results: Vec<Vec<(Vec<String>, Vec<f64>, f64)>> = source_indices
            .par_iter()
            .map(|(_source_name, source_idx)| {
                let source_idx = *source_idx;

                // Check if we already have enough paths globally
                if total_paths.load(Ordering::Relaxed) >= max_total {
                    return vec![];
                }

                let mut paths = Vec::new();
                let mut stack: Vec<(NodeIndex, Vec<NodeIndex>, Vec<f64>, f64)> = Vec::new();

                // Initialize with source node
                stack.push((source_idx, vec![source_idx], vec![], 0.0));

                while let Some((current, path, weights, volume)) = stack.pop() {
                    // Check if we have enough paths from this source
                    if paths.len() >= max_paths_per_source {
                        break;
                    }

                    // Check global limit
                    if total_paths.load(Ordering::Relaxed) >= max_total {
                        break;
                    }

                    // If path is long enough and ends at a target, check if it's a layering pattern
                    if path.len() >= min_length && target_set.contains(&current) {
                        // Check volume threshold
                        if volume >= min_volume && !weights.is_empty() {
                            // Check coefficient of variation
                            let mean: f64 = weights.iter().copied().sum::<f64>() / weights.len() as f64;
                            if mean > 0.0 {
                                let variance: f64 = weights.iter()
                                    .map(|w| (*w - mean).powi(2))
                                    .sum::<f64>() / weights.len() as f64;
                                let cv = variance.sqrt() / mean;

                                if cv <= max_cv {
                                    // Convert path to node names
                                    let path_names: Vec<String> = path
                                        .iter()
                                        .filter_map(|idx| self.reverse_map.get(idx).cloned())
                                        .collect();

                                    if path_names.len() == path.len() {
                                        paths.push((path_names, weights.clone(), volume));
                                        total_paths.fetch_add(1, Ordering::Relaxed);
                                    }
                                }
                            }
                        }
                    }

                    // Don't extend if path is at max length
                    if path.len() >= max_length {
                        continue;
                    }

                    // Explore outgoing edges
                    for edge in self.directed.edges(current) {
                        let next = edge.target();
                        let weight = *edge.weight();

                        // Skip if already in path (simple path)
                        if path.contains(&next) {
                            continue;
                        }

                        let mut new_path = path.clone();
                        new_path.push(next);

                        let mut new_weights = weights.clone();
                        new_weights.push(weight);

                        let new_volume = volume + weight;

                        stack.push((next, new_path, new_weights, new_volume));
                    }
                }

                paths
            })
            .collect();

        // Flatten results
        results.into_iter().flatten().collect()
    }

    /// Find simple cycles in parallel across SCCs
    /// Uses DFS per SCC with length bound - parallelized with rayon
    #[pyo3(signature = (max_length=10, max_cycles_per_scc=100))]
    pub fn simple_cycles(&self, max_length: usize, max_cycles_per_scc: usize) -> Vec<Vec<String>> {
        // First get SCCs using Kosaraju
        let sccs = kosaraju_scc(&self.directed);

        // Filter SCCs with more than 1 node (only these can have cycles)
        let multi_node_sccs: Vec<Vec<NodeIndex>> = sccs.into_iter()
            .filter(|scc| scc.len() > 1)
            .collect();

        // Process SCCs in parallel using rayon
        let all_cycles: Vec<Vec<Vec<String>>> = multi_node_sccs
            .par_iter()
            .map(|scc| {
                find_cycles_in_scc_helper(&self.directed, &self.reverse_map, scc, max_length, max_cycles_per_scc)
            })
            .collect();

        // Flatten results
        all_cycles.into_iter().flatten().collect()
    }

    /// Find cycles with volume ranking - prioritizes high-value suspicious cycles
    /// Returns: Vec<(cycle_nodes, total_volume, suspicion_score)>
    ///
    /// Pareto principle: Focus on cycles that move the most money
    /// suspicion_score = volume * log(1 + cycle_length)  (longer cycles + more money = more suspicious)
    #[pyo3(signature = (max_length=10, max_cycles=100, min_volume=0.0))]
    pub fn find_suspicious_cycles(
        &self,
        max_length: usize,
        max_cycles: usize,
        min_volume: f64,
    ) -> Vec<(Vec<String>, f64, f64)> {
        // Get all cycles first
        let sccs = kosaraju_scc(&self.directed);
        let multi_node_sccs: Vec<Vec<NodeIndex>> = sccs.into_iter()
            .filter(|scc| scc.len() > 1)
            .collect();

        // Build edge weight lookup for fast volume calculation
        let mut edge_weights: HBMap<(NodeIndex, NodeIndex), f64> = HBMap::new();
        for edge in self.directed.edge_references() {
            edge_weights.insert((edge.source(), edge.target()), *edge.weight());
        }

        // Find cycles with volume in parallel
        let cycles_with_volume: Vec<Vec<(Vec<String>, f64, f64)>> = multi_node_sccs
            .par_iter()
            .map(|scc| {
                find_cycles_with_volume_helper(
                    &self.directed,
                    &self.reverse_map,
                    &edge_weights,
                    scc,
                    max_length,
                    max_cycles / multi_node_sccs.len().max(1) + 1,
                    min_volume,
                )
            })
            .collect();

        // Flatten and sort by suspicion score (descending)
        let mut all_cycles: Vec<(Vec<String>, f64, f64)> = cycles_with_volume
            .into_iter()
            .flatten()
            .collect();

        // Sort by suspicion score (volume * log(length))
        all_cycles.sort_by(|a, b| b.2.partial_cmp(&a.2).unwrap_or(std::cmp::Ordering::Equal));

        // Take top N
        all_cycles.truncate(max_cycles);
        all_cycles
    }

    /// Calculate risk proximity scores for all nodes based on distance to risk addresses
    /// Returns: Vec<(node_name, risk_score, closest_risk_address, distance)>
    ///
    /// Score = sum(1 / (distance^decay)) for all risk addresses within max_distance
    /// Higher score = closer to more bad actors = more suspicious
    #[pyo3(signature = (risk_addresses, max_distance=3, decay_factor=1.0))]
    pub fn calculate_risk_proximity(
        &self,
        risk_addresses: Vec<String>,
        max_distance: usize,
        decay_factor: f64,
    ) -> Vec<(String, f64, String, usize)> {
        // Build undirected adjacency
        let adj: HBMap<NodeIndex, Vec<NodeIndex>> = {
            let mut adj: HBMap<NodeIndex, Vec<NodeIndex>> = HBMap::new();
            for edge in self.directed.edge_references() {
                let s = edge.source();
                let t = edge.target();
                adj.entry(s).or_default().push(t);
                adj.entry(t).or_default().push(s);
            }
            adj
        };

        // Convert risk addresses to indices
        let risk_indices: Vec<(String, NodeIndex)> = risk_addresses
            .into_iter()
            .filter_map(|s| self.node_map.get(&s).map(|&idx| (s, idx)))
            .collect();

        if risk_indices.is_empty() {
            return vec![];
        }

        // Parallel BFS from each risk address
        let risk_distances: Vec<HBMap<NodeIndex, (usize, String)>> = risk_indices
            .par_iter()
            .map(|(risk_name, risk_idx)| {
                let mut distances: HBMap<NodeIndex, (usize, String)> = HBMap::new();
                let mut queue: VecDeque<(NodeIndex, usize)> = VecDeque::new();

                queue.push_back((*risk_idx, 0));
                distances.insert(*risk_idx, (0, risk_name.clone()));

                while let Some((node, dist)) = queue.pop_front() {
                    if dist >= max_distance {
                        continue;
                    }

                    if let Some(neighbors) = adj.get(&node) {
                        for &neighbor in neighbors {
                            if !distances.contains_key(&neighbor) {
                                let new_dist = dist + 1;
                                distances.insert(neighbor, (new_dist, risk_name.clone()));
                                if new_dist < max_distance {
                                    queue.push_back((neighbor, new_dist));
                                }
                            }
                        }
                    }
                }
                distances
            })
            .collect();

        // Aggregate scores per node
        let mut node_scores: HBMap<NodeIndex, (f64, String, usize)> = HBMap::new();

        for distances in risk_distances {
            for (node_idx, (dist, risk_name)) in distances {
                if dist == 0 {
                    continue;  // Skip the risk address itself
                }

                let contribution = 1.0 / (dist as f64).powf(decay_factor);

                node_scores.entry(node_idx)
                    .and_modify(|(score, closest, closest_dist)| {
                        *score += contribution;
                        if dist < *closest_dist {
                            *closest = risk_name.clone();
                            *closest_dist = dist;
                        }
                    })
                    .or_insert((contribution, risk_name.clone(), dist));
            }
        }

        // Convert to output format and sort by score
        let mut results: Vec<(String, f64, String, usize)> = node_scores
            .into_iter()
            .filter_map(|(idx, (score, closest, dist))| {
                self.reverse_map.get(&idx).map(|name| (name.clone(), score, closest, dist))
            })
            .collect();

        // Sort by risk score (descending)
        results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        results
    }

    /// Detect temporal bursts for all addresses in parallel
    ///
    /// Takes per-address timestamp data and finds bursts using sliding window analysis.
    /// Returns: Vec<(address, burst_start_ms, burst_end_ms, tx_count, volume, intensity, z_score, counterparties)>
    ///
    /// Parallelized across addresses using rayon.
    #[pyo3(signature = (addresses, timestamps_list, volumes_list, counterparties_list, time_window_ms, min_transactions=10, min_intensity=3.0, z_threshold=2.0))]
    pub fn detect_temporal_bursts(
        &self,
        addresses: Vec<String>,
        timestamps_list: Vec<Vec<i64>>,
        volumes_list: Vec<Vec<f64>>,
        counterparties_list: Vec<Vec<String>>,
        time_window_ms: i64,
        min_transactions: usize,
        min_intensity: f64,
        z_threshold: f64,
    ) -> Vec<(String, i64, i64, usize, f64, f64, f64, Vec<String>)> {
        if addresses.is_empty() {
            return vec![];
        }

        // Process each address in parallel
        let results: Vec<Option<(String, i64, i64, usize, f64, f64, f64, Vec<String>)>> = addresses
            .into_par_iter()
            .zip(timestamps_list.into_par_iter())
            .zip(volumes_list.into_par_iter())
            .zip(counterparties_list.into_par_iter())
            .map(|(((addr, timestamps), volumes), counterparties)| {
                if timestamps.len() < min_transactions {
                    return None;
                }

                // Sort by timestamp (create indices to sort all arrays together)
                let mut indices: Vec<usize> = (0..timestamps.len()).collect();
                indices.sort_by_key(|&i| timestamps[i]);

                let sorted_ts: Vec<i64> = indices.iter().map(|&i| timestamps[i]).collect();
                let sorted_vol: Vec<f64> = indices.iter().map(|&i| volumes[i]).collect();
                let sorted_cp: Vec<String> = indices.iter().map(|&i| counterparties[i].clone()).collect();

                // Calculate baseline rate
                let total_span = sorted_ts.last().unwrap_or(&0) - sorted_ts.first().unwrap_or(&0);
                if total_span <= 0 {
                    return None;
                }

                let baseline_rate = sorted_ts.len() as f64 / (total_span as f64 / time_window_ms as f64);

                // Find best burst using sliding window with binary search
                let mut best_burst: Option<(i64, i64, usize, f64, f64, f64, Vec<String>)> = None;
                let mut best_z_score = 0.0f64;

                for i in 0..sorted_ts.len() {
                    let window_start = sorted_ts[i];
                    let window_end = window_start + time_window_ms;

                    // Binary search for end of window
                    let end_idx = sorted_ts.partition_point(|&ts| ts <= window_end);
                    let window_count = end_idx - i;

                    if window_count < min_transactions {
                        continue;
                    }

                    let window_rate = window_count as f64;
                    let intensity = window_rate / baseline_rate.max(0.1);

                    if intensity < min_intensity {
                        continue;
                    }

                    // Calculate z-score
                    let expected = baseline_rate;
                    let std_dev = expected.sqrt().max(0.1);
                    let z_score = (window_count as f64 - expected) / std_dev;

                    if z_score < z_threshold {
                        continue;
                    }

                    if z_score > best_z_score {
                        best_z_score = z_score;
                        let window_volume: f64 = sorted_vol[i..end_idx].iter().sum();
                        let actual_end = sorted_ts[end_idx - 1];
                        let window_counterparties: Vec<String> = sorted_cp[i..end_idx]
                            .iter()
                            .cloned()
                            .collect::<std::collections::HashSet<_>>()
                            .into_iter()
                            .collect();

                        best_burst = Some((
                            window_start,
                            actual_end,
                            window_count,
                            window_volume,
                            intensity,
                            z_score,
                            window_counterparties,
                        ));
                    }
                }

                best_burst.map(|(start, end, count, vol, intensity, z, cp)| {
                    (addr, start, end, count, vol, intensity, z, cp)
                })
            })
            .collect();

        // Filter out None values
        results.into_iter().flatten().collect()
    }

    /// Find chain bursts in timestamped edge data (for BurstDetector)
    ///
    /// Takes edge instances with timestamps and finds sequential chains A→B→C→D
    /// that occur within a tight time window.
    ///
    /// Returns: Vec<(chain_addresses, total_volume, start_ts, end_ts)>
    ///
    /// This is optimized for temporal analysis:
    /// - Processes time windows in parallel with rayon
    /// - Uses efficient DFS with pruning
    /// - Returns chains sorted by volume
    #[pyo3(signature = (sources, targets, timestamps, volumes, time_window_ms, min_chain_length=5, max_chains=1000))]
    pub fn find_chain_bursts(
        &self,
        sources: Vec<String>,
        targets: Vec<String>,
        timestamps: Vec<i64>,
        volumes: Vec<f64>,
        time_window_ms: i64,
        min_chain_length: usize,
        max_chains: usize,
    ) -> Vec<(Vec<String>, f64, i64, i64)> {
        if sources.len() < min_chain_length {
            return vec![];
        }

        // Build edge instances sorted by timestamp
        let mut edges: Vec<(String, String, i64, f64, usize)> = sources
            .into_iter()
            .zip(targets)
            .zip(timestamps)
            .zip(volumes)
            .enumerate()
            .map(|(idx, (((s, t), ts), v))| (s, t, ts, v, idx))
            .collect();

        edges.sort_by_key(|e| e.2);

        // Create time windows to process in parallel
        // Sample windows at regular intervals to cover the data
        let n_edges = edges.len();
        if n_edges == 0 {
            return vec![];
        }

        let first_ts = edges[0].2;
        let last_ts = edges[n_edges - 1].2;
        let total_span = last_ts - first_ts;

        // Create window start points - sample every quarter of the window
        let window_step = (time_window_ms / 4).max(1);
        let n_windows = ((total_span / window_step) as usize + 1).min(10000);

        let window_starts: Vec<i64> = (0..n_windows)
            .map(|i| first_ts + (i as i64) * window_step)
            .collect();

        // Process windows in parallel
        let results: Vec<Vec<(Vec<String>, f64, i64, i64)>> = window_starts
            .par_iter()
            .map(|&window_start| {
                let window_end = window_start + time_window_ms;

                // Find edges in this window using binary search
                let start_idx = edges.partition_point(|e| e.2 < window_start);
                let end_idx = edges.partition_point(|e| e.2 <= window_end);

                if end_idx - start_idx < min_chain_length {
                    return vec![];
                }

                let window_edges: Vec<_> = edges[start_idx..end_idx].to_vec();

                // Build adjacency for this window
                let mut adj: HBMap<&str, Vec<(&str, usize, f64, i64)>> = HBMap::new();
                let mut in_degree: HBMap<&str, usize> = HBMap::new();
                let mut all_nodes: HBSet<&str> = HBSet::new();

                for (src, tgt, ts, vol, idx) in &window_edges {
                    adj.entry(src.as_str()).or_default().push((tgt.as_str(), *idx, *vol, *ts));
                    *in_degree.entry(tgt.as_str()).or_insert(0) += 1;
                    all_nodes.insert(src.as_str());
                    all_nodes.insert(tgt.as_str());
                }

                // Find sources (in_degree=0 in this window)
                let sources: Vec<&str> = adj.keys()
                    .filter(|&&n| !in_degree.contains_key(n) || in_degree[n] == 0)
                    .copied()
                    .collect();

                let mut chains = Vec::new();
                let max_per_window = max_chains / n_windows.max(1) + 1;

                // DFS from each source with iterative approach and pruning
                for source in sources {
                    if chains.len() >= max_per_window {
                        break;
                    }

                    // Stack: (current_node, path, used_edge_indices, total_volume, min_ts, max_ts)
                    let mut stack: Vec<(&str, Vec<&str>, HBSet<usize>, f64, i64, i64)> = Vec::new();
                    stack.push((source, vec![source], HBSet::new(), 0.0, i64::MAX, i64::MIN));

                    while let Some((current, path, used, volume, min_ts, max_ts)) = stack.pop() {
                        if chains.len() >= max_per_window {
                            break;
                        }

                        // Record chain if long enough
                        if path.len() >= min_chain_length {
                            chains.push((
                                path.iter().map(|s| s.to_string()).collect(),
                                volume,
                                min_ts,
                                max_ts,
                            ));
                        }

                        // Limit path depth
                        if path.len() >= 20 {
                            continue;
                        }

                        // Explore neighbors
                        if let Some(neighbors) = adj.get(current) {
                            for &(next, edge_idx, edge_vol, edge_ts) in neighbors {
                                if used.contains(&edge_idx) || path.contains(&next) {
                                    continue;
                                }

                                let mut new_path = path.clone();
                                new_path.push(next);

                                let mut new_used = used.clone();
                                new_used.insert(edge_idx);

                                stack.push((
                                    next,
                                    new_path,
                                    new_used,
                                    volume + edge_vol,
                                    min_ts.min(edge_ts),
                                    max_ts.max(edge_ts),
                                ));
                            }
                        }
                    }
                }

                chains
            })
            .collect();

        // Flatten, deduplicate by normalized key, and sort by volume
        let mut all_chains: Vec<(Vec<String>, f64, i64, i64)> = Vec::new();
        let mut seen: HBSet<String> = HBSet::new();

        for chains in results {
            for chain in chains {
                // Create dedup key from first 3 nodes + start timestamp
                let key = format!(
                    "{}_{}_{}_{}",
                    chain.0.first().unwrap_or(&String::new()),
                    chain.0.get(1).unwrap_or(&String::new()),
                    chain.0.get(2).unwrap_or(&String::new()),
                    chain.2
                );

                if !seen.contains(&key) {
                    seen.insert(key);
                    all_chains.push(chain);
                }
            }
        }

        // Sort by volume descending and truncate
        all_chains.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
        all_chains.truncate(max_chains);

        all_chains
    }

    /// Get high-volume nodes (nodes with total flow above percentile threshold)
    /// Returns: Vec<(node_name, total_volume)>
    pub fn get_high_volume_nodes(&self, percentile: f64) -> Vec<(String, f64)> {
        // Calculate total volume for each node
        let node_volumes: Vec<(NodeIndex, f64)> = self.directed
            .node_indices()
            .map(|node| {
                let in_vol: f64 = self.directed
                    .edges_directed(node, petgraph::Direction::Incoming)
                    .map(|e| e.weight())
                    .sum();
                let out_vol: f64 = self.directed
                    .edges(node)
                    .map(|e| e.weight())
                    .sum();
                (node, in_vol + out_vol)
            })
            .collect();

        if node_volumes.is_empty() {
            return vec![];
        }

        // Sort by volume to find percentile threshold
        let mut volumes: Vec<f64> = node_volumes.iter().map(|(_, v)| *v).collect();
        volumes.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

        let threshold_idx = ((percentile / 100.0) * volumes.len() as f64) as usize;
        let threshold = volumes.get(threshold_idx.min(volumes.len() - 1)).copied().unwrap_or(0.0);

        // Filter and convert to names
        node_volumes
            .into_iter()
            .filter(|(_, vol)| *vol >= threshold)
            .filter_map(|(idx, vol)| {
                self.reverse_map.get(&idx).map(|name| (name.clone(), vol))
            })
            .collect()
    }
}

/// Helper function for volume-weighted cycle detection
fn find_cycles_with_volume_helper(
    graph: &DiGraph<String, f64>,
    reverse_map: &HBMap<NodeIndex, String>,
    edge_weights: &HBMap<(NodeIndex, NodeIndex), f64>,
    scc: &[NodeIndex],
    max_length: usize,
    max_cycles: usize,
    min_volume: f64,
) -> Vec<(Vec<String>, f64, f64)> {
    let mut cycles = Vec::new();
    let scc_set: HBSet<NodeIndex> = scc.iter().copied().collect();

    for &start in scc {
        if cycles.len() >= max_cycles {
            break;
        }

        // DFS with volume tracking: (current_node, path, path_volume)
        let mut stack: Vec<(NodeIndex, Vec<NodeIndex>, f64)> = vec![(start, vec![start], 0.0)];
        let mut found_from_start = 0;
        let max_per_start = max_cycles / scc.len().max(1) + 1;

        while let Some((current, path, volume)) = stack.pop() {
            if found_from_start >= max_per_start || path.len() > max_length {
                continue;
            }

            for edge in graph.edges(current) {
                let next = edge.target();
                if !scc_set.contains(&next) {
                    continue;
                }

                let edge_vol = edge_weights.get(&(current, next)).copied().unwrap_or(0.0);
                let new_volume = volume + edge_vol;

                // Found cycle back to start
                if next == start && path.len() >= 2 {
                    if new_volume >= min_volume {
                        let cycle_names: Vec<String> = path
                            .iter()
                            .filter_map(|idx| reverse_map.get(idx).cloned())
                            .collect();

                        if cycle_names.len() == path.len() {
                            // Suspicion score: volume * log(1 + length)
                            let suspicion = new_volume * (1.0 + path.len() as f64).ln();
                            cycles.push((cycle_names, new_volume, suspicion));
                            found_from_start += 1;

                            if cycles.len() >= max_cycles {
                                return cycles;
                            }
                        }
                    }
                    continue;
                }

                if path.contains(&next) {
                    continue;
                }

                let mut new_path = path.clone();
                new_path.push(next);
                stack.push((next, new_path, new_volume));
            }
        }
    }

    cycles
}

/// Helper function for parallel cycle detection (outside impl to avoid PyO3 exposure)
fn find_cycles_in_scc_helper(
    graph: &DiGraph<String, f64>,
    reverse_map: &HBMap<NodeIndex, String>,
    scc: &[NodeIndex],
    max_length: usize,
    max_cycles: usize,
) -> Vec<Vec<String>> {
    let mut cycles = Vec::new();
    let scc_set: HBSet<NodeIndex> = scc.iter().copied().collect();

    // Try starting from each node in SCC
    for &start in scc {
        if cycles.len() >= max_cycles {
            break;
        }

        // DFS to find cycles back to start
        let mut stack: Vec<(NodeIndex, Vec<NodeIndex>)> = vec![(start, vec![start])];
        let mut found_from_start = 0;
        let max_per_start = max_cycles / scc.len().max(1) + 1;

        while let Some((current, path)) = stack.pop() {
            if found_from_start >= max_per_start {
                break;
            }

            if path.len() > max_length {
                continue;
            }

            // Check outgoing edges
            for edge in graph.edges(current) {
                let next = edge.target();

                // Must stay within SCC
                if !scc_set.contains(&next) {
                    continue;
                }

                // Found a cycle back to start
                if next == start && path.len() >= 2 {
                    let cycle_names: Vec<String> = path
                        .iter()
                        .filter_map(|idx| reverse_map.get(idx).cloned())
                        .collect();

                    if cycle_names.len() == path.len() {
                        cycles.push(cycle_names);
                        found_from_start += 1;

                        if cycles.len() >= max_cycles {
                            return cycles;
                        }
                    }
                    continue;
                }

                // Don't revisit nodes in current path (simple cycle)
                if path.contains(&next) {
                    continue;
                }

                // Extend path
                let mut new_path = path.clone();
                new_path.push(next);
                stack.push((next, new_path));
            }
        }
    }

    cycles
}

// ============================================================================
// Data I/O Functions (exposed to Python)
// ============================================================================

/// Load transfers from parquet and build timestamp data for burst detection.
/// Returns dict: address -> list of (timestamp, volume, counterparty) tuples.
#[pyfunction]
fn load_transfers_and_build_timestamps(
    parquet_path: String,
    start_timestamp_ms: i64,
    end_timestamp_ms: i64,
    filter_addresses: Vec<String>,
) -> PyResult<HashMap<String, Vec<(i64, f64, String)>>> {
    let transfers = data_io::load_transfers_from_parquet(
        &parquet_path,
        start_timestamp_ms,
        end_timestamp_ms,
    ).map_err(|e| pyo3::exceptions::PyIOError::new_err(e))?;

    let filter_set: HBSet<String> = filter_addresses.into_iter().collect();
    let timestamp_data = data_io::build_timestamp_data(&transfers, &filter_set);

    // Convert to std HashMap for Python
    Ok(timestamp_data.into_iter()
        .map(|(addr, records)| {
            let py_records: Vec<(i64, f64, String)> = records.into_iter()
                .map(|r| (r.timestamp, r.volume, r.counterparty))
                .collect();
            (addr, py_records)
        })
        .collect())
}

/// Load money flows from parquet and return edge arrays for graph building.
/// Returns (sources, targets, weights) tuple.
#[pyfunction]
fn load_money_flows(
    parquet_path: String,
    start_timestamp_ms: i64,
    end_timestamp_ms: i64,
) -> PyResult<(Vec<String>, Vec<String>, Vec<f64>)> {
    data_io::load_money_flows_from_parquet(
        &parquet_path,
        start_timestamp_ms,
        end_timestamp_ms,
    ).map_err(|e| pyo3::exceptions::PyIOError::new_err(e))
}

/// Load transfers from parquet.
/// Returns list of (from_address, to_address, timestamp, amount_usd) tuples.
#[pyfunction]
fn load_transfers(
    parquet_path: String,
    start_timestamp_ms: i64,
    end_timestamp_ms: i64,
) -> PyResult<Vec<(String, String, i64, f64)>> {
    let transfers = data_io::load_transfers_from_parquet(
        &parquet_path,
        start_timestamp_ms,
        end_timestamp_ms,
    ).map_err(|e| pyo3::exceptions::PyIOError::new_err(e))?;

    Ok(transfers.into_iter()
        .map(|t| (t.from_address, t.to_address, t.block_timestamp, t.amount_usd))
        .collect())
}

/// Python module - embedded as chainswarm_analyzers_baseline._rust
#[pymodule]
fn _rust(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<RustGraph>()?;
    // Data I/O functions
    m.add_function(wrap_pyfunction!(load_transfers_and_build_timestamps, m)?)?;
    m.add_function(wrap_pyfunction!(load_money_flows, m)?)?;
    m.add_function(wrap_pyfunction!(load_transfers, m)?)?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_scc() {
        let mut g = RustGraph::new();
        g.add_edges(
            vec!["a".into(), "b".into(), "c".into()],
            vec!["b".into(), "c".into(), "a".into()],
            vec![1.0, 1.0, 1.0],
        );
        let sccs = g.strongly_connected_components();
        assert_eq!(sccs.len(), 1);
        assert_eq!(sccs[0].len(), 3);
    }

    #[test]
    fn test_triangle() {
        let mut g = RustGraph::new();
        // Triangle: a-b, b-c, c-a
        g.add_edges(
            vec!["a".into(), "b".into(), "c".into()],
            vec!["b".into(), "c".into(), "a".into()],
            vec![1.0, 1.0, 1.0],
        );
        assert_eq!(g.triangle_count(), 1);
    }

    #[test]
    fn test_pagerank() {
        let mut g = RustGraph::new();
        g.add_edges(
            vec!["a".into(), "b".into()],
            vec!["b".into(), "a".into()],
            vec![1.0, 1.0],
        );
        let pr = g.pagerank(0.85, 100, 1e-6);
        assert!(pr.len() == 2);
    }

    #[test]
    fn test_bfs() {
        let mut g = RustGraph::new();
        // Chain: a -> b -> c -> d
        g.add_edges(
            vec!["a".into(), "b".into(), "c".into()],
            vec!["b".into(), "c".into(), "d".into()],
            vec![1.0, 1.0, 1.0],
        );
        let distances = g.bfs("a".into(), None);
        assert_eq!(distances.get("a"), Some(&0));
        assert_eq!(distances.get("b"), Some(&1));
        assert_eq!(distances.get("c"), Some(&2));
        assert_eq!(distances.get("d"), Some(&3));
    }

    #[test]
    fn test_bfs_cutoff() {
        let mut g = RustGraph::new();
        // Chain: a -> b -> c -> d
        g.add_edges(
            vec!["a".into(), "b".into(), "c".into()],
            vec!["b".into(), "c".into(), "d".into()],
            vec![1.0, 1.0, 1.0],
        );
        let distances = g.bfs("a".into(), Some(2));
        assert_eq!(distances.get("a"), Some(&0));
        assert_eq!(distances.get("b"), Some(&1));
        assert_eq!(distances.get("c"), Some(&2));
        assert_eq!(distances.get("d"), None); // Beyond cutoff
    }

    #[test]
    fn test_find_layering_paths() {
        let mut g = RustGraph::new();
        // Create a simple layering path: a -> b -> c -> d with consistent weights
        g.add_edges(
            vec!["a".into(), "b".into(), "c".into()],
            vec!["b".into(), "c".into(), "d".into()],
            vec![100.0, 100.0, 100.0],  // Consistent weights (low CV)
        );

        let paths = g.find_layering_paths(
            vec!["a".into()],      // sources
            vec!["d".into()],      // targets
            3,                      // min_length
            10,                     // max_length
            10,                     // max_paths_per_source
            100.0,                  // min_volume
            0.5,                    // max_cv
        );

        assert_eq!(paths.len(), 1);
        assert_eq!(paths[0].0, vec!["a", "b", "c", "d"]);
        assert_eq!(paths[0].2, 300.0);  // total volume
    }

    #[test]
    fn test_high_volume_nodes() {
        let mut g = RustGraph::new();
        // Create graph with varying volumes
        g.add_edges(
            vec!["a".into(), "b".into(), "c".into()],
            vec!["b".into(), "c".into(), "d".into()],
            vec![1000.0, 100.0, 10.0],
        );

        let high_vol = g.get_high_volume_nodes(50.0);
        // Top 50% should include high-volume nodes
        assert!(!high_vol.is_empty());
    }
}
