dyntri-core 0.10.2

//! Average sphere distance algorithms
//!
//! The fastest implementations are [`asd()`] and [`asd_double()`], for the single-sphere
//! and double-sphere average sphere distance respectively.
//! If one can accept some statistical error [`asd_sampled()`] provides an estimate of the
//! average sphere distance with controllable accuracy against speed
//!
//! An implementation that could make use of parallelized BFS is [`asd_multi()`].

use std::collections::VecDeque;

use itertools::Itertools;
use log::debug;
use rand::seq::IndexedRandom;
use rand::{Rng, seq::SliceRandom};

use super::bfs::{bfs, multi_bfs_with_dist};

use super::bfs::bfs_sphere_sizehint;
use crate::graph::Graph;

/// Returns the single-sphere average sphere distance.
///
/// The `Vec` corresponds to the average sphere distance at delta = 1, 2, 3, ..., delta_max.
pub fn asd(graph: &Graph, origin: usize, delta_max: u16) -> Vec<f64> {
    // Construct starting spheres for all delta up to delta_max
    let bfs_result = bfs(graph, origin, delta_max);
    let dist_sphere = bfs_result.distance_list;
    let spheres = bfs_result.spheres;
    let sphere_vol = bfs_result.sphere_volumes;

    // Initialize ASD vector (make use of one-pass average algorithm with norm)
    let mut asd = vec![0f64; delta_max as usize];
    // Do not start at 0 but at sphere_vol to account for the missed 0 distance geodesics
    let mut norm = sphere_vol[1..=delta_max as usize].to_owned();

    let mut visited = vec![false; graph.len()];
    for start_node in spheres {
        // Get the delta of the current starting point
        let delta = dist_sphere[start_node];
        // ASD for delta=0 is not well-defined, skip
        if delta == 0 {
            continue;
        }

        visited.fill(false);
        visited[start_node] = true;

        // Use a Deque structure to keep track of nodes still to visit and their distance
        // TODO: Figure out good initial value (8*delta is ideal for square grid)
        let mut queue = VecDeque::with_capacity(delta as usize * 8);
        queue.push_back((start_node, 0));

        // Get (mutable) references to values corresponding to current `delta`
        let current_asd = asd.get_mut(delta as usize - 1).unwrap();
        let current_norm = norm.get_mut(delta as usize - 1).unwrap();
        let svol = *sphere_vol.get(delta as usize).unwrap();

        // Perform breadth-first search
        'bfs: loop {
            let Some((node, r)) = queue.pop_front() else {
                break; // Stop if queue is empty (all nodes have been searched)
            };
            for &nbr in graph.get_neighbours(node) {
                let visited_nbr = &mut visited[nbr];
                if *visited_nbr {
                    continue; // Do not reconsider already visited node
                }
                *visited_nbr = true;
                queue.push_back((nbr, r + 1));

                // Compute ASD contribution
                if dist_sphere[nbr] != delta {
                    continue; // Only contribute nodes on delta sphere
                }

                // Use one-pass algorithm
                *current_norm += 1;
                *current_asd += ((r + 1) as f64 - *current_asd) / (*current_norm as f64);
                // Stop BFS is all nodes in `delta` sphere have been reached
                if *current_norm % (svol - 1) == 1 {
                    break 'bfs;
                }
            }
        }
    }

    asd
}

/// Returns the single-sphere average sphere distance.
///
/// The `Vec` corresponds to the average sphere distance at each delta in `deltas`.
pub fn asd_deltas(graph: &Graph, origin: usize, deltas: &[u16]) -> Vec<f64> {
    if deltas.is_empty() {
        return Vec::new();
    }
    let delta_max = deltas.iter().copied().max().unwrap();
    // Construct starting spheres for all delta up to delta_max
    let bfs_result = bfs(graph, origin, delta_max);
    let sphere_idx = bfs_result.sphere_idx;
    let dist_sphere = bfs_result.distance_list;
    let spheres = bfs_result.spheres;

    let mut asd: Vec<f64> = Vec::with_capacity(deltas.len());

    // // Initialize ASD vector (make use of one-pass average algorithm with norm)
    // let mut asd = vec![0f64; delta_max as usize];
    // // Do not start at 0 but at sphere_vol to account for the missed 0 distance geodesics
    // let mut norm = sphere_vol[1..=delta_max as usize].to_owned();

    let mut visited = vec![false; graph.len()];
    for &delta in deltas {
        let sphere = &spheres[sphere_idx[delta as usize]..sphere_idx[delta as usize + 1]];
        // Initialize ASD (make use of one-pass average algorithm with norm)
        let mut asd_delta = 0f64;
        // Do not start at 0 but at sphere_vol to account for the missed 0 distance geodesics
        let mut norm = sphere.len();
        for start_node in sphere.iter().copied() {
            visited.fill(false);
            visited[start_node] = true;

            // Use a Deque structure to keep track of nodes still to visit and their distance
            // TODO: Figure out good initial value (8*delta is ideal for square grid)
            let mut queue = VecDeque::with_capacity(delta as usize * 8);
            queue.push_back((start_node, 0));

            // Perform breadth-first search
            loop {
                let Some((node, r)) = queue.pop_front() else {
                    break; // Stop if queue is empty (all nodes have been searched)
                };
                if r >= 2 * delta {
                    break; // All nodes should have been found
                }
                for &nbr in graph.get_neighbours(node) {
                    let visited_nbr = &mut visited[nbr];
                    if *visited_nbr {
                        continue; // Do not reconsider already visited node
                    }
                    *visited_nbr = true;
                    queue.push_back((nbr, r + 1));

                    // Compute ASD contribution
                    if dist_sphere[nbr] != delta {
                        continue; // Only contribute nodes on delta sphere
                    }

                    // Use one-pass algorithm
                    norm += 1;
                    let diff = (r + 1) as f64 - asd_delta;
                    asd_delta += diff / (norm as f64);
                    // Stop BFS is all nodes in `delta` sphere have been reached
                }
            }
        }
        asd.push(asd_delta);
    }

    asd
}

/// Returns the average sphere distance of a single sphere.
///
/// The `Vec` corresponds to the average sphere distance at delta = 1, 2, 3, ..., delta_max.
/// This implementation makes use of multiple BFSs in chunks and would be suited for a parallel BFS
/// be it on GPU or CPU.
pub fn asd_multi(graph: &Graph, origin: usize, delta_max: u16) -> Vec<f64> {
    // You can choose this such that distances will fit in memory
    // On CPU it can be significantly smaller, it does not matter much
    const MAX_CHUNCK_SIZE: usize = 50;

    // Establish spheres
    let bfs0 = bfs(graph, origin, delta_max);
    let spheres = &bfs0.spheres;
    let ball_volumes = bfs0.get_ball_volumes();

    let mut asd = vec![0f64; delta_max as usize];
    let nchunks = spheres.len().div_ceil(MAX_CHUNCK_SIZE);
    let nsize = graph.len();
    let mut distances = std::iter::repeat_n(u16::MAX, MAX_CHUNCK_SIZE * nsize).collect_vec();

    for k in 0..nchunks {
        let chunck_start = k * MAX_CHUNCK_SIZE;
        let chunck_end = usize::min((k + 1) * MAX_CHUNCK_SIZE, spheres.len());
        let chunck_size = chunck_end - chunck_start;

        if k > 0 {
            distances.fill(u16::MAX);
        }
        // Get partial distance matrix
        multi_bfs_with_dist(
            graph,
            &spheres[chunck_start..chunck_end],
            2 * delta_max,
            &mut distances,
        );

        for i in 0..chunck_size {
            let inode = spheres[chunck_start + i];
            let delta = bfs0.distance_list[inode] as usize;
            if delta == 0 {
                debug!("Delta = 0 origin point found, skipping");
                continue;
            }
            let sphere_start = ball_volumes[delta - 1];
            let sphere_end = ball_volumes[delta];
            #[allow(clippy::needless_range_loop)]
            for j in sphere_start..sphere_end {
                let jnode = spheres[j];
                let r = distances[i * nsize + jnode];
                asd[delta - 1] += r as f64;
            }
        }
    }
    asd.into_iter()
        .zip(bfs0.sphere_volumes.into_iter().skip(1))
        .map(|(asd, svol)| asd / (svol as f64).powi(2))
        .collect()
}

/// Returns the two-sphere ASD between two given spheres
///
/// The first sphere is given by its nodes `sphere1` and the second sphere is given by its origin
/// `origin2`, its volume `vol_sphere2` and the distance list of its interior `dist_sphere2`
///
/// Note: this computation is the most effective if `sphere1` is the smaller sphere
pub fn asd_double_delta(
    graph: &Graph,
    origin2: usize,
    sphere1: &[usize],
    vol_sphere2: usize,
    dist_sphere2: &[u16],
    delta: u16,
) -> f64 {
    // Compute ASD for origin1 and origin2
    let mut asd = delta as f64;
    let mut norm = vol_sphere2;
    for start_node in sphere1.iter().copied() {
        if start_node == origin2 {
            continue;
        }
        let mut visited = vec![false; graph.len()];
        visited[start_node] = true;

        let mut queue = VecDeque::with_capacity(delta as usize * 8);
        queue.push_back((start_node, 0));
        if dist_sphere2[start_node] == delta {
            norm += 1;
            asd -= asd / (norm as f64);
        }

        'bfs: loop {
            let Some((node, r)) = queue.pop_front() else {
                break; // Stop if queue is empty (all nodes have been searched)
            };
            if r >= 3 * delta {
                break; // All nodes should have been found
            }
            for &nbr in graph.get_neighbours(node) {
                let visited_nbr = &mut visited[nbr];
                if *visited_nbr {
                    continue; // Do not reconsider already visited node
                }
                *visited_nbr = true;
                queue.push_back((nbr, r + 1));

                // Compute ASD contribution
                if dist_sphere2[nbr] != delta {
                    continue; // Only contribute nodes on delta sphere
                }

                // Use one-pass algorithm
                norm += 1;
                asd += ((r + 1) as f64 - asd) / (norm as f64);

                // Exit early if all points have been found
                if norm >= vol_sphere2 * sphere1.len() {
                    break 'bfs;
                }
            }
        }
    }
    assert_eq!(norm, vol_sphere2 * sphere1.len());
    asd
}

/// Returns the two-sphere average sphere distance.
///
/// The ASD is computed with one sphere at the given `origin` and the other is sampled randomly
/// from the found sphere. Note that this means the point-pairs are not sampled randomly, but
/// with a probability proportional to the inverse of the sphere volume.
///
/// The resulting `Vec` corresponds to the average sphere distance at delta = 1, 2, 3, ..., delta_max.
pub fn asd_double<R: Rng + ?Sized>(
    graph: &Graph,
    origin: usize,
    delta_max: u16,
    rng: &mut R,
) -> Vec<f64> {
    let bfs_base = bfs(graph, origin, delta_max);
    let dist_sphere = &bfs_base.distance_list[..];
    // let (dist_sphere, spheres, _, ball_vol) = bfs(graph, origin, delta_max);

    let mut asd = Vec::with_capacity(delta_max as usize);

    for delta in 1..=delta_max {
        let sphere_origin = bfs_base.get_sphere(delta);

        let origin_other = *sphere_origin
            .choose(rng)
            .unwrap_or_else(|| panic!("There are no nodes at distance {}", delta));
        let bfs_sphere_other = bfs_sphere_sizehint(graph, origin_other, delta, sphere_origin.len());
        let dist_sphere_other = &bfs_sphere_other.distance_list[..];
        let sphere_other = &bfs_sphere_other.sphere[..];
        // Order the spheres such that BFS is performed from the smallest sphere
        let ((_, dist_sphere2), (sphere1, sphere2), (_, origin2)) =
            if sphere_origin.len() <= sphere_other.len() {
                (
                    (dist_sphere, dist_sphere_other),
                    (sphere_origin, sphere_other),
                    (origin, origin_other),
                )
            } else {
                (
                    (dist_sphere_other, dist_sphere),
                    (sphere_other, sphere_origin),
                    (origin_other, origin),
                )
            };
        let asd_delta =
            asd_double_delta(graph, origin2, sphere1, sphere2.len(), dist_sphere2, delta);

        asd.push(asd_delta);
    }

    asd
}

/// Returns and estimate of the average sphere distance of a single sphere by sampling.
///
/// The `Vec` corresponds to the average sphere distance at delta = 1, 2, 3, ..., delta_max.
/// The sampling uses `sample_size` number of starting nodes in each `delta` sphere to estimate
/// the ASD at that `delta`.
///
/// Note: if the `sample_size` is larger that the volume of the sphere, all nodes are used and the
/// result is exact. Also the sampling is performed without repetition.
pub fn asd_sampled<R: Rng + ?Sized>(
    graph: &Graph,
    origin: usize,
    delta_max: u16,
    sample_size: usize,
    rng: &mut R,
) -> (Vec<f64>, Vec<f64>) {
    // Construct starting spheres for all delta up to delta_max
    let bfs_result = bfs(graph, origin, delta_max);
    let sphere_idx = bfs_result.sphere_idx;
    let dist_sphere = bfs_result.distance_list;
    let mut spheres = bfs_result.spheres;

    let mut asd = Vec::with_capacity(delta_max as usize);
    let mut asd_var = Vec::with_capacity(delta_max as usize);

    let mut visited = vec![false; graph.len()];
    for delta in 1..=delta_max {
        let sphere = &mut spheres[sphere_idx[delta as usize]..sphere_idx[delta as usize + 1]];
        let exact = sphere.len() <= sample_size;
        let (nodes, _) = sphere.partial_shuffle(rng, sample_size);
        // Initialize ASD (make use of one-pass average algorithm with norm)
        let mut asd_delta = 0f64;
        let mut asd_var_delta = 0f64;
        // Do not start at 0 but at sphere_vol to account for the missed 0 distance geodesics
        let mut norm = nodes.len();
        for start_node in nodes.iter().copied() {
            visited.fill(false);
            visited[start_node] = true;

            // Use a Deque structure to keep track of nodes still to visit and their distance
            // TODO: Figure out good initial value (8*delta is ideal for square grid)
            let mut queue = VecDeque::with_capacity(delta as usize * 8);
            queue.push_back((start_node, 0));

            // Perform breadth-first search
            loop {
                let Some((node, r)) = queue.pop_front() else {
                    break; // Stop if queue is empty (all nodes have been searched)
                };
                if r >= 2 * delta {
                    break; // All nodes should have been found
                }
                for &nbr in graph.get_neighbours(node) {
                    let visited_nbr = &mut visited[nbr];
                    if *visited_nbr {
                        continue; // Do not reconsider already visited node
                    }
                    *visited_nbr = true;
                    queue.push_back((nbr, r + 1));

                    // Compute ASD contribution
                    if dist_sphere[nbr] != delta {
                        continue; // Only contribute nodes on delta sphere
                    }

                    // Use one-pass algorithm
                    norm += 1;
                    let diff = (r + 1) as f64 - asd_delta;
                    let diffn = diff / (norm as f64);
                    asd_delta += diffn;
                    asd_var_delta += diffn * diff * ((norm - 1) as f64);
                }
            }
        }
        asd.push(asd_delta);
        if exact {
            asd_var.push(0.0)
        } else {
            let sample_var = asd_var_delta / ((norm - 1) as f64);
            // Store the variance of the mean
            asd_var.push(sample_var / (norm as f64));
        }
    }

    (asd, asd_var)
}

/// Returns the average distance between each point on the sphere and the sphere.
///
/// The outer `Vec` corresponds to the average distances at delta = 1, 2, 3, ..., delta_max.
/// The inner `Vec` corresponds to each point in the sphere at the corresponing delta, where
/// the value is the average distance of this point to the sphere at the same delta.
///
/// Note: if one takes the average of all values at a given delta, this is exactly what
/// the average sphere distance is. This is used to see how the distance is distributed
/// and thus judge how well sampling works.
pub fn asd_distr(graph: &Graph, origin: usize, delta_max: u16) -> Vec<Vec<f64>> {
    // Construct starting spheres for all delta up to delta_max
    let bfs_results = bfs(graph, origin, delta_max);
    let dist_sphere = bfs_results.distance_list;
    let spheres = bfs_results.spheres;
    let sphere_vol = bfs_results.sphere_volumes;

    // Initialize ASD vector (make use of one-pass average algorithm with norm)
    let mut asd = (1..=(delta_max as usize))
        .map(|delta| Vec::with_capacity(sphere_vol[delta]))
        .collect_vec();

    let mut visited = vec![false; graph.len()];
    for start_node in spheres {
        // Get the delta of the current starting point
        let delta = dist_sphere[start_node];
        // ASD for delta=0 is not well-defined, skip
        if delta == 0 {
            continue;
        }

        visited.fill(false);
        visited[start_node] = true;

        // Use a Deque structure to keep track of nodes still to visit and their distance
        // TODO: Figure out good initial value (8*delta is ideal for square grid)
        let mut queue = VecDeque::with_capacity(delta as usize * 8);
        queue.push_back((start_node, 0));

        // Get (mutable) references to values corresponding to current `delta`
        let mut current_asd: f64 = 0.0;
        // Start at 1 to account for the missed geodesic
        let mut norm: usize = 1;
        let svol = *sphere_vol.get(delta as usize).unwrap();

        // Perform breadth-first search
        'bfs: loop {
            let Some((node, r)) = queue.pop_front() else {
                break; // Stop if queue is empty (all nodes have been searched)
            };
            for &nbr in graph.get_neighbours(node) {
                let visited_nbr = &mut visited[nbr];
                if *visited_nbr {
                    continue; // Do not reconsider already visited node
                }
                *visited_nbr = true;
                queue.push_back((nbr, r + 1));

                // Compute ASD contribution
                if dist_sphere[nbr] != delta {
                    continue; // Only contribute nodes on delta sphere
                }
                // Use one-pass algorithm
                norm += 1;
                current_asd += ((r + 1) as f64 - current_asd) / (norm as f64);
                // Stop BFS is all nodes in `delta` sphere have been reached
                if norm >= svol {
                    break 'bfs;
                }
            }
        }

        asd[delta as usize - 1].push(current_asd);
    }

    asd
}

/// Returns the average distance between each point on the smallest sphere,
/// and the other sphere, determined from origin.
///
/// This does the some thing as [`asd_distr()`] but for [`asd_double()`]
pub fn asd_distr_double<R: Rng + ?Sized>(
    graph: &Graph,
    origin: usize,
    delta_max: u16,
    rng: &mut R,
) -> Vec<Vec<f64>> {
    let bfs_base = bfs(graph, origin, delta_max);
    let dist_sphere = &bfs_base.distance_list[..];

    let mut asd = Vec::with_capacity(delta_max as usize);

    for delta in 1..=delta_max {
        let sphere_origin = bfs_base.get_sphere(delta);

        let origin_other = *sphere_origin
            .choose(rng)
            .unwrap_or_else(|| panic!("There are no nodes at distance {}", delta));
        let bfs_sphere_other = bfs_sphere_sizehint(graph, origin_other, delta, sphere_origin.len());
        let dist_sphere_other = &bfs_sphere_other.distance_list[..];
        let sphere_other = &bfs_sphere_other.sphere[..];
        // Order the spheres such that BFS is performed from the smallest sphere
        let ((_, dist_sphere2), (sphere1, sphere2), (_, origin2)) =
            if sphere_origin.len() <= sphere_other.len() {
                (
                    (dist_sphere, dist_sphere_other),
                    (sphere_origin, sphere_other),
                    (origin, origin_other),
                )
            } else {
                (
                    (dist_sphere_other, dist_sphere),
                    (sphere_other, sphere_origin),
                    (origin_other, origin),
                )
            };
        let asd_delta =
            asd_distr_double_delta(graph, origin2, sphere1, sphere2.len(), dist_sphere2, delta);

        asd.push(asd_delta);
    }
    asd
}

/// Returns the average distance between each point on sphere1, and sphere2.
///
/// See [`asd_distr()`] for more details.
pub fn asd_distr_double_delta(
    graph: &Graph,
    origin2: usize,
    sphere1: &[usize],
    vol_sphere2: usize,
    dist_sphere2: &[u16],
    delta: u16,
) -> Vec<f64> {
    // Compute ASD for origin1 and origin2
    let mut asd: Vec<f64> = Vec::with_capacity(sphere1.len());
    asd.push(delta as f64);
    for start_node in sphere1.iter().copied() {
        if start_node == origin2 {
            continue;
        }
        let mut current_asd: f64 = 0.0;
        let mut norm: usize = 0;

        let mut visited = vec![false; graph.len()];
        visited[start_node] = true;

        let mut queue = VecDeque::with_capacity(delta as usize * 8);
        queue.push_back((start_node, 0));
        if dist_sphere2[start_node] == delta {
            norm += 1;
            current_asd -= current_asd / (norm as f64);
        }

        'bfs: loop {
            let Some((node, r)) = queue.pop_front() else {
                break; // Stop if queue is empty (all nodes have been searched)
            };
            if r >= 3 * delta {
                break; // All nodes should have been found
            }
            for &nbr in graph.get_neighbours(node) {
                let visited_nbr = &mut visited[nbr];
                if *visited_nbr {
                    continue; // Do not reconsider already visited node
                }
                *visited_nbr = true;
                queue.push_back((nbr, r + 1));

                // Compute ASD contribution
                if dist_sphere2[nbr] != delta {
                    continue; // Only contribute nodes on delta sphere
                }

                // Use one-pass algorithm
                norm += 1;
                current_asd += ((r + 1) as f64 - current_asd) / (norm as f64);

                // Exit early if all points have been found
                if norm >= vol_sphere2 * sphere1.len() {
                    break 'bfs;
                }
            }
        }
        asd.push(current_asd);
    }
    asd
}