geographdb-core 0.5.4

use crate::algorithms::four_d::GraphNode4D;
use std::collections::HashMap;

#[derive(Debug, Clone)]
pub struct InfoNode {
    pub id: u64,
    pub entropy: f32,
}

#[derive(Debug, Clone)]
pub struct InfoEdge {
    pub src: u64,
    pub dst: u64,
    /// D_KL(μ_src ∥ μ_dst)
    pub kl_uv: f32,
    /// D_KL(μ_dst ∥ μ_src)
    pub kl_vu: f32,
    /// Jensen-Shannon divergence (symmetric, bounded in [0, ln 2])
    pub js_div: f32,
    /// Fisher-Rao geodesic distance on the probability simplex
    pub fisher_rao: f32,
    pub entropy_src: f32,
    pub entropy_dst: f32,
}

/// Information-geometric measures for every node and undirected edge.
///
/// For each node u the lazy random walk measure μ_u^α is built:
///   μ_u(u) = α,  μ_u(v) = (1−α)/deg(u) for each out-neighbour v.
///
/// For each edge (u,v) this function computes KL divergence (both directions),
/// Jensen-Shannon divergence, and the Fisher-Rao geodesic distance between
/// μ_u^α and μ_v^α.
pub fn info_geometry(nodes: &[GraphNode4D], alpha: f32) -> (Vec<InfoNode>, Vec<InfoEdge>) {
    let node_map: HashMap<u64, &GraphNode4D> = nodes.iter().map(|n| (n.id, n)).collect();

    let info_nodes = nodes
        .iter()
        .map(|node| {
            let mu = build_measure(&node_map, node.id, alpha);
            InfoNode {
                id: node.id,
                entropy: entropy(&mu),
            }
        })
        .collect();

    let mut info_edges = Vec::new();
    for node in nodes {
        let u = node.id;
        for edge in &node.successors {
            let v = edge.dst;
            if v <= u {
                continue; // each undirected edge once
            }
            let mu = build_measure(&node_map, u, alpha);
            let nu = build_measure(&node_map, v, alpha);
            info_edges.push(InfoEdge {
                src: u,
                dst: v,
                kl_uv: kl_divergence(&mu, &nu),
                kl_vu: kl_divergence(&nu, &mu),
                js_div: js_divergence(&mu, &nu),
                fisher_rao: fisher_rao_dist(&mu, &nu),
                entropy_src: entropy(&mu),
                entropy_dst: entropy(&nu),
            });
        }
    }

    (info_nodes, info_edges)
}

/// Lazy random walk measure (same construction as in ricci.rs).
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
}

/// Shannon entropy H(μ) = −Σ p(x) ln p(x).
fn entropy(mu: &[(u64, f32)]) -> f32 {
    mu.iter()
        .filter(|&&(_, p)| p > 0.0)
        .map(|&(_, p)| -p * p.ln())
        .sum()
}

/// KL divergence D_KL(p ∥ q) = Σ p(x) ln(p(x)/q(x)).
/// Returns f32::INFINITY if q(x) = 0 for any x where p(x) > 0.
fn kl_divergence(p: &[(u64, f32)], q: &[(u64, f32)]) -> f32 {
    let q_map: HashMap<u64, f32> = q.iter().cloned().collect();
    let mut kl = 0.0_f32;
    for &(id, pi) in p {
        if pi <= 0.0 {
            continue;
        }
        let qi = q_map.get(&id).copied().unwrap_or(0.0);
        if qi <= 0.0 {
            return f32::INFINITY;
        }
        kl += pi * (pi / qi).ln();
    }
    kl
}

/// Fisher-Rao geodesic distance on the probability simplex:
///   d_FR(p,q) = 2 · arccos(Σ_x √(p(x)·q(x)))
///
/// Equals 0 when p = q and π when supports are disjoint.
fn fisher_rao_dist(p: &[(u64, f32)], q: &[(u64, f32)]) -> f32 {
    let q_map: HashMap<u64, f32> = q.iter().cloned().collect();
    let bc: f32 = p
        .iter()
        .filter(|&&(_, pi)| pi > 0.0)
        .map(|&(id, pi)| {
            let qi = q_map.get(&id).copied().unwrap_or(0.0);
            (pi * qi).sqrt()
        })
        .sum();
    2.0 * bc.clamp(0.0, 1.0).acos()
}

/// Jensen-Shannon divergence JSD(p,q) = H(m) − ½(H(p)+H(q))  where m = (p+q)/2.
fn js_divergence(p: &[(u64, f32)], q: &[(u64, f32)]) -> f32 {
    let mut m_map: HashMap<u64, f32> = HashMap::new();
    for &(id, pi) in p {
        *m_map.entry(id).or_insert(0.0) += pi * 0.5;
    }
    for &(id, qi) in q {
        *m_map.entry(id).or_insert(0.0) += qi * 0.5;
    }
    let m: Vec<(u64, f32)> = m_map.into_iter().collect();
    (entropy(&m) - (entropy(p) + entropy(q)) * 0.5).max(0.0)
}

#[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(),
        }
    }

    // ── entropy ───────────────────────────────────────────────────────────────

    #[test]
    fn test_entropy_uniform_four_items() {
        // Uniform over 4 items → H = ln(4)
        let mu = vec![(0u64, 0.25), (1, 0.25), (2, 0.25), (3, 0.25)];
        let h = entropy(&mu);
        let expected = (4.0_f32).ln();
        assert!(
            (h - expected).abs() < 1e-5,
            "H(uniform4) expected {expected:.6}, got {h:.6}"
        );
    }

    #[test]
    fn test_entropy_point_mass_is_zero() {
        let mu = vec![(0u64, 1.0)];
        assert!(entropy(&mu).abs() < 1e-6, "H(point mass) must be 0");
    }

    // ── KL divergence ─────────────────────────────────────────────────────────

    #[test]
    fn test_kl_identical_distributions_zero() {
        let p = vec![(0u64, 0.5), (1, 0.3), (2, 0.2)];
        assert!(kl_divergence(&p, &p).abs() < 1e-5, "D_KL(p‖p) must be 0");
    }

    #[test]
    fn test_kl_orthogonal_distributions_infinite() {
        let p = vec![(0u64, 0.5), (1, 0.5)];
        let q = vec![(2u64, 0.5), (3, 0.5)];
        assert!(
            kl_divergence(&p, &q).is_infinite(),
            "D_KL with disjoint support must be ∞"
        );
    }

    #[test]
    fn test_kl_triangle_measures_positive() {
        // μ_0 = {0:0.5, 1:0.25, 2:0.25}, μ_1 = {0:0.25, 1:0.5, 2:0.25}
        // D_KL = 0.25·ln(2) ≈ 0.1733
        let mu0 = vec![(0u64, 0.5), (1, 0.25), (2, 0.25)];
        let mu1 = vec![(1u64, 0.5), (0, 0.25), (2, 0.25)];
        let kl = kl_divergence(&mu0, &mu1);
        let expected = 0.25 * 2.0_f32.ln(); // ≈ 0.1733
        assert!(
            (kl - expected).abs() < 1e-4,
            "D_KL(μ_0‖μ_1) for triangle expected {expected:.4}, got {kl:.4}"
        );
    }

    // ── Fisher-Rao distance ───────────────────────────────────────────────────

    #[test]
    fn test_fisher_rao_identical_distributions_zero() {
        let p = vec![(0u64, 0.5), (1, 0.5)];
        assert!(fisher_rao_dist(&p, &p).abs() < 1e-5, "d_FR(p,p) must be 0");
    }

    #[test]
    fn test_fisher_rao_orthogonal_distributions_pi() {
        // Disjoint support → Bhattacharyya coefficient = 0 → d_FR = 2·arccos(0) = π
        let p = vec![(0u64, 1.0)];
        let q = vec![(1u64, 1.0)];
        let d = fisher_rao_dist(&p, &q);
        assert!(
            (d - std::f32::consts::PI).abs() < 1e-5,
            "d_FR with disjoint support expected π, got {d:.6}"
        );
    }

    // ── Jensen-Shannon divergence ─────────────────────────────────────────────

    #[test]
    fn test_js_divergence_symmetry() {
        let p = vec![(0u64, 0.6), (1, 0.4)];
        let q = vec![(0u64, 0.2), (1, 0.8)];
        let jspq = js_divergence(&p, &q);
        let jsqp = js_divergence(&q, &p);
        assert!(
            (jspq - jsqp).abs() < 1e-6,
            "JSD(p,q) and JSD(q,p) must be equal; got {jspq:.6} vs {jsqp:.6}"
        );
    }

    // ── info_geometry integration ─────────────────────────────────────────────

    #[test]
    fn test_info_geometry_two_node_identical_measures() {
        // Two-node graph 0↔1: α=0.5 → μ_0 = {0:0.5, 1:0.5} = μ_1
        let nodes = vec![node(0, &[1]), node(1, &[0])];
        let (info_nodes, info_edges) = info_geometry(&nodes, 0.5);
        assert_eq!(info_nodes.len(), 2);
        assert_eq!(info_edges.len(), 1);
        let e = &info_edges[0];
        assert!(
            e.kl_uv.abs() < 1e-5,
            "identical measures → kl_uv=0, got {}",
            e.kl_uv
        );
        assert!(
            e.fisher_rao.abs() < 1e-5,
            "identical measures → FR=0, got {}",
            e.fisher_rao
        );
        assert!(
            e.js_div.abs() < 1e-5,
            "identical measures → JS=0, got {}",
            e.js_div
        );
    }

    #[test]
    fn test_info_geometry_triangle_positive_divergences() {
        let nodes = vec![node(0, &[1, 2]), node(1, &[0, 2]), node(2, &[0, 1])];
        let (info_nodes, info_edges) = info_geometry(&nodes, 0.5);
        assert_eq!(info_nodes.len(), 3);
        assert_eq!(info_edges.len(), 3);
        // Measures are not identical → all divergences > 0
        for e in &info_edges {
            assert!(e.kl_uv > 0.0, "triangle kl_uv must be > 0");
            assert!(e.fisher_rao > 0.0, "triangle fisher_rao must be > 0");
            assert!(e.fisher_rao < std::f32::consts::PI, "fisher_rao < π");
            assert!(e.js_div > 0.0, "triangle js_div must be > 0");
        }
        // Entropy of each node: H({self:0.5, n1:0.25, n2:0.25}) = 3/2·ln(2) ≈ 1.0397
        for n in &info_nodes {
            let expected = 1.5 * 2.0_f32.ln();
            assert!(
                (n.entropy - expected).abs() < 1e-4,
                "triangle node entropy expected {expected:.4}, got {:.4}",
                n.entropy
            );
        }
    }
}