foxhunt/ml/src/tgnn/graph.rs

//! Market graph implementation for TGGN
//!
//! Optimized graph structure for market microstructure representation
//! with cache-friendly operations and temporal decay.

use std::sync::RwLock;

use dashmap::DashMap;
use petgraph::graph::NodeIndex;
use petgraph::{Directed, Graph};
use serde::{Deserialize, Serialize};
use tracing::debug;

use super::{MarketEdge, NodeId, NodeType};
use crate::{MLError, PRECISION_FACTOR};

/// Graph statistics
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct GraphStats {
    pub node_count: usize,
    pub edge_count: usize,
    pub density: f64,
    pub average_degree: f64,
    pub max_degree: usize,
    pub connected_components: usize,
}

/// High-performance market graph for TGGN
#[derive(Debug)]
pub struct MarketGraph {
    /// Underlying petgraph structure
    graph: RwLock<Graph<NodeData, EdgeData, Directed>>,

    /// Node mapping for fast lookup
    node_mapping: DashMap<NodeId, NodeIndex>,

    /// Node features cache
    node_features: DashMap<NodeId, Vec<f64>>,

    /// Edge cache for fast neighbor lookup
    edge_cache: DashMap<NodeId, Vec<NodeId>>,

    /// Maximum nodes allowed
    pub max_nodes: usize,

    /// Maximum edges allowed  
    pub max_edges: usize,

    /// Current statistics
    stats: RwLock<GraphStats>,
}

#[derive(Debug, Clone)]
struct NodeData {
    node_id: NodeId,
    #[allow(dead_code)]
    features: Vec<f64>,
    timestamp: u64,
}

#[derive(Debug, Clone)]
struct EdgeData {
    edge: MarketEdge,
    #[allow(dead_code)]
    source: NodeId,
    #[allow(dead_code)]
    target: NodeId,
}

impl MarketGraph {
    /// Create new market graph
    pub fn new(max_nodes: usize, max_edges: usize) -> Result<Self, MLError> {
        Ok(Self {
            graph: RwLock::new(Graph::with_capacity(max_nodes, max_edges)),
            node_mapping: DashMap::new(),
            node_features: DashMap::new(),
            edge_cache: DashMap::new(),
            max_nodes,
            max_edges,
            stats: RwLock::new(GraphStats::default()),
        })
    }

    /// Add node to graph
    pub fn add_node(&self, node_id: NodeId, features: Vec<f64>) -> Result<(), MLError> {
        // Check capacity
        if self.node_mapping.len() >= self.max_nodes {
            return Err(MLError::ResourceLimit {
                resource: "graph_nodes".to_string(),
                limit: self.max_nodes,
            });
        }

        let node_data = NodeData {
            node_id: node_id.clone(),
            features: features.clone(),
            timestamp: std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap_or_default()
                .as_nanos() as u64,
        };

        let mut graph = self.graph.write().map_err(|_| MLError::ConcurrencyError {
            operation: "graph_write".to_string(),
        })?;

        let node_index = graph.add_node(node_data);
        self.node_mapping.insert(node_id.clone(), node_index);
        self.node_features.insert(node_id.clone(), features);

        // Update stats
        drop(graph);
        self.update_stats()?;

        debug!("Added node {:?} with index {:?}", node_id, node_index);
        Ok(())
    }

    /// Remove node from graph
    pub fn remove_node(&self, node_id: &NodeId) -> Result<(), MLError> {
        if let Some((_, node_index)) = self.node_mapping.remove(node_id) {
            let mut graph = self.graph.write().map_err(|_| MLError::ConcurrencyError {
                operation: "graph_write".to_string(),
            })?;

            graph.remove_node(node_index);
            self.node_features.remove(node_id);
            self.edge_cache.remove(node_id);

            // Clean up edge cache references
            self.edge_cache.retain(|_, neighbors| {
                neighbors.retain(|n| n != node_id);
                !neighbors.is_empty()
            });

            drop(graph);
            self.update_stats()?;
            debug!("Removed node {:?}", node_id);
        }
        Ok(())
    }

    /// Add edge between nodes
    pub fn add_edge(
        &self,
        source: &NodeId,
        target: &NodeId,
        edge: MarketEdge,
    ) -> Result<(), MLError> {
        // Check capacity
        let graph_guard = self.graph.read().map_err(|_| MLError::ConcurrencyError {
            operation: "graph_read".to_string(),
        })?;

        if graph_guard.edge_count() >= self.max_edges {
            return Err(MLError::ResourceLimit {
                resource: "graph_edges".to_string(),
                limit: self.max_edges,
            });
        }
        drop(graph_guard);

        let source_idx = self
            .node_mapping
            .get(source)
            .ok_or_else(|| MLError::GraphError {
                message: format!("Source node {:?} not found", source),
            })?;

        let target_idx = self
            .node_mapping
            .get(target)
            .ok_or_else(|| MLError::GraphError {
                message: format!("Target node {:?} not found", target),
            })?;

        let edge_data = EdgeData {
            edge,
            source: source.clone(),
            target: target.clone(),
        };

        let mut graph = self.graph.write().map_err(|_| MLError::ConcurrencyError {
            operation: "graph_write".to_string(),
        })?;

        graph.add_edge(*source_idx, *target_idx, edge_data);

        // Update edge cache
        self.edge_cache
            .entry(source.clone())
            .or_insert_with(Vec::new)
            .push(target.clone());

        drop(graph);
        self.update_stats()?;
        debug!("Added edge from {:?} to {:?}", source, target);
        Ok(())
    }

    /// Get node features
    pub fn get_node_features(&self, node_id: &NodeId) -> Option<Vec<f64>> {
        self.node_features.get(node_id).map(|f| f.clone())
    }

    /// Update node features
    pub fn update_node_features(
        &self,
        node_id: &NodeId,
        features: Vec<f64>,
    ) -> Result<(), MLError> {
        if let Some(mut node_features) = self.node_features.get_mut(node_id) {
            *node_features = features;
            debug!("Updated features for node {:?}", node_id);
            Ok(())
        } else {
            Err(MLError::GraphError {
                message: format!("Node {:?} not found", node_id),
            })
        }
    }

    /// Get neighbors of a node
    pub fn get_neighbors(&self, node_id: &NodeId) -> Option<Vec<NodeId>> {
        self.edge_cache
            .get(node_id)
            .map(|neighbors| neighbors.clone())
    }

    /// Get edge weight between nodes
    pub fn get_edge_weight(&self, source: &NodeId, target: &NodeId) -> Option<f64> {
        let graph = self.graph.read().ok()?;
        let source_idx = *self.node_mapping.get(source)?;
        let target_idx = *self.node_mapping.get(target)?;

        if let Some(edge_idx) = graph.find_edge(source_idx, target_idx) {
            let edge_data = graph.edge_weight(edge_idx)?;
            // Normalize weight to 0-1 range
            Some(edge_data.edge.weight as f64 / PRECISION_FACTOR as f64)
        } else {
            None
        }
    }

    /// Get number of nodes
    pub fn node_count(&self) -> usize {
        self.node_mapping.len()
    }

    /// Get number of edges
    pub fn edge_count(&self) -> usize {
        self.graph.read().map(|g| g.edge_count()).unwrap_or(0)
    }

    /// Clear temporal data based on age
    pub fn clear_temporal_data(
        &mut self,
        current_time: u64,
        decay_factor: f64,
    ) -> Result<(), MLError> {
        let graph = self.graph.write().map_err(|_| MLError::ConcurrencyError {
            operation: "graph_write".to_string(),
        })?;

        let mut nodes_to_remove = Vec::new();

        // Check node ages and mark for removal if too old
        for node_index in graph.node_indices() {
            if let Some(node_data) = graph.node_weight(node_index) {
                let age = current_time.saturating_sub(node_data.timestamp);
                let age_seconds = age as f64 / 1_000_000_000.0;
                let decay = decay_factor.powf(age_seconds);

                // Remove nodes that have decayed below threshold
                if decay < 0.01 {
                    nodes_to_remove.push(node_data.node_id.clone());
                }
            }
        }

        drop(graph);

        // Remove old nodes
        for node_id in nodes_to_remove {
            self.remove_node(&node_id)?;
        }

        // Apply temporal decay to edges
        let mut graph = self.graph.write().map_err(|_| MLError::ConcurrencyError {
            operation: "graph_write".to_string(),
        })?;

        for edge_index in graph.edge_indices() {
            if let Some(edge_data) = graph.edge_weight_mut(edge_index) {
                edge_data.edge.apply_temporal_decay(current_time);
            }
        }

        drop(graph);
        self.update_stats()?;
        Ok(())
    }

    /// Get nodes by type
    pub fn get_nodes_by_type(&self, node_type: NodeType) -> Vec<NodeId> {
        self.node_mapping
            .iter()
            .filter(|entry| entry.key().node_type == node_type)
            .map(|entry| entry.key().clone())
            .collect()
    }

    /// Find shortest path between nodes (simplified Dijkstra)
    pub fn shortest_path(&self, source: &NodeId, target: &NodeId) -> Option<Vec<NodeId>> {
        let graph = self.graph.read().ok()?;
        let source_idx = *self.node_mapping.get(source)?;
        let target_idx = *self.node_mapping.get(target)?;

        use petgraph::algo::dijkstra;
        let node_map = dijkstra(&*graph, source_idx, Some(target_idx), |_| 1);

        if node_map.contains_key(&target_idx) {
            // Reconstruct path (simplified)
            let mut path = Vec::new();

            // This is a simplified path reconstruction
            // A full implementation would track predecessors
            for entry in self.node_mapping.iter() {
                let node_id = entry.key();
                let node_idx = entry.value();
                if *node_idx == source_idx {
                    path.insert(0, node_id.clone());
                } else if *node_idx == target_idx {
                    path.push(node_id.clone());
                } else if node_map.contains_key(node_idx) {
                    path.insert(path.len().saturating_sub(1), node_id.clone());
                }
            }

            Some(path)
        } else {
            None
        }
    }

    /// Get graph statistics
    pub fn get_stats(&self) -> GraphStats {
        self.stats
            .read()
            .map(|stats| stats.clone())
            .unwrap_or_default()
    }

    /// Update internal statistics
    fn update_stats(&self) -> Result<(), MLError> {
        let mut stats = self.stats.write().map_err(|_| MLError::ConcurrencyError {
            operation: "stats_write".to_string(),
        })?;

        let node_count = self.node_count();
        let edge_count = self.edge_count();

        stats.node_count = node_count;
        stats.edge_count = edge_count;
        stats.density = if node_count > 1 {
            (2.0 * edge_count as f64) / (node_count as f64 * (node_count - 1) as f64)
        } else {
            0.0
        };
        stats.average_degree = if node_count > 0 {
            (2.0 * edge_count as f64) / node_count as f64
        } else {
            0.0
        };

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tgnn::EdgeType;
    // use crate::safe_operations; // DISABLED - module not found

    #[test]
    fn test_graph_creation() -> Result<(), MLError> {
        let graph = MarketGraph::new(100, 500)?;
        assert_eq!(graph.max_nodes, 100);
        assert_eq!(graph.max_edges, 500);
        assert_eq!(graph.node_count(), 0);
        assert_eq!(graph.edge_count(), 0);
        Ok(())
    }

    #[test]
    fn test_node_operations() -> Result<(), MLError> {
        let graph = MarketGraph::new(10, 20)?;

        let node1 = NodeId::price_level(100);
        let node2 = NodeId::price_level(101);

        // Add nodes
        graph.add_node(node1.clone(), vec![1.0, 2.0])?;
        graph.add_node(node2.clone(), vec![3.0, 4.0])?;

        assert_eq!(graph.node_count(), 2);

        // Check features
        let features = graph.get_node_features(&node1).unwrap();
        assert_eq!(features, vec![1.0, 2.0]);

        // Update features
        graph.update_node_features(&node1, vec![5.0, 6.0])?;
        let updated_features = graph.get_node_features(&node1).unwrap();
        assert_eq!(updated_features, vec![5.0, 6.0]);

        // Remove node
        graph.remove_node(&node1)?;
        assert_eq!(graph.node_count(), 1);
        assert!(graph.get_node_features(&node1).is_none());
        Ok(())
    }

    #[test]
    fn test_edge_operations() -> Result<(), MLError> {
        let graph = MarketGraph::new(10, 20)?;

        let node1 = NodeId::price_level(100);
        let node2 = NodeId::price_level(101);

        graph.add_node(node1.clone(), vec![1.0])?;
        graph.add_node(node2.clone(), vec![2.0])?;

        let edge = MarketEdge::new(EdgeType::PriceProximity, 5000, 0.8);
        graph.add_edge(&node1, &node2, edge)?;

        assert_eq!(graph.edge_count(), 1);

        // Check edge weight
        let weight = graph.get_edge_weight(&node1, &node2).unwrap();
        assert!((weight - 0.5).abs() < 0.1); // 5000 / 10000 = 0.5

        // Check neighbors
        let neighbors = graph.get_neighbors(&node1).unwrap();
        assert_eq!(neighbors.len(), 1);
        assert_eq!(neighbors[0], node2);
        Ok(())
    }

    #[test]
    fn test_graph_stats() -> Result<(), MLError> {
        let graph = MarketGraph::new(10, 20)?;

        let node1 = NodeId::price_level(100);
        let node2 = NodeId::price_level(101);
        let node3 = NodeId::price_level(102);

        graph.add_node(node1.clone(), vec![1.0])?;
        graph.add_node(node2.clone(), vec![2.0])?;
        graph.add_node(node3.clone(), vec![3.0])?;

        let edge1 = MarketEdge::new(EdgeType::PriceProximity, 1000, 0.8);
        let edge2 = MarketEdge::new(EdgeType::LiquidityFlow, 2000, 0.9);

        graph.add_edge(&node1, &node2, edge1)?;
        graph.add_edge(&node2, &node3, edge2)?;

        let stats = graph.get_stats();
        assert_eq!(stats.node_count, 3);
        assert_eq!(stats.edge_count, 2);
        assert!(stats.density > 0.0);
        assert!(stats.average_degree > 0.0);
        Ok(())
    }

    #[test]
    fn test_nodes_by_type() -> Result<(), MLError> {
        let graph = MarketGraph::new(10, 20)?;

        let price_node = NodeId::price_level(100);
        let mm_node = NodeId::market_maker("test_mm");

        graph.add_node(price_node.clone(), vec![1.0])?;
        graph.add_node(mm_node.clone(), vec![2.0])?;

        let price_nodes = graph.get_nodes_by_type(NodeType::PriceLevel);
        let mm_nodes = graph.get_nodes_by_type(NodeType::MarketMaker);

        assert_eq!(price_nodes.len(), 1);
        assert_eq!(mm_nodes.len(), 1);
        assert_eq!(price_nodes[0], price_node);
        assert_eq!(mm_nodes[0], mm_node);
        Ok(())
    }

    #[test]
    fn test_shortest_path() -> Result<(), MLError> {
        let graph = MarketGraph::new(10, 20)?;

        let node1 = NodeId::price_level(100);
        let node2 = NodeId::price_level(101);
        let node3 = NodeId::price_level(102);

        graph.add_node(node1.clone(), vec![1.0])?;
        graph.add_node(node2.clone(), vec![2.0])?;
        graph.add_node(node3.clone(), vec![3.0])?;

        let edge1 = MarketEdge::new(EdgeType::PriceProximity, 1000, 0.8);
        let edge2 = MarketEdge::new(EdgeType::PriceProximity, 2000, 0.9);

        graph.add_edge(&node1, &node2, edge1)?;
        graph.add_edge(&node2, &node3, edge2)?;

        let path = graph.shortest_path(&node1, &node3).unwrap();
        assert_eq!(path.len(), 3);
        assert_eq!(path[0], node1);
        assert_eq!(path[1], node2);
        assert_eq!(path[2], node3);
        Ok(())
    }
}