fast-nnt 0.4.0

use ndarray::Array2;
use petgraph::Undirected;
use petgraph::graph::{Graph, NodeIndex};
use std::cmp::min;
use std::time::Instant;

use super::pair_search::{Component, PairSearchState, avg_d_comp_comp, avg_d_p_comp};

/// Closest-pair circular ordering (Bryant & Huson 2023).
/// Input distances are 0-based (shape n x n), symmetric, zero diagonal.
/// Returns a 1-based cycle with a leading 0 sentinel.
pub fn compute_order_huson_2023(dist: &Array2<f64>) -> anyhow::Result<Vec<usize>> {
    let n_tax = dist.nrows();
    anyhow::ensure!(
        n_tax == dist.ncols(),
        "Distance matrix must be square ({}x{})",
        n_tax,
        dist.ncols()
    );
    if n_tax <= 3 {
        return Ok(create_array_upward_count(n_tax));
    }

    let mut graph: Graph<usize, (), Undirected> = Graph::default();
    let mut node_map: Vec<NodeIndex> = vec![NodeIndex::end(); n_tax + 1];
    let mut components: Vec<Component> = (1..=n_tax).map(Component::singleton).collect();
    for (t, slot) in node_map.iter_mut().enumerate().skip(1) {
        *slot = graph.add_node(t);
    }

    // 1-based padded copy: row/col 0 unused, mutated during agglomeration.
    let mut d = Array2::<f64>::zeros((n_tax + 1, n_tax + 1));
    for i in 1..=n_tax {
        for j in 1..=n_tax {
            d[[i, j]] = dist[[i - 1, j - 1]];
        }
    }
    let mut pair_search = PairSearchState::new(&components, &d);

    let ordering_started = Instant::now();
    let mut pair_search_sec = 0.0_f64;
    let mut iters = 0usize;
    let mut next_progress_report = n_tax.saturating_sub(250);
    while components.len() >= 2 {
        iters += 1;

        let t_pair = Instant::now();
        let (ip, iq) = pair_search.select_closest_pair(&components, &d);
        pair_search_sec += t_pair.elapsed().as_secs_f64();
        let m_old = components.len();
        let mut old_avg_ip = vec![0.0_f64; m_old];
        let mut old_avg_iq = vec![0.0_f64; m_old];
        for k in 0..m_old {
            if k == ip || k == iq {
                continue;
            }
            old_avg_ip[k] = avg_d_comp_comp(&d, &components[k], &components[ip]);
            old_avg_iq[k] = avg_d_comp_comp(&d, &components[k], &components[iq]);
        }

        let (p_comp, q_comp) = (components[ip], components[iq]);
        debug!("Selected: P={} Q={}", p_comp.first(), q_comp.first());
        debug!("Sizes: P={} Q={}", p_comp.size(), q_comp.size());
        if p_comp.size() == 1 && q_comp.size() == 1 {
            let p = p_comp.first();
            let q = q_comp.first();
            graph.add_edge(node_map[p], node_map[q], ());

            let new_component = Component::pair(p, q);
            components[ip] = new_component;
            debug!(
                "First case: P={} Q={} NewComponent={:?}",
                p, q, components[ip]
            );
            components.remove(iq);
        } else if p_comp.size() == 1 && q_comp.size() == 2 {
            let p = p_comp.first();
            let q = select_closest_1_vs_2(ip, iq, &d, &components);
            let qb = q_comp.other(q).expect("case 2: Q must be a pair");

            d[[p, qb]] = (d[[p, qb]] + d[[q, qb]] + d[[p, q]]) / 3.0;
            d[[qb, p]] = d[[p, qb]];

            for (i, comp) in components.iter().enumerate() {
                if i == ip || i == iq {
                    continue;
                }

                for &r_opt in comp.values().iter() {
                    if let Some(r) = r_opt
                        && r != p
                        && r != q
                        && r != qb
                    {
                        d[[p, r]] = (2.0 * d[[p, r]] + d[[q, r]]) / 3.0;
                        d[[r, p]] = d[[p, r]];
                        d[[qb, r]] = (2.0 * d[[qb, r]] + d[[q, r]]) / 3.0;
                        d[[r, qb]] = d[[qb, r]];
                    }
                }
            }

            graph.add_edge(node_map[p], node_map[q], ());
            let new_component = Component::pair(p, qb);
            components[ip] = new_component;
            debug!(
                "Second case: P={} Q={} NewComponent={:?}",
                p, q, components[ip]
            );
            components.remove(iq);
        } else if p_comp.size() == 2 && q_comp.size() == 2 {
            let (p, q) = select_closest_2_vs_2(ip, iq, &d, &components);
            let pb = p_comp.other(p).expect("case 3: P must be a pair");
            let qb = q_comp.other(q).expect("case 3: Q must be a pair");

            d[[pb, qb]] =
                (d[[pb, p]] + d[[pb, q]] + d[[pb, qb]] + d[[p, q]] + d[[p, qb]] + d[[q, qb]]) / 6.0;
            d[[qb, pb]] = d[[pb, qb]];

            for (i, &other) in components.iter().enumerate() {
                if i == ip || i == iq {
                    continue;
                }

                for &r_opt in other.values().iter() {
                    if let Some(r) = r_opt
                        && r != p
                        && r != q
                        && r != pb
                        && r != qb
                    {
                        let pb_r = d[[pb, r]] / 2.0 + d[[p, r]] / 3.0 + d[[q, r]] / 6.0;
                        let qb_r = d[[p, r]] / 6.0 + d[[q, r]] / 3.0 + d[[qb, r]] / 2.0;
                        d[[pb, r]] = pb_r;
                        d[[r, pb]] = pb_r;
                        d[[qb, r]] = qb_r;
                        d[[r, qb]] = qb_r;
                    }
                }
            }

            graph.add_edge(node_map[p], node_map[q], ());
            let new_component = Component::pair(pb, qb);
            debug!(
                "Third case: P={} Q={} NewComponent={:?}",
                p, q, new_component
            );
            components[ip] = new_component;
            components.remove(iq);
        } else {
            unreachable!(
                "Components can only be size 1 or 2, got |P|={} and |Q|={}",
                p_comp.size(),
                q_comp.size()
            );
        }
        pair_search.update_after_merge(&components, &d, ip, iq, &old_avg_ip, &old_avg_iq);

        if n_tax >= 1_000 && components.len() <= next_progress_report {
            info!(
                "Closest-pair ordering progress: {} components remaining (iteration {}, elapsed {:.1}s)",
                components.len(),
                iters,
                ordering_started.elapsed().as_secs_f64()
            );
            next_progress_report = components.len().saturating_sub(250);
        }
    }

    let p = components[0].first();
    let q = components[0].second();
    graph.add_edge(node_map[p], node_map[q], ());
    debug!("Final edge: P={} Q={}", p, q);
    info!(
        "Closest-pair ordering finished in {:.3}s over {} iterations (closest-pair search {:.3}s)",
        ordering_started.elapsed().as_secs_f64(),
        iters,
        pair_search_sec
    );

    Ok(extract_ordering(&graph, &node_map))
}

fn create_array_upward_count(n: usize) -> Vec<usize> {
    let mut v = Vec::with_capacity(n + 1);
    v.push(0);
    v.extend(1..=n);
    v
}

fn select_closest_1_vs_2(ip: usize, iq: usize, d: &Array2<f64>, components: &[Component]) -> usize {
    let m = components.len();
    let p = components[ip].first();
    let q1 = components[iq].first();
    let q2 = components[iq].second();

    // Serial accumulation keeps the FP order deterministic across runs.
    let mut p_r_extra = 0.0;
    let mut q1_r_extra = 0.0;
    let mut q2_r_extra = 0.0;
    for i in (0..m).filter(|&i| i != ip && i != iq) {
        let other = components[i];
        p_r_extra += avg_d_p_comp(d, p, &other);
        q1_r_extra += avg_d_p_comp(d, q1, &other);
        q2_r_extra += avg_d_p_comp(d, q2, &other);
    }

    let p_r = d[[q1, p]] + d[[q2, p]] + p_r_extra;
    let q1_r = d[[q1, q2]] + d[[q1, p]] + q1_r_extra;
    let q2_r = d[[q1, q2]] + d[[q2, p]] + q2_r_extra;

    let mm1 = m as f64 - 1.0;
    let q1p_adj = mm1 * d[[q1, p]] - q1_r - p_r;
    let q2p_adj = mm1 * d[[q2, p]] - q2_r - p_r;

    debug!(
        "1x2 @ P={} Q={} -> q1p_adj:{:.9} q2p_adj:{:.9}",
        p, q1, q1p_adj, q2p_adj
    );

    if q1p_adj <= q2p_adj { q1 } else { q2 }
}

fn select_closest_2_vs_2(
    ip: usize,
    iq: usize,
    d: &Array2<f64>,
    components: &[Component],
) -> (usize, usize) {
    let m = components.len();

    let p1 = components[ip].first();
    let p2 = components[ip].second();
    let q1 = components[iq].first();
    let q2 = components[iq].second();

    // Serial accumulation keeps the FP order deterministic across runs.
    let mut p1_r_extra = 0.0;
    let mut p2_r_extra = 0.0;
    let mut q1_r_extra = 0.0;
    let mut q2_r_extra = 0.0;
    for i in (0..m).filter(|&i| i != ip && i != iq) {
        let other = components[i];
        p1_r_extra += avg_d_p_comp(d, p1, &other);
        p2_r_extra += avg_d_p_comp(d, p2, &other);
        q1_r_extra += avg_d_p_comp(d, q1, &other);
        q2_r_extra += avg_d_p_comp(d, q2, &other);
    }

    let p1_r = d[[p1, q1]] + d[[p1, q2]] + p1_r_extra;
    let p2_r = d[[p2, q1]] + d[[p2, q2]] + p2_r_extra;
    let q1_r = d[[p1, q1]] + d[[p2, q1]] + q1_r_extra;
    let q2_r = d[[p1, q2]] + d[[p2, q2]] + q2_r_extra;

    let m_f = m as f64;
    let p1q1 = m_f * d[[p1, q1]] - p1_r - q1_r;
    let p2q1 = m_f * d[[p2, q1]] - p2_r - q1_r;
    let p1q2 = m_f * d[[p1, q2]] - p1_r - q2_r;
    let p2q2 = m_f * d[[p2, q2]] - p2_r - q2_r;

    debug!(
        "2x2 inputs P={{{}, {}}} Q={{{}, {}}} | \
        D: p1q1={:.12}, p2q1={:.12}, p1q2={:.12}, p2q2={:.12} | \
        R: p1R={:.12}, p2R={:.12}, q1R={:.12}, q2R={:.12}",
        p1,
        p2,
        q1,
        q2,
        d[[p1, q1]],
        d[[p2, q1]],
        d[[p1, q2]],
        d[[p2, q2]],
        p1_r,
        p2_r,
        q1_r,
        q2_r
    );

    match rank_of_min(&[p1q1, p2q1, p1q2, p2q2]) {
        0 => (p1, q1),
        1 => (p2, q1),
        2 => (p1, q2),
        _ => (p2, q2),
    }
}

#[inline]
fn rank_of_min(vals: &[f64]) -> usize {
    let mut idx = 0usize;
    for i in 1..vals.len() {
        if vals[i] < vals[idx] {
            idx = i;
        }
    }
    idx
}

/// Walk the degree-2 cycle graph from taxon 1, taking the smaller-labelled
/// neighbour first so the direction is deterministic.
fn extract_ordering(graph: &Graph<usize, (), Undirected>, node_map: &[NodeIndex]) -> Vec<usize> {
    let n = node_map.len() - 1;
    let mut order: Vec<usize> = Vec::with_capacity(n + 1);
    order.push(0);
    debug!("Node Map: {:?}", node_map);
    if n == 0 {
        return order;
    }
    if n <= 3 {
        for t in 1..=n {
            order.push(t);
        }
        return order;
    }

    let v1 = node_map[1];
    let neigh: Vec<NodeIndex> = graph.neighbors(v1).collect();
    assert!(
        neigh.len() == 2,
        "ordering graph must be a simple cycle: node 1 should have degree 2"
    );
    let mut prev = v1;
    let mut cur = min(neigh[0], neigh[1]);
    debug!(
        "Starting at v1={} cur={}: neighbours {:?}",
        graph[v1], graph[cur], neigh
    );

    order.push(graph[v1]);
    while order.len() - 1 < n {
        order.push(graph[cur]);
        let it = graph.neighbors(cur).collect::<Vec<_>>();
        debug!("prev {:?} -> cur {:?}: neighbours {:?}", prev, cur, it);
        let a = it[0];
        let b = it[1];
        let nxt = if a == prev { b } else { a };
        prev = cur;
        cur = nxt;
    }
    order
}