geographdb-core 0.5.4

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

#[derive(Debug, Clone)]
pub struct RicciEdge {
    pub src: u64,
    pub dst: u64,
    pub curvature: f32,
    pub w1: f32,
}

/// Ollivier-Ricci curvature for every undirected edge in the graph.
///
/// κ(u,v) = 1 - W₁(μ_u^α, μ_v^α) / d(u,v)
///
/// where μ_u^α is the lazy random walk measure: mass α stays at u,
/// mass (1-α)/deg(u) spreads uniformly to each out-neighbour.
///
/// `alpha` ∈ [0,1) is the laziness parameter (0.5 is the standard choice).
pub fn ollivier_ricci(nodes: &[GraphNode4D], alpha: f32) -> Vec<RicciEdge> {
    let node_map: HashMap<u64, &GraphNode4D> = nodes.iter().map(|n| (n.id, n)).collect();
    let mut results = Vec::new();

    for node in nodes {
        let u = node.id;
        for edge in &node.successors {
            let v = edge.dst;
            if v <= u {
                continue; // process each undirected edge exactly once
            }
            let dist_u = bfs_dist(&node_map, u);
            let d_uv = match dist_u.get(&v) {
                Some(&d) if d > 0.0 => d,
                _ => continue,
            };

            let mu = build_measure(&node_map, u, alpha);
            let nu = build_measure(&node_map, v, alpha);
            let w1 = wasserstein1(&mu, &nu, &node_map);
            let curvature = 1.0 - w1 / d_uv;
            results.push(RicciEdge {
                src: u,
                dst: v,
                curvature,
                w1,
            });
        }
    }

    results
}

/// Build a directed curvature lookup from an undirected Ollivier-Ricci computation.
///
/// For every undirected edge {u,v} the map contains both `(u,v)` and `(v,u)`
/// entries so that walkers can look up curvature in either direction.
pub fn curvature_map(nodes: &[GraphNode4D], alpha: f32) -> HashMap<(u64, u64), f32> {
    let mut map = HashMap::new();
    for edge in ollivier_ricci(nodes, alpha) {
        map.insert((edge.src, edge.dst), edge.curvature);
        map.insert((edge.dst, edge.src), edge.curvature);
    }
    map
}

/// Fast Dice neighbor-overlap curvature map.
///
/// κ(u,v) = 2|N(u) ∩ N(v)| / (|N(u)| + |N(v)|) - 1
///
/// This is position-independent and uses fixed-size bitsets per node, making it
/// feasible for dense 50K+ token graphs where exact Ollivier-Ricci would need a
/// BFS per edge.
///
/// `min_edge_weight` filters weak edges: only successors with
/// `edge.weight >= min_edge_weight` are included in the neighborhoods.
/// For count-weighted graphs use the desired count threshold; for PMI-weighted
/// graphs set this to 0.0 to keep all edges.
pub fn curvature_map_fast(nodes: &[GraphNode4D], min_edge_weight: f32) -> HashMap<(u64, u64), f32> {
    let threshold = min_edge_weight.max(0.0);
    let n = nodes.len();
    if n == 0 {
        return HashMap::new();
    }
    let blocks = n.div_ceil(64);

    // Map node id -> dense index for bitset addressing.
    let id_to_idx: HashMap<u64, usize> = nodes
        .iter()
        .enumerate()
        .map(|(i, node)| (node.id, i))
        .collect();

    // Build adjacency bitsets from edges passing the weight threshold.
    let mut adj_bits: Vec<Vec<u64>> = vec![vec![0u64; blocks]; n];
    for (i, node) in nodes.iter().enumerate() {
        for e in &node.successors {
            if e.weight < threshold {
                continue;
            }
            if let Some(&j) = id_to_idx.get(&e.dst) {
                if i != j {
                    let block = j / 64;
                    let bit = j % 64;
                    adj_bits[i][block] |= 1u64 << bit;
                }
            }
        }
    }

    // Precompute degrees from bitsets.
    let degrees: Vec<usize> = adj_bits
        .iter()
        .map(|bits| bits.iter().map(|b| b.count_ones() as usize).sum())
        .collect();

    // Compute κ for every directed edge.
    let mut map = HashMap::new();
    for (i, node) in nodes.iter().enumerate() {
        if degrees[i] == 0 {
            continue;
        }
        for e in &node.successors {
            if e.weight < threshold {
                continue;
            }
            let Some(&j) = id_to_idx.get(&e.dst) else {
                continue;
            };
            if i == j || degrees[j] == 0 {
                continue;
            }
            let mut common = 0usize;
            for (b, row_i) in adj_bits[i].iter().enumerate() {
                common += (row_i & adj_bits[j][b]).count_ones() as usize;
            }
            let deg_sum = degrees[i] + degrees[j];
            let kappa = if deg_sum == 0 {
                0.0f32
            } else {
                2.0f32 * common as f32 / deg_sum as f32 - 1.0f32
            };
            map.insert((node.id, e.dst), kappa);
            map.insert((e.dst, node.id), kappa);
        }
    }

    map
}

/// Lazy random walk measure at `id` with laziness `alpha`.
fn build_measure(node_map: &HashMap<u64, &GraphNode4D>, id: u64, alpha: f32) -> Vec<(u64, f32)> {
    let node = match node_map.get(&id) {
        Some(n) => n,
        None => return vec![(id, 1.0)],
    };
    let deg = node.successors.len();
    if deg == 0 {
        return vec![(id, 1.0)];
    }
    let mut measure = Vec::with_capacity(deg + 1);
    measure.push((id, alpha));
    let w = (1.0 - alpha) / deg as f32;
    for e in &node.successors {
        measure.push((e.dst, w));
    }
    measure
}

/// BFS shortest-path distances (unit edge weights) from `start`.
fn bfs_dist(node_map: &HashMap<u64, &GraphNode4D>, start: u64) -> HashMap<u64, f32> {
    let mut dist = HashMap::new();
    let mut queue = VecDeque::new();
    dist.insert(start, 0.0_f32);
    queue.push_back(start);
    while let Some(cur) = queue.pop_front() {
        let d = dist[&cur];
        if let Some(node) = node_map.get(&cur) {
            for edge in &node.successors {
                if let std::collections::hash_map::Entry::Vacant(e) = dist.entry(edge.dst) {
                    e.insert(d + 1.0);
                    queue.push_back(edge.dst);
                }
            }
        }
    }
    dist
}

/// Wasserstein-1 distance between `mu` and `nu` using the graph metric.
fn wasserstein1(
    mu: &[(u64, f32)],
    nu: &[(u64, f32)],
    node_map: &HashMap<u64, &GraphNode4D>,
) -> f32 {
    let m = mu.len();
    let n = nu.len();

    // Build cost matrix: pairwise BFS distances.
    let mut cost = vec![0.0_f32; m * n];
    for (i, &(s, _)) in mu.iter().enumerate() {
        let dists = bfs_dist(node_map, s);
        for (j, &(t, _)) in nu.iter().enumerate() {
            cost[i * n + j] = dists.get(&t).copied().unwrap_or(f32::INFINITY);
        }
    }

    let supply: Vec<f32> = mu.iter().map(|&(_, w)| w).collect();
    let demand: Vec<f32> = nu.iter().map(|&(_, w)| w).collect();

    solve_transport(&supply, &demand, &cost, m, n)
}

/// Transportation simplex: North-West corner initialisation + MODI potentials
/// + stepping-stone loop improvement.
fn solve_transport(supply: &[f32], demand: &[f32], cost: &[f32], m: usize, n: usize) -> f32 {
    // Replace ∞ with big-M so the NW corner never produces NaN.
    // The optimal plan avoids big-M cells since finite alternatives always exist
    // for probability distributions with overlapping support on a connected graph.
    let max_finite = cost
        .iter()
        .filter(|&&c| c.is_finite())
        .cloned()
        .fold(1.0_f32, f32::max);
    let big_m = max_finite * 10.0 * (m + n) as f32;
    let cost_buf: Vec<f32> = cost
        .iter()
        .map(|&c| if c.is_finite() { c } else { big_m })
        .collect();
    let cost = cost_buf.as_slice();

    let mut s = supply.to_vec();
    let mut d = demand.to_vec();
    let mut alloc = vec![0.0_f32; m * n];
    let mut basic = vec![false; m * n];
    let eps = 1e-8_f32;

    // North-West corner with degenerate-case handling to maintain m+n-1 basics.
    let (mut ii, mut jj) = (0_usize, 0_usize);
    while ii < m && jj < n {
        let x = s[ii].min(d[jj]);
        alloc[ii * n + jj] = x;
        basic[ii * n + jj] = true;
        s[ii] -= x;
        d[jj] -= x;

        let si_done = s[ii] < eps;
        let dj_done = d[jj] < eps;

        if si_done && dj_done {
            ii += 1;
            if ii < m && jj < n {
                // Degenerate: place a 0-alloc basic so we have m+n-1 basics.
                basic[ii * n + jj] = true;
            }
            jj += 1;
        } else if si_done {
            ii += 1;
        } else {
            jj += 1;
        }
    }

    // MODI + stepping-stone loop (max 50 pivots for safety).
    for _ in 0..50 {
        // Compute dual potentials u[i] + v[j] = cost[i,j] for each basic (i,j).
        let mut u = vec![f32::NAN; m];
        let mut v = vec![f32::NAN; n];
        u[0] = 0.0;
        let mut changed = true;
        while changed {
            changed = false;
            for i in 0..m {
                for j in 0..n {
                    if !basic[i * n + j] {
                        continue;
                    }
                    let c = cost[i * n + j];
                    if !u[i].is_nan() && v[j].is_nan() {
                        v[j] = c - u[i];
                        changed = true;
                    } else if u[i].is_nan() && !v[j].is_nan() {
                        u[i] = c - v[j];
                        changed = true;
                    }
                }
            }
        }

        // Find most negative reduced cost among non-basic cells.
        let mut best_rc = -1e-6_f32;
        let (mut enter_i, mut enter_j) = (m, n); // sentinel = no entering cell
        for i in 0..m {
            for j in 0..n {
                if basic[i * n + j] || u[i].is_nan() || v[j].is_nan() {
                    continue;
                }
                let rc = cost[i * n + j] - u[i] - v[j];
                if rc < best_rc {
                    best_rc = rc;
                    enter_i = i;
                    enter_j = j;
                }
            }
        }

        if enter_i == m {
            break; // optimal
        }

        // Find the unique cycle in the basis tree that includes (enter_i, enter_j).
        let tree_cells = find_tree_path(&basic, m, n, enter_i, enter_j);
        if tree_cells.is_empty() {
            break; // degenerate fallback — shouldn't happen with correct initialisation
        }

        // Loop: entering cell (+) followed by tree path alternating -, +, -, ...
        let mut loop_cells = vec![(enter_i, enter_j)];
        loop_cells.extend_from_slice(&tree_cells);

        // θ = min allocation at minus-positions (odd indices).
        // Track only the first tie as the leaving variable to avoid removing
        // multiple cells from the basis in a single degenerate (θ=0) pivot.
        let mut theta = f32::INFINITY;
        let mut leave_k = loop_cells.len(); // sentinel
        for (k, &(i, j)) in loop_cells.iter().enumerate() {
            if k % 2 == 0 {
                continue;
            }
            let a = alloc[i * n + j];
            if a < theta {
                theta = a;
                leave_k = k;
            }
        }
        if leave_k == loop_cells.len() {
            break;
        }

        for (k, &(i, j)) in loop_cells.iter().enumerate() {
            if k % 2 == 0 {
                alloc[i * n + j] += theta;
                basic[i * n + j] = true;
            } else {
                alloc[i * n + j] -= theta;
                if k == leave_k {
                    // Only one cell leaves the basis per pivot.
                    alloc[i * n + j] = alloc[i * n + j].max(0.0);
                    basic[i * n + j] = false;
                }
            }
        }
    }

    // Compute total transport cost.
    (0..m * n).map(|k| alloc[k] * cost[k]).sum()
}

/// BFS on the basis spanning tree to find the path from row-node `start_row`
/// to col-node `end_col`. Returns the sequence of (row, col) basic cells
/// along that path, ordered from start to end.
///
/// The bipartite basis tree has m + n nodes: rows 0..m and cols m..m+n.
fn find_tree_path(
    basic: &[bool],
    m: usize,
    n: usize,
    start_row: usize,
    end_col: usize,
) -> Vec<(usize, usize)> {
    let total = m + n;
    let end_node = m + end_col;

    let mut visited = vec![false; total];
    let mut parent: Vec<Option<usize>> = vec![None; total];

    visited[start_row] = true;
    let mut queue = VecDeque::from([start_row]);

    while let Some(cur) = queue.pop_front() {
        if cur == end_node {
            // Reconstruct path of (row, col) pairs.
            let mut edge_pairs = Vec::new();
            let mut node = cur;
            while let Some(par) = parent[node] {
                // Each edge in the bipartite tree corresponds to one basic cell.
                let (row, col) = if par < m {
                    (par, node - m) // row→col edge
                } else {
                    (node, par - m) // col→row edge (reversed direction)
                };
                edge_pairs.push((row, col));
                node = par;
            }
            edge_pairs.reverse();
            return edge_pairs;
        }

        if cur < m {
            // Row node: expand to connected col nodes via basic cells.
            for j in 0..n {
                let next = m + j;
                if !visited[next] && basic[cur * n + j] {
                    visited[next] = true;
                    parent[next] = Some(cur);
                    queue.push_back(next);
                }
            }
        } else {
            // Col node: expand to connected row nodes via basic cells.
            let j = cur - m;
            for i in 0..m {
                if !visited[i] && basic[i * n + j] {
                    visited[i] = true;
                    parent[i] = Some(cur);
                    queue.push_back(i);
                }
            }
        }
    }

    Vec::new() // no path (shouldn't happen with a valid basis)
}

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

    fn node(id: u64, neighbors: &[u64]) -> GraphNode4D {
        GraphNode4D {
            id,
            x: id as f32,
            y: 0.0,
            z: 0.0,
            begin_ts: 0,
            end_ts: 0,
            properties: GraphProperties::default(),
            successors: neighbors
                .iter()
                .map(|&dst| TemporalEdge {
                    dst,
                    weight: 1.0,
                    begin_ts: 0,
                    end_ts: 0,
                })
                .collect(),
        }
    }

    fn find_edge(edges: &[RicciEdge], src: u64, dst: u64) -> Option<&RicciEdge> {
        edges.iter().find(|e| e.src == src && e.dst == dst)
    }

    // ── Two-node graph ────────────────────────────────────────────────────────
    // Identical neighbourhood distributions → W₁ = 0 → κ = 1.
    #[test]
    fn test_ricci_two_node_complete_curvature_one() {
        let nodes = vec![node(0, &[1]), node(1, &[0])];
        let edges = ollivier_ricci(&nodes, 0.5);
        let e = find_edge(&edges, 0, 1).expect("edge (0,1) missing");
        assert!(
            (e.curvature - 1.0).abs() < 1e-5,
            "expected κ=1, got {}",
            e.curvature
        );
        assert!(e.w1.abs() < 1e-5, "expected W₁=0, got {}", e.w1);
    }

    // ── Triangle graph ────────────────────────────────────────────────────────
    // Positive curvature: κ = 0.75 with α=0.5.
    #[test]
    fn test_ricci_triangle_positive_curvature() {
        let nodes = vec![node(0, &[1, 2]), node(1, &[0, 2]), node(2, &[0, 1])];
        let edges = ollivier_ricci(&nodes, 0.5);
        let e = find_edge(&edges, 0, 1).expect("edge (0,1) missing");
        // W₁ = 0.25 (transport 0.25 from node-0's self-mass to node-1's self-mass at cost 1).
        assert!(
            (e.w1 - 0.25).abs() < 1e-4,
            "triangle W₁ expected 0.25, got {}",
            e.w1
        );
        assert!(
            (e.curvature - 0.75).abs() < 1e-4,
            "triangle κ expected 0.75, got {}",
            e.curvature
        );
    }

    // ── Path graph P4 middle edge ─────────────────────────────────────────────
    // Zero curvature on the middle edge (0-1-2-3): κ(1,2) = 0.
    #[test]
    fn test_ricci_path_graph_middle_edge_zero() {
        let nodes = vec![
            node(0, &[1]),
            node(1, &[0, 2]),
            node(2, &[1, 3]),
            node(3, &[2]),
        ];
        let edges = ollivier_ricci(&nodes, 0.5);
        let e = find_edge(&edges, 1, 2).expect("edge (1,2) missing");
        assert!(
            e.curvature.abs() < 1e-4,
            "path middle edge κ expected 0, got {}",
            e.curvature
        );
    }

    // ── 4-cycle C4 ────────────────────────────────────────────────────────────
    // C4 with α=0.5 has κ = 0.5 for every edge.
    #[test]
    fn test_ricci_c4_cycle_half_curvature() {
        // 0-1-3-2-0 (square)
        let nodes = vec![
            node(0, &[1, 2]),
            node(1, &[0, 3]),
            node(2, &[0, 3]),
            node(3, &[1, 2]),
        ];
        let edges = ollivier_ricci(&nodes, 0.5);
        for e in &edges {
            assert!(
                (e.curvature - 0.5).abs() < 1e-4,
                "C4 edge ({},{}) κ expected 0.5, got {}",
                e.src,
                e.dst,
                e.curvature
            );
        }
        assert_eq!(edges.len(), 4, "C4 should have 4 undirected edges");
    }

    // ── Negative curvature: tree graph ────────────────────────────────────────
    // A 3-regular Bethe lattice edge has κ < 0.  Here a simple star graph with
    // 4 leaves has negative curvature on each leaf edge with α=0.
    #[test]
    fn test_ricci_star_graph_negative_curvature() {
        // Star: centre=0, leaves=1,2,3,4
        let nodes = vec![
            node(0, &[1, 2, 3, 4]),
            node(1, &[0]),
            node(2, &[0]),
            node(3, &[0]),
            node(4, &[0]),
        ];
        let edges = ollivier_ricci(&nodes, 0.0);
        // Edge (0,1): μ_0 spreads over {1,2,3,4} uniformly (0.25 each),
        // μ_1 = {0: 1.0}. W₁ = 0.25*d(1,0) + 0.25*d(2,0) + 0.25*d(3,0) + 0.25*d(4,0)
        // But that's not right — μ_0 has masses 0.25 going to leaves, μ_1 has all mass at centre.
        // W₁ = 0.75 * d(neighbour_of_0, centre) = 0.75 * 1 = 0.75  (3/4 of mass travels dist 1)
        // Wait: source {1:0.25, 2:0.25, 3:0.25, 4:0.25}, sink {0:1.0}
        // All mass goes to 0: W₁ = 0.25*(d(1,0)+d(2,0)+d(3,0)+d(4,0)) = 0.25*4 = 1.0
        // κ = 1 - 1.0/1 = 0  (edge between a hub and leaf has κ=0 in star with α=0)
        // Actually for edge (0,1): leaves that are not 1 must travel through 0 (d=2), leaf 1→0 is d=1.
        // Source: {1:0.25, 2:0.25, 3:0.25, 4:0.25}. Sink: {0:1.0}.
        // d(1,0)=1, d(2,0)=1, d(3,0)=1, d(4,0)=1. W₁ = 1.0.  κ=0.
        //
        // For edge (1,2) (leaf-to-leaf, not adjacent): not in graph.
        //
        // Hmm, κ=0 not negative. Let me instead check that leaf nodes with α=0 give
        // non-positive curvature and the curvature is well-defined.
        // A star-graph leaf edge with α=0 has κ = 1 - W₁/d.
        // W₁ = 1 (as computed), d=1, so κ=0 for (0,i).
        // Not negative in this case.
        //
        // Let me just verify all edges are computed and curvature = 0 for star with α=0.
        assert_eq!(edges.len(), 4);
        for e in &edges {
            assert!(
                e.curvature.abs() < 1e-4,
                "star edge ({},{}) expected κ≈0 with α=0, got {}",
                e.src,
                e.dst,
                e.curvature
            );
        }
    }

    // ── Laziness parameter effect ─────────────────────────────────────────────
    // With α=1 (fully lazy), μ_u = δ_u and W₁ = d(u,v) → κ = 0 always.
    #[test]
    fn test_ricci_fully_lazy_curvature_zero() {
        let nodes = vec![node(0, &[1, 2]), node(1, &[0, 2]), node(2, &[0, 1])];
        let edges = ollivier_ricci(&nodes, 1.0);
        for e in &edges {
            assert!(
                e.curvature.abs() < 1e-4,
                "fully lazy edge ({},{}) expected κ=0, got {}",
                e.src,
                e.dst,
                e.curvature
            );
        }
    }
}