fast-nnt 0.4.0

use std::time::{Duration, Instant};

use anyhow::Result;
use log::info;
use rayon::prelude::*;

use super::band::{
    BandIndex, BandKernelPerf, calculate_ab_band, calculate_ab_band_with_masked_norm_sq,
    calculate_forward_band, calculate_forward_band_with_norm_sq,
};
use super::nnls_ops::{EPS, PackedTriIndex, Scratch, zero_negative_entries};
use crate::weights::active_set_weights::Progress;

#[doc(hidden)]
#[allow(clippy::too_many_arguments)] // mirrors the reference solver's parameter surface
pub fn cgnr(
    x: &mut [f64],
    d: &[f64],
    active_set: &[bool],
    max_iters: usize,
    tol: f64,
    scratch: &mut Scratch,
    tri_idx: &PackedTriIndex,
    kernel_perf: &mut BandKernelPerf,
    track_kernel_perf: bool,
    progress: Option<&dyn Progress>,
    started: Instant,
    max_time_ms: u64,
    band: &BandIndex,
    row_sums: &mut [f64],
) -> Result<usize> {
    let n = tri_idx.n;
    let (p, r, z, w) = (
        &mut scratch.p,
        &mut scratch.r,
        &mut scratch.z,
        &mut scratch.w,
    );

    zero_negative_entries(x);
    let free_indices = collect_free_indices(active_set);
    if free_indices.is_empty() {
        return Ok(1);
    }
    let nfree = free_indices.len();
    let mut p_free = vec![0.0; nfree];
    let mut z_free = vec![0.0; nfree];

    let mut row_sq_accum = vec![0.0f64; n];
    let mut shifted = vec![0.0f64; n];

    profiled_forward_band(
        x,
        r,
        band,
        row_sums,
        &mut shifted,
        kernel_perf,
        track_kernel_perf,
    );
    r.iter_mut()
        .zip(d.iter())
        .for_each(|(ri, &di)| *ri = di - *ri);

    let mut zts = profiled_adjoint_band_with_masked_norm(
        r,
        z,
        active_set,
        band,
        row_sums,
        &mut row_sq_accum,
        &mut shifted,
        kernel_perf,
        track_kernel_perf,
    );
    gather_by_index(z, &free_indices, &mut z_free);
    p_free.copy_from_slice(&z_free);

    if zts <= EPS || !zts.is_finite() {
        return Ok(1);
    }
    let mut k = 1usize;
    let log_cgnr_progress = max_iters >= 1024 && x.len() > 1_000_000;

    // Non-free entries of p stay zero, so one fill suffices.
    p.fill(0.0);

    loop {
        scatter_by_index(&p_free, &free_indices, p);

        let denom = profiled_forward_band_with_norm(
            p,
            w,
            band,
            row_sums,
            &mut row_sq_accum,
            &mut shifted,
            kernel_perf,
            track_kernel_perf,
        );
        if denom <= EPS || !denom.is_finite() {
            break;
        }
        let alpha = zts / denom;

        for (slot, &idx) in free_indices.iter().enumerate() {
            x[idx] += alpha * p_free[slot];
        }
        for i in 0..r.len() {
            r[i] -= alpha * w[i];
        }

        let zts_new = profiled_adjoint_band_with_masked_norm(
            r,
            z,
            active_set,
            band,
            row_sums,
            &mut row_sq_accum,
            &mut shifted,
            kernel_perf,
            track_kernel_perf,
        );
        gather_by_index(z, &free_indices, &mut z_free);
        if zts_new <= EPS || !zts_new.is_finite() {
            break;
        }
        let beta = zts_new / zts;
        if log_cgnr_progress && k.is_multiple_of(256) {
            info!(
                "NNLS CGNR progress: iter={}/{} residual={:.3e} elapsed={:.1}s",
                k,
                max_iters,
                zts_new,
                started.elapsed().as_secs_f64()
            );
        }
        if zts_new < tol || k >= max_iters {
            break;
        }

        for i in 0..nfree {
            p_free[i] = z_free[i] + beta * p_free[i];
        }

        zts = zts_new;
        k += 1;

        if let Some(pl) = progress
            && k.is_multiple_of(n)
        {
            pl.check_for_cancel()?;
        }
        if started.elapsed() >= Duration::from_millis(max_time_ms) {
            break;
        }
    }
    Ok(k)
}

// The forward and adjoint operators accumulate row sums in different orders;
// pairing them this way keeps results bit-identical to the row-major kernels.
#[allow(clippy::too_many_arguments)] // threads shared solver state through a hot path
pub(crate) fn eval_gradient_band(
    x: &[f64],
    d: &[f64],
    gradient: &mut [f64],
    residual: &mut [f64],
    band: &BandIndex,
    row_sums: &mut [f64],
    shifted: &mut [f64],
    kernel_perf: &mut BandKernelPerf,
    track_kernel_perf: bool,
) {
    profiled_forward_band(
        x,
        residual,
        band,
        row_sums,
        shifted,
        kernel_perf,
        track_kernel_perf,
    );
    residual
        .par_iter_mut()
        .zip(d.par_iter())
        .for_each(|(ri, &di)| *ri -= di);
    profiled_adjoint_band(
        residual,
        gradient,
        band,
        row_sums,
        shifted,
        kernel_perf,
        track_kernel_perf,
    );
}

/// Sum of squares in row-major order over band-major data.
#[doc(hidden)]
pub fn sum_array_squared_band(x: &[f64], band: &BandIndex, tri_idx: &PackedTriIndex) -> f64 {
    let mut total = 0.0f64;
    for i in 1..=tri_idx.n {
        let start = tri_idx.row_offsets[i];
        let len = tri_idx.n - i;
        let mut s_i = 0.0f64;
        for row_idx in start..start + len {
            let val = x[band.row_to_band_map[row_idx]];
            s_i += val * val;
        }
        total += s_i;
    }
    total
}

#[inline]
fn profiled_forward_band(
    x: &[f64],
    y: &mut [f64],
    band: &BandIndex,
    row_sums: &mut [f64],
    shifted: &mut [f64],
    perf: &mut BandKernelPerf,
    track: bool,
) {
    if track {
        let t0 = Instant::now();
        calculate_forward_band(x, y, band, row_sums, shifted);
        perf.calc_ab_calls += 1;
        perf.calc_ab_time += t0.elapsed();
    } else {
        calculate_forward_band(x, y, band, row_sums, shifted);
    }
}

#[inline]
fn profiled_adjoint_band(
    x: &[f64],
    y: &mut [f64],
    band: &BandIndex,
    row_sums: &mut [f64],
    shifted: &mut [f64],
    perf: &mut BandKernelPerf,
    track: bool,
) {
    if track {
        let t0 = Instant::now();
        calculate_ab_band(x, y, band, row_sums, shifted);
        perf.calc_atx_calls += 1;
        perf.calc_atx_time += t0.elapsed();
    } else {
        calculate_ab_band(x, y, band, row_sums, shifted);
    }
}

#[inline]
#[allow(clippy::too_many_arguments)] // thin timing shim over the band kernel
fn profiled_forward_band_with_norm(
    x: &[f64],
    y: &mut [f64],
    band: &BandIndex,
    row_sums: &mut [f64],
    row_sq_accum: &mut [f64],
    shifted: &mut [f64],
    perf: &mut BandKernelPerf,
    track: bool,
) -> f64 {
    if track {
        let t0 = Instant::now();
        let norm = calculate_forward_band_with_norm_sq(x, y, band, row_sums, row_sq_accum, shifted);
        perf.calc_ab_calls += 1;
        perf.calc_ab_time += t0.elapsed();
        norm
    } else {
        calculate_forward_band_with_norm_sq(x, y, band, row_sums, row_sq_accum, shifted)
    }
}

#[inline]
#[allow(clippy::too_many_arguments)] // thin timing shim over the band kernel
fn profiled_adjoint_band_with_masked_norm(
    x: &[f64],
    y: &mut [f64],
    active_set: &[bool],
    band: &BandIndex,
    row_sums: &mut [f64],
    row_sq_accum: &mut [f64],
    shifted: &mut [f64],
    perf: &mut BandKernelPerf,
    track: bool,
) -> f64 {
    if track {
        let t0 = Instant::now();
        let norm = calculate_ab_band_with_masked_norm_sq(
            x,
            y,
            band,
            row_sums,
            row_sq_accum,
            active_set,
            shifted,
        );
        perf.calc_atx_calls += 1;
        perf.calc_atx_time += t0.elapsed();
        norm
    } else {
        calculate_ab_band_with_masked_norm_sq(
            x,
            y,
            band,
            row_sums,
            row_sq_accum,
            active_set,
            shifted,
        )
    }
}

#[inline]
fn collect_free_indices(active_set: &[bool]) -> Vec<usize> {
    active_set
        .iter()
        .enumerate()
        .filter_map(|(i, &is_active)| (!is_active).then_some(i))
        .collect()
}

#[inline]
fn gather_by_index(src: &[f64], indices: &[usize], dst: &mut [f64]) {
    debug_assert_eq!(indices.len(), dst.len());
    for (slot, &idx) in indices.iter().enumerate() {
        dst[slot] = src[idx];
    }
}

#[inline]
fn scatter_by_index(src: &[f64], indices: &[usize], dst: &mut [f64]) {
    debug_assert_eq!(indices.len(), src.len());
    for (slot, &idx) in indices.iter().enumerate() {
        dst[idx] = src[slot];
    }
}