geographdb-core 0.5.4

use crate::algorithms::four_d::GraphNode4D;
use std::collections::{HashMap, HashSet, VecDeque};

/// Alexandrov interval I[src,dst] = {w : src →* w →* dst} in the temporal causal order.
#[derive(Debug, Clone)]
pub struct CausalInterval {
    pub src: u64,
    pub dst: u64,
    /// |I[src,dst]| including both endpoints
    pub volume: usize,
    /// Length of the longest directed path from src to dst (causal proper time)
    pub proper_time: u32,
}

/// Aggregate statistics for the causal set defined by the temporal subgraph.
#[derive(Debug, Clone)]
pub struct CausalStats {
    pub n_nodes: usize,
    pub n_related_pairs: usize,
    /// Myrheim-Meyer dimension estimate (log-log regression slope of volume vs proper_time).
    /// NaN if fewer than 2 pairs with proper_time ≥ 2 exist.
    pub mm_dimension: f32,
    pub mean_volume: f32,
    pub mean_proper_time: f32,
}

/// Forward and reverse adjacency maps using only temporal edges (edge.end_ts > edge.begin_ts).
fn build_causal_adj(nodes: &[GraphNode4D]) -> (HashMap<u64, Vec<u64>>, HashMap<u64, Vec<u64>>) {
    let mut fwd: HashMap<u64, Vec<u64>> = HashMap::with_capacity(nodes.len());
    let mut rev: HashMap<u64, Vec<u64>> = HashMap::with_capacity(nodes.len());
    for node in nodes {
        fwd.entry(node.id).or_default();
        rev.entry(node.id).or_default();
        for edge in &node.successors {
            if edge.end_ts > edge.begin_ts {
                fwd.entry(node.id).or_default().push(edge.dst);
                rev.entry(edge.dst).or_default().push(node.id);
            }
        }
    }
    (fwd, rev)
}

/// BFS reachability from `start` following `adj`. Returns empty if start is not in adj.
fn reachable(adj: &HashMap<u64, Vec<u64>>, start: u64) -> HashSet<u64> {
    if !adj.contains_key(&start) {
        return HashSet::new();
    }
    let mut visited = HashSet::new();
    let mut queue = VecDeque::from([start]);
    while let Some(cur) = queue.pop_front() {
        if !visited.insert(cur) {
            continue;
        }
        if let Some(neighbors) = adj.get(&cur) {
            for &nb in neighbors {
                if !visited.contains(&nb) {
                    queue.push_back(nb);
                }
            }
        }
    }
    visited
}

/// Longest directed path from `src` to `dst` within `interval` using DAG DP.
/// Topological order is derived from node.begin_ts (strictly increasing along temporal edges).
fn longest_path(
    node_map: &HashMap<u64, &GraphNode4D>,
    fwd_adj: &HashMap<u64, Vec<u64>>,
    interval: &HashSet<u64>,
    src: u64,
    dst: u64,
) -> u32 {
    if src == dst {
        return 0;
    }
    let mut topo: Vec<u64> = interval.iter().cloned().collect();
    topo.sort_by_key(|&id| node_map.get(&id).map(|n| n.begin_ts).unwrap_or(0));

    let mut dist: HashMap<u64, u32> = HashMap::new();
    dist.insert(src, 0);

    for &u in &topo {
        let d_u = match dist.get(&u) {
            Some(&d) => d,
            None => continue,
        };
        if let Some(neighbors) = fwd_adj.get(&u) {
            for &v in neighbors {
                if interval.contains(&v) {
                    let e = dist.entry(v).or_insert(0);
                    if d_u + 1 > *e {
                        *e = d_u + 1;
                    }
                }
            }
        }
    }

    dist.get(&dst).copied().unwrap_or(0)
}

/// Computes all Alexandrov intervals in the temporal causal order.
///
/// Only temporal edges (edge.end_ts > edge.begin_ts) define causal precedence.
/// For each pair (src, dst) where dst is reachable from src, the interval
/// I[src,dst] = forward_reach(src) ∩ backward_reach(dst), inclusive of endpoints.
pub fn causal_intervals(nodes: &[GraphNode4D]) -> Vec<CausalInterval> {
    let node_map: HashMap<u64, &GraphNode4D> = nodes.iter().map(|n| (n.id, n)).collect();
    let (fwd_adj, rev_adj) = build_causal_adj(nodes);

    let fwd_reach: HashMap<u64, HashSet<u64>> = nodes
        .iter()
        .map(|n| (n.id, reachable(&fwd_adj, n.id)))
        .collect();
    let bwd_reach: HashMap<u64, HashSet<u64>> = nodes
        .iter()
        .map(|n| (n.id, reachable(&rev_adj, n.id)))
        .collect();

    let mut result = Vec::new();
    for node in nodes {
        let u = node.id;
        let Some(fwd_u) = fwd_reach.get(&u) else {
            continue;
        };
        for &v in fwd_u {
            if v == u {
                continue;
            }
            let Some(bwd_v) = bwd_reach.get(&v) else {
                continue;
            };
            let interval: HashSet<u64> = fwd_u.intersection(bwd_v).cloned().collect();
            let volume = interval.len();
            let tau = longest_path(&node_map, &fwd_adj, &interval, u, v);
            result.push(CausalInterval {
                src: u,
                dst: v,
                volume,
                proper_time: tau,
            });
        }
    }
    result
}

/// Aggregate causal-set statistics including the Myrheim-Meyer dimension estimate.
pub fn causal_stats(nodes: &[GraphNode4D]) -> CausalStats {
    let intervals = causal_intervals(nodes);
    let n_nodes = nodes.len();
    let n_related_pairs = intervals.len();

    if intervals.is_empty() {
        return CausalStats {
            n_nodes,
            n_related_pairs: 0,
            mm_dimension: f32::NAN,
            mean_volume: 0.0,
            mean_proper_time: 0.0,
        };
    }

    let mean_volume =
        intervals.iter().map(|i| i.volume as f32).sum::<f32>() / n_related_pairs as f32;
    let mean_proper_time =
        intervals.iter().map(|i| i.proper_time as f32).sum::<f32>() / n_related_pairs as f32;

    // Log-log regression slope: ln(volume) ~ mm_dim * ln(proper_time)
    // Only pairs with proper_time ≥ 2 contribute (avoid log(1) clustering noise).
    let long_pairs: Vec<(f32, f32)> = intervals
        .iter()
        .filter(|i| i.proper_time >= 2)
        .map(|i| ((i.proper_time as f32).ln(), (i.volume as f32).ln()))
        .collect();

    let mm_dimension = if long_pairs.len() < 2 {
        f32::NAN
    } else {
        let n = long_pairs.len() as f32;
        let mean_x = long_pairs.iter().map(|(x, _)| x).sum::<f32>() / n;
        let mean_y = long_pairs.iter().map(|(_, y)| y).sum::<f32>() / n;
        let cov = long_pairs
            .iter()
            .map(|(x, y)| (x - mean_x) * (y - mean_y))
            .sum::<f32>();
        let var_x = long_pairs
            .iter()
            .map(|(x, _)| (x - mean_x).powi(2))
            .sum::<f32>();
        if var_x < 1e-10 {
            f32::NAN
        } else {
            cov / var_x
        }
    };

    CausalStats {
        n_nodes,
        n_related_pairs,
        mm_dimension,
        mean_volume,
        mean_proper_time,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::algorithms::four_d::{GraphNode4D, GraphProperties, TemporalEdge};
    use crate::algorithms::tnet4d::build_tnet4d;

    fn temporal_edge(dst: u64, begin_ts: u64, end_ts: u64) -> TemporalEdge {
        TemporalEdge {
            dst,
            weight: 1.0,
            begin_ts,
            end_ts,
        }
    }

    fn spatial_edge(dst: u64, ts: u64) -> TemporalEdge {
        TemporalEdge {
            dst,
            weight: 1.0,
            begin_ts: ts,
            end_ts: ts,
        }
    }

    fn node(id: u64, t: u64, succs: Vec<TemporalEdge>) -> GraphNode4D {
        GraphNode4D {
            id,
            x: id as f32,
            y: 0.0,
            z: 0.0,
            begin_ts: t,
            end_ts: t + 1,
            properties: GraphProperties::default(),
            successors: succs,
        }
    }

    fn make_chain(n: usize) -> Vec<GraphNode4D> {
        (0..n as u64)
            .map(|i| {
                let succs = if i + 1 < n as u64 {
                    vec![temporal_edge(i + 1, i, i + 1)]
                } else {
                    vec![]
                };
                node(i, i, succs)
            })
            .collect()
    }

    // ── reachable ─────────────────────────────────────────────────────────────

    #[test]
    fn test_reachable_single_chain() {
        // 0→1→2→3: reachable from 0 via temporal edges = {0,1,2,3}
        let nodes = make_chain(4);
        let (fwd, _) = build_causal_adj(&nodes);
        let reach = reachable(&fwd, 0);
        assert!(reach.contains(&1), "0 must reach 1");
        assert!(reach.contains(&2), "0 must reach 2");
        assert!(reach.contains(&3), "0 must reach 3");
        assert!(!reach.contains(&4), "4 does not exist");
    }

    #[test]
    fn test_reachable_missing_start_is_empty() {
        // Node not in adjacency map → empty reachable set
        let nodes: Vec<GraphNode4D> = vec![];
        let (fwd, _) = build_causal_adj(&nodes);
        let reach = reachable(&fwd, 42);
        assert!(reach.is_empty(), "missing start → empty reachable set");
    }

    // ── causal_intervals ─────────────────────────────────────────────────────

    #[test]
    fn test_interval_chain_inclusive() {
        // 0→1→2→3: I[0,3] = {0,1,2,3}, volume=4, proper_time=3
        let nodes = make_chain(4);
        let intervals = causal_intervals(&nodes);
        let e = intervals
            .iter()
            .find(|i| i.src == 0 && i.dst == 3)
            .expect("interval (0,3) must exist");
        assert_eq!(e.volume, 4, "I[0,3] volume must be 4");
        assert_eq!(e.proper_time, 3, "proper_time(0,3) must be 3");
    }

    #[test]
    fn test_interval_direct_edge() {
        // 0→1: I[0,1] = {0,1}, volume=2, proper_time=1
        let nodes = make_chain(2);
        let intervals = causal_intervals(&nodes);
        assert_eq!(intervals.len(), 1);
        assert_eq!(intervals[0].volume, 2);
        assert_eq!(intervals[0].proper_time, 1);
    }

    #[test]
    fn test_interval_diamond() {
        // 0→1, 0→2, 1→3, 2→3: I[0,3] = {0,1,2,3}, volume=4, proper_time=2
        let nodes = vec![
            node(0, 0, vec![temporal_edge(1, 0, 1), temporal_edge(2, 0, 1)]),
            node(1, 1, vec![temporal_edge(3, 1, 2)]),
            node(2, 1, vec![temporal_edge(3, 1, 2)]),
            node(3, 2, vec![]),
        ];
        let intervals = causal_intervals(&nodes);
        let e = intervals
            .iter()
            .find(|i| i.src == 0 && i.dst == 3)
            .expect("interval (0,3) must exist");
        assert_eq!(e.volume, 4, "I[0,3] volume must be 4");
        assert_eq!(e.proper_time, 2, "proper_time(0,3) = 2 hops");
    }

    #[test]
    fn test_spatial_edges_ignored() {
        // Two nodes connected only by spatial edges: no causal intervals
        let nodes = vec![
            node(0, 0, vec![spatial_edge(1, 0)]),
            node(1, 0, vec![spatial_edge(0, 0)]),
        ];
        let intervals = causal_intervals(&nodes);
        assert!(
            intervals.is_empty(),
            "spatial-only graph must produce no causal intervals"
        );
    }

    #[test]
    fn test_causal_intervals_count_chain() {
        // Chain 0→1→2→3→4: C(5,2) = 10 causal pairs
        let nodes = make_chain(5);
        let intervals = causal_intervals(&nodes);
        assert_eq!(
            intervals.len(),
            10,
            "5-node chain must yield 10 causal intervals"
        );
    }

    // ── causal_stats ──────────────────────────────────────────────────────────

    #[test]
    fn test_mm_dimension_chain_approx_one() {
        // A pure temporal chain is topologically 1D; d̂ should be in [0.5, 1.5]
        let nodes = make_chain(10);
        let stats = causal_stats(&nodes);
        assert!(
            stats.mm_dimension >= 0.5 && stats.mm_dimension <= 1.5,
            "1D chain MM dimension expected in [0.5,1.5], got {}",
            stats.mm_dimension
        );
    }

    #[test]
    fn test_causal_stats_tnet4d_single_site() {
        // tnet4d(1,1,1,5): 1 site × 5 layers → 5-node chain, C(5,2)=10 pairs
        let nodes = build_tnet4d(1, 1, 1, 5);
        let stats = causal_stats(&nodes);
        assert_eq!(stats.n_nodes, 5);
        assert_eq!(stats.n_related_pairs, 10, "5-node chain → 10 pairs");
        assert!(
            stats.mm_dimension.is_finite(),
            "MM dimension must be finite"
        );
        assert!(
            stats.mm_dimension >= 0.5 && stats.mm_dimension <= 1.5,
            "single-site temporal chain should be ~1D, got {}",
            stats.mm_dimension
        );
    }

    #[test]
    fn test_causal_intervals_tnet4d_parallel() {
        // tnet4d(2,2,1,3): 4 sites × 3 layers → 4 parallel 3-chains, 4×C(3,2)=12 pairs
        let nodes = build_tnet4d(2, 2, 1, 3);
        let intervals = causal_intervals(&nodes);
        assert_eq!(
            intervals.len(),
            12,
            "4 parallel 3-chains → 12 causal intervals"
        );
    }
}