fast-nnt 0.4.0

use anyhow::Result;
use clap::Args;
use fixedbitset::FixedBitSet;
use log::info;
use ndarray::Array2;
use std::{
    cmp::{Ordering, max},
    time::Instant,
};

use super::band::{BandIndex, BandKernelPerf, avg_band_time_us, band_to_row, row_to_band};
use super::cgnr::eval_gradient_band;
use super::nnls_ops::sum_array_squared_indexed;
use crate::splits::asplit::ASplit;

pub use super::cgnr::{cgnr, sum_array_squared_band};
pub use super::nnls_ops::{
    EPS, PackedTriIndex, Scratch, calc_ainv_y, calc_atx_indexed, get_active_entries,
    mask_elements_branchless, sum_array_squared, zero_negative_entries,
};

pub trait Progress {
    fn check_for_cancel(&self) -> Result<()>;
}

#[derive(Args, Clone, Debug)]
pub struct NNLSParams {
    /// Include only split weights > cutoff (trivial splits always included).
    #[arg(long, default_value = "1e-4")]
    pub cutoff: f64,
    /// Stop if squared projected gradient < this (0 = derive from the data).
    #[arg(long, default_value = "0")]
    pub proj_grad_bound: f64,
    /// Hard iteration cap (outer loops).
    #[arg(long, default_value = "5000")]
    pub max_iterations: usize,
    /// Wall-clock cap in ms (use `u64::MAX` to disable).
    #[arg(long, default_value_t = u64::MAX)]
    pub max_time_ms: u64,
    /// CGNR inner iteration cap (0 = derive from the problem size).
    #[arg(long, default_value = "0")]
    pub cgnr_iterations: usize,
    /// CGNR convergence tolerance (0 = half the projected-gradient bound).
    #[arg(long, default_value = "0")]
    pub cgnr_tolerance: f64,
    /// Fraction of offending (negative) coords to push into the active set per correction.
    #[arg(long, default_value = "0.4")]
    pub active_set_rho: f64,
}

impl Default for NNLSParams {
    fn default() -> Self {
        Self {
            cutoff: 1e-4,
            proj_grad_bound: 0.0,
            max_iterations: 5000,
            max_time_ms: u64::MAX,
            cgnr_iterations: 0,
            cgnr_tolerance: 0.0,
            active_set_rho: 0.4,
        }
    }
}

/// Flattened weights (mainly for validation and debugging).
#[derive(Clone, Debug, PartialEq)]
pub struct SplitWeights {
    /// x[(i,j)] for 1 <= i < j <= n, flattened by blocks (1,2..n), (2,3..n), ...
    pub x: Vec<f64>,
}

/// - `cycle`: NeighborNet circular order (1-based with leading `0` sentinel)
/// - `distances`: symmetric `n x n` matrix in original taxon indices (0-based)
pub fn compute_asplits(
    cycle: &[usize],
    distances: &Array2<f64>,
    params: &mut NNLSParams,
    progress: Option<&dyn Progress>,
) -> Result<Vec<ASplit>> {
    let (_weights, pairs) = compute_use_1d(cycle, distances, params, progress)?;
    let n = cycle.len() - 1;

    let splits = pairs
        .into_iter()
        .map(|(a, w)| ASplit::from_a_ntax_with_weight(a, n, w))
        .collect();

    Ok(splits)
}

/// Core solver (vector form), also returns `(A_bitset, weight)` for building ASplit.
pub fn compute_use_1d(
    cycle: &[usize],
    distances: &Array2<f64>,
    params: &mut NNLSParams,
    progress: Option<&dyn Progress>,
) -> Result<(SplitWeights, Vec<(FixedBitSet, f64)>)> {
    let n = cycle.len() - 1;
    if n == 0 || n == 1 {
        return Ok((SplitWeights { x: vec![] }, vec![]));
    }
    if n == 2 {
        let d_12 = distances[[cycle[1] - 1, cycle[2] - 1]];
        let mut a = FixedBitSet::with_capacity(n + 1);
        a.grow(n + 1);
        a.set(cycle[1], true);
        return Ok((SplitWeights { x: vec![d_12] }, vec![(a, d_12.max(0.0))]));
    }

    let npairs = n * (n - 1) / 2;
    let tri_idx = PackedTriIndex::new(n);
    let mut d = vec![0.0; npairs];
    {
        let mut idx = 0usize;
        for i in 1..=n {
            for j in (i + 1)..=n {
                d[idx] = distances[[cycle[i] - 1, cycle[j] - 1]];
                idx += 1;
            }
        }
    }

    // Reuse `x` as workspace for A^T d to avoid an extra allocation.
    let mut x = vec![0.0; npairs];
    calc_atx_indexed(&d, &mut x, &tri_idx);
    if params.proj_grad_bound <= 0.0 {
        let norm_atd = sum_array_squared_indexed(&x, &tri_idx).sqrt();
        params.proj_grad_bound = (1e-4 * norm_atd).powi(2);
    }

    calc_ainv_y(&d, &mut x, n);
    let min_val = x.iter().copied().fold(f64::INFINITY, f64::min);
    if min_val < 0.0 {
        let start = Instant::now();
        zero_negative_entries(&mut x);
        let mut active = vec![false; npairs];
        get_active_entries(&x, &mut active);

        let mut scratch = Scratch::new(npairs);
        let mut splits_idx = vec![0usize; npairs];
        let mut order_idx: Vec<usize> = (0..npairs).collect();
        let mut vals = vec![0.0; npairs];
        let mut xstar = vec![0.0; npairs];

        if params.cgnr_tolerance <= 0.0 {
            params.cgnr_tolerance = params.proj_grad_bound / 2.0;
        }
        if params.cgnr_iterations == 0 {
            // A high inner-iteration cap helps small problems converge, but
            // makes large inputs stall.
            params.cgnr_iterations = if n <= 1_000 { max(50, npairs) } else { 50 };
        }

        let band = BandIndex::new(n);
        let mut d_band = vec![0.0; npairs];
        row_to_band(&d, &mut d_band, &band);

        let mut x_band = vec![0.0; npairs];
        row_to_band(&x, &mut x_band, &band);

        let mut active_band = vec![false; npairs];
        for (row_idx, &band_idx) in band.row_to_band_map.iter().enumerate() {
            active_band[band_idx] = active[row_idx];
        }

        active_set_method(
            &mut x_band,
            &mut xstar,
            &d_band,
            &tri_idx,
            params,
            &mut active_band,
            &mut scratch,
            &mut splits_idx,
            &mut order_idx,
            &mut vals,
            progress,
            start,
            &band,
        )?;

        band_to_row(&x_band, &mut x, &band);
    }

    let mut out_pairs = Vec::with_capacity(npairs);
    let mut idx = 0usize;
    for i in 1..=n {
        let mut a = FixedBitSet::with_capacity(n + 1);
        a.grow(n + 1);
        for j in (i + 1)..=n {
            a.set(cycle[j - 1], true);
            let w = x[idx].max(0.0);

            let size_a = a.ones().filter(|&t| t != 0).count();
            if w > params.cutoff || size_a == 1 || size_a == n - 1 {
                out_pairs.push((a.clone(), w));
            }
            idx += 1;
        }
    }

    Ok((SplitWeights { x }, out_pairs))
}

#[allow(clippy::too_many_arguments)] // mirrors the reference solver's parameter surface
fn active_set_method(
    x: &mut [f64],
    xstar: &mut [f64],
    d: &[f64],
    tri_idx: &PackedTriIndex,
    params: &mut NNLSParams,
    active: &mut [bool],
    scratch: &mut Scratch,
    tmp_splits: &mut [usize],
    idx_sorted: &mut [usize],
    vals: &mut [f64],
    progress: Option<&dyn Progress>,
    started: Instant,
    band: &BandIndex,
) -> Result<()> {
    let n = tri_idx.n;
    let npairs = x.len();
    debug_assert_eq!(npairs, tri_idx.npairs);
    let mut k_outer = 0usize;
    let outer_started = Instant::now();
    let mut next_info_iter = 10usize;
    let small_problem = npairs <= 4096;
    let track_kernel_perf = !small_problem;
    let mut kernel_perf = BandKernelPerf::default();
    let mut row_sums = vec![0.0; n];
    let mut shifted = vec![0.0; n];

    let start = Instant::now();
    loop {
        debug!("Outer Loop {}", k_outer);
        loop {
            xstar[..npairs].copy_from_slice(x);

            let iters = cgnr(
                xstar,
                d,
                active,
                params.cgnr_iterations,
                params.cgnr_tolerance,
                scratch,
                tri_idx,
                &mut kernel_perf,
                track_kernel_perf,
                progress,
                started,
                params.max_time_ms,
                band,
                &mut row_sums,
            )?;

            k_outer += 1;

            let ok = feasible_move_active_set(
                x,
                &xstar[..npairs],
                active,
                tmp_splits,
                idx_sorted,
                vals,
                n,
                params,
            );

            debug!(
                "\t{}\t{}\t{}\t{}",
                k_outer,
                params.cgnr_iterations,
                params.active_set_rho,
                (Instant::now() - start).as_millis()
            );

            if ok && iters < params.cgnr_iterations {
                break;
            }
            if k_outer > params.max_iterations {
                return Ok(());
            }
        }

        x.copy_from_slice(&xstar[..npairs]);
        let (p, r) = (&mut scratch.p, &mut scratch.r);

        if !small_problem && k_outer >= next_info_iter {
            info!(
                "NNLS progress: outer_iter={} elapsed={:.1}s active_fraction={:.3} ab_calls={} atx_calls={} ab_avg_us={:.1} atx_avg_us={:.1}",
                k_outer,
                outer_started.elapsed().as_secs_f64(),
                active.iter().filter(|&&a| a).count() as f64 / npairs as f64,
                kernel_perf.calc_ab_calls,
                kernel_perf.calc_atx_calls,
                avg_band_time_us(kernel_perf.calc_ab_time, kernel_perf.calc_ab_calls),
                avg_band_time_us(kernel_perf.calc_atx_time, kernel_perf.calc_atx_calls),
            );
            next_info_iter += 10;
        }

        eval_gradient_band(
            x,
            d,
            p,
            r,
            band,
            &mut row_sums,
            &mut shifted,
            &mut kernel_perf,
            track_kernel_perf,
        );

        p.iter_mut().zip(x.iter()).for_each(|(gi, &xi)| {
            if xi == 0.0 {
                // At the bound only components pointing into the feasible
                // region count toward the projected gradient.
                *gi = gi.min(0.0);
            }
        });

        let pg = sum_array_squared_band(p, band, tri_idx);
        if (pg - params.proj_grad_bound) < -EPS {
            return Ok(());
        }

        debug!(
            "Projected gradient squared = {}\t target = {}\tNumber iterations = {}",
            pg, params.proj_grad_bound, k_outer
        );

        let mut imin = 0usize;
        let mut pmin = 0.0f64;
        for i in 0..npairs {
            if active[i] && p[i] < pmin {
                pmin = p[i];
                imin = i;
            }
        }
        if pmin < 0.0 {
            active[imin] = false;
        }
    }
}

/// Returns `true` if `xstar` has no negative entries (it is copied into `x`);
/// otherwise marks the worst offenders active, moves `x` toward `xstar` by the
/// smallest blocking step, and returns `false`.
#[allow(clippy::too_many_arguments)] // mirrors the reference solver's parameter surface
pub fn feasible_move_active_set(
    x: &mut [f64],
    xstar: &[f64],
    active_set: &mut [bool],
    splits: &mut [usize],
    indices: &mut [usize],
    vals: &mut [f64],
    n: usize,
    params: &NNLSParams,
) -> bool {
    let npairs = n * (n - 1) / 2;

    assert_eq!(x.len(), npairs, "x len must be n*(n-1)/2");
    assert_eq!(xstar.len(), npairs, "xstar len must be n*(n-1)/2");
    assert_eq!(active_set.len(), npairs, "active_set len must be n*(n-1)/2");
    assert!(
        splits.len() >= npairs && indices.len() >= npairs && vals.len() >= npairs,
        "splits/indices/vals must have capacity >= n*(n-1)/2"
    );

    let mut count = 0usize;
    for i in 0..npairs {
        if xstar[i] < 0.0 {
            vals[count] = x[i] / (x[i] - xstar[i]);
            splits[count] = i;
            indices[count] = count;
            count += 1;
        }
    }

    if count == 0 {
        x.copy_from_slice(xstar);
        return true;
    }

    let mut min_idx = 0usize;
    for i in 1..count {
        if cmp_vals_indices(vals, i, min_idx) == Ordering::Less {
            min_idx = i;
        }
    }
    let tmin = vals[min_idx];

    let num_to_make_active = max(1, (count as f64 * params.active_set_rho).ceil() as usize);

    if num_to_make_active >= count {
        for &pos in &splits[..count] {
            active_set[pos] = true;
        }
    } else {
        let kth = num_to_make_active - 1;
        let (left, pivot, _right) =
            indices[..count].select_nth_unstable_by(kth, |&a, &b| cmp_vals_indices(vals, a, b));
        for &slot in left.iter() {
            active_set[splits[slot]] = true;
        }
        active_set[splits[*pivot]] = true;
    }

    for i in 0..npairs {
        if active_set[i] {
            x[i] = 0.0;
        } else {
            x[i] = (1.0 - tmin) * x[i] + tmin * xstar[i];
        }
    }

    false
}

#[inline]
fn cmp_vals_indices(vals: &[f64], a: usize, b: usize) -> Ordering {
    let va = vals[a];
    let vb = vals[b];
    let ord = match (va.is_nan(), vb.is_nan()) {
        (true, true) => Ordering::Equal,
        (true, false) => Ordering::Greater,
        (false, true) => Ordering::Less,
        (false, false) => va.total_cmp(&vb),
    };
    if ord == Ordering::Equal {
        // Tie-break by insertion order for determinism.
        a.cmp(&b)
    } else {
        ord
    }
}