fast-nnt 0.4.0

use anyhow::Result;
use fixedbitset::FixedBitSet;
use log::info;
use ndarray::Array2;
use std::time::{Duration, Instant};

use super::band::{
    BandIndex, BandKernelPerf, avg_band_time_us, band_to_row, calculate_ab_band,
    calculate_atx_band, row_to_band,
};
use crate::splits::asplit::ASplit;

pub const DEFAULT_CUTOFF: f64 = 1e-6;

#[derive(Debug, Clone, serde::Serialize, Default)]
pub struct CGSolveStats {
    pub outer_iterations: usize,
    pub cg_calls: usize,
    pub cg_iters_total: usize,
    pub max_cg_iters_single_call: usize,
    pub stagnation_exits: usize,
    pub inner_loop_activations: usize,
    pub collapse_events: usize,
    pub collapsed_indices_total: usize,
    pub kernel_ab_calls: usize,
    pub kernel_atx_calls: usize,
    pub elapsed_secs: f64,
}

type KernelPerf = BandKernelPerf;

#[derive(Debug)]
struct SolveProgress {
    started: Instant,
    next_log: Instant,
    next_inner_log: Instant,
    outer_iter: usize,
    cg_calls: usize,
    cg_iters_total: usize,
    collapse_events: usize,
    collapsed_indices_total: usize,
    max_cg_iters_single_call: usize,
    inner_loop_activations: usize,
    stagnation_exits: usize,
}

impl SolveProgress {
    fn new(now: Instant) -> Self {
        Self {
            started: now,
            next_log: now + Duration::from_secs(5),
            next_inner_log: now + Duration::from_secs(10),
            outer_iter: 0,
            cg_calls: 0,
            cg_iters_total: 0,
            collapse_events: 0,
            collapsed_indices_total: 0,
            max_cg_iters_single_call: 0,
            inner_loop_activations: 0,
            stagnation_exits: 0,
        }
    }
}

#[derive(Clone, Copy, Debug)]
struct CGOptions {
    cutoff: f64,
}

impl Default for CGOptions {
    fn default() -> Self {
        Self {
            cutoff: DEFAULT_CUTOFF,
        }
    }
}

/// Split weights via active-set NNLS with a conjugate-gradient inner solver.
/// `cycle` is 1-based with a leading 0 sentinel.
pub fn compute_splits(
    cycle: &[usize],
    distances: &Array2<f64>,
) -> Result<(Vec<ASplit>, CGSolveStats)> {
    compute_weighted_splits(cycle, distances, CGOptions::default())
}

fn compute_weighted_splits(
    cycle: &[usize],
    distances: &Array2<f64>,
    options: CGOptions,
) -> Result<(Vec<ASplit>, CGSolveStats)> {
    let n = cycle.len().saturating_sub(1);
    if n <= 1 {
        return Ok((Vec::new(), CGSolveStats::default()));
    }
    if n == 2 {
        let d_ij = distances[[cycle[1] - 1, cycle[2] - 1]];
        if d_ij > 0.0 {
            let mut a = FixedBitSet::with_capacity(n + 1);
            a.grow(n + 1);
            a.set(cycle[1], true);
            return Ok((
                vec![ASplit::from_a_ntax_with_weight(a, n, d_ij)],
                CGSolveStats::default(),
            ));
        }
        return Ok((Vec::new(), CGSolveStats::default()));
    }

    let npairs = n * (n - 1) / 2;
    let d = setup_d(distances, cycle, n);

    let mut x = vec![0.0; npairs];
    let stats = run_active_conjugate(n, &d, &mut x);

    let mut splits = Vec::new();
    let mut index = 0usize;
    for i in 0..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);
            if x[index] > options.cutoff {
                splits.push(ASplit::from_a_ntax_with_weight(a.clone(), n, x[index]));
            }
            index += 1;
        }
    }
    Ok((splits, stats))
}

fn setup_d(distances: &Array2<f64>, cycle: &[usize], n: usize) -> Vec<f64> {
    let mut d = vec![0.0; n * (n - 1) / 2];
    let mut index = 0usize;
    for i in 0..n {
        for j in (i + 1)..n {
            d[index] = distances[[cycle[i + 1] - 1, cycle[j + 1] - 1]];
            index += 1;
        }
    }
    d
}

fn run_active_conjugate(n: usize, d: &[f64], x: &mut [f64]) -> CGSolveStats {
    let collapse_many_negs = true;
    let npairs = d.len();
    let band = BandIndex::new(n);
    let track_kernel_perf = npairs > 4096;
    let mut kernel_perf = KernelPerf::default();
    let started = Instant::now();
    let mut progress = SolveProgress::new(started);
    debug_assert_eq!(
        x.len(),
        npairs,
        "Vectors d,x have different dimensions: d.len()={}, x.len()={}",
        npairs,
        x.len()
    );

    // Unconstrained LS in row-major (uses row-major recurrence)
    run_unconstrained_ls(n, d, x);
    if x.iter().all(|&val| val >= 0.0) {
        return CGSolveStats {
            elapsed_secs: started.elapsed().as_secs_f64(),
            ..Default::default()
        };
    }

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

    let mut buf = vec![0.0; npairs];
    let mut row_sums = vec![0.0; n];
    let mut shifted = vec![0.0; n];
    let mut atwd = vec![0.0; npairs];
    {
        let mut y_band = vec![0.0; npairs];
        row_to_band(d, &mut y_band, &band);
        calculate_atx_band(&y_band, &mut atwd, &band, &mut row_sums, &mut shifted);
    }

    // CG scratch vectors (all band-major)
    let mut r = vec![0.0; npairs];
    let mut wtmp = vec![0.0; npairs];
    let mut p = vec![0.0; npairs];
    let mut old_x = vec![1.0; npairs];
    let mut negative_values = Vec::new();
    let mut selected_negative_indices = Vec::new();

    let mut active = vec![false; npairs];
    let fixed_active = vec![false; npairs];
    let mut free_indices: Vec<usize> = (0..npairs).collect();

    let mut first_pass = true;
    loop {
        progress.outer_iter += 1;
        let tol_factor = 1e-8;

        loop {
            if !first_pass {
                free_indices = (0..npairs).filter(|&i| !active[i]).collect();
                circular_conjugate_grads_band_direct(
                    npairs,
                    &mut r,
                    &mut wtmp,
                    &mut p,
                    &atwd,
                    &active,
                    &mut x_band,
                    &band,
                    &mut buf,
                    &mut row_sums,
                    &mut shifted,
                    &mut progress,
                    &mut kernel_perf,
                    track_kernel_perf,
                    &free_indices,
                    tol_factor,
                );
            }
            first_pass = false;

            if collapse_many_negs
                && worst_indices_reuse(
                    &x_band,
                    0.6,
                    &mut negative_values,
                    &mut selected_negative_indices,
                )
            {
                progress.collapse_events += 1;
                progress.collapsed_indices_total += selected_negative_indices.len();
                for &idx in &selected_negative_indices {
                    x_band[idx] = 0.0;
                    active[idx] = true;
                }
                free_indices = (0..npairs).filter(|&i| !active[i]).collect();
                circular_conjugate_grads_band_direct(
                    npairs,
                    &mut r,
                    &mut wtmp,
                    &mut p,
                    &atwd,
                    &active,
                    &mut x_band,
                    &band,
                    &mut buf,
                    &mut row_sums,
                    &mut shifted,
                    &mut progress,
                    &mut kernel_perf,
                    track_kernel_perf,
                    &free_indices,
                    tol_factor,
                );
            }

            let mut min_i: Option<usize> = None;
            let mut min_xi = -1.0_f64;
            for &i in &free_indices {
                if x_band[i] < 0.0 {
                    let xi = old_x[i] / (old_x[i] - x_band[i]);
                    if min_i.is_none() || xi < min_xi {
                        min_i = Some(i);
                        min_xi = xi;
                    }
                }
            }

            match min_i {
                None => break,
                Some(min_i) => {
                    for &i in &free_indices {
                        old_x[i] += min_xi * (x_band[i] - old_x[i]);
                    }
                    active[min_i] = true;
                    x_band[min_i] = 0.0;
                    progress.inner_loop_activations += 1;
                }
            }
        }

        // Gradient check: r = A^T W A x (direct band-major)
        if track_kernel_perf {
            let t0 = Instant::now();
            calculate_ab_band(&x_band, &mut buf, &band, &mut row_sums, &mut shifted);
            kernel_perf.calc_ab_calls += 1;
            kernel_perf.calc_ab_time += t0.elapsed();
        } else {
            calculate_ab_band(&x_band, &mut buf, &band, &mut row_sums, &mut shifted);
        }
        if track_kernel_perf {
            let t0 = Instant::now();
            calculate_atx_band(&buf, &mut r, &band, &mut row_sums, &mut shifted);
            kernel_perf.calc_atx_calls += 1;
            kernel_perf.calc_atx_time += t0.elapsed();
        } else {
            calculate_atx_band(&buf, &mut r, &band, &mut row_sums, &mut shifted);
        }

        if track_kernel_perf && Instant::now() >= progress.next_log {
            let active_count = active.iter().filter(|&&f| f).count();
            info!(
                "CG inference progress: outer_iter={} elapsed={:.1}s cg_calls={} cg_iters_total={} max_cg_single={} active={} inner_activations={} stagnation_exits={} collapse_events={} collapsed_total={} ab_calls={} atx_calls={} ab_avg_us={:.1} atx_avg_us={:.1}",
                progress.outer_iter,
                progress.started.elapsed().as_secs_f64(),
                progress.cg_calls,
                progress.cg_iters_total,
                progress.max_cg_iters_single_call,
                active_count,
                progress.inner_loop_activations,
                progress.stagnation_exits,
                progress.collapse_events,
                progress.collapsed_indices_total,
                kernel_perf.calc_ab_calls,
                kernel_perf.calc_atx_calls,
                avg_time_us(kernel_perf.calc_ab_time, kernel_perf.calc_ab_calls),
                avg_time_us(kernel_perf.calc_atx_time, kernel_perf.calc_atx_calls),
            );
            progress.next_log = Instant::now() + Duration::from_secs(5);
        }

        let mut min_i: Option<usize> = None;
        let mut min_grad = 1.0_f64;
        for i in 0..npairs {
            r[i] -= atwd[i];
            if active[i] && !fixed_active[i] {
                let grad_ij = r[i];
                if min_i.is_none() || grad_ij < min_grad {
                    min_i = Some(i);
                    min_grad = grad_ij;
                }
            }
        }

        let release = if min_grad > -0.0001 { None } else { min_i };
        if release.is_none() {
            if track_kernel_perf {
                info!(
                    "CG inference complete: outer_iters={} elapsed={:.1}s cg_calls={} cg_iters_total={} max_cg_single={} inner_activations={} stagnation_exits={} collapse_events={} collapsed_total={}",
                    progress.outer_iter,
                    progress.started.elapsed().as_secs_f64(),
                    progress.cg_calls,
                    progress.cg_iters_total,
                    progress.max_cg_iters_single_call,
                    progress.inner_loop_activations,
                    progress.stagnation_exits,
                    progress.collapse_events,
                    progress.collapsed_indices_total,
                );
            }
            // Convert x_band back to row-major
            band_to_row(&x_band, x, &band);
            return CGSolveStats {
                outer_iterations: progress.outer_iter,
                cg_calls: progress.cg_calls,
                cg_iters_total: progress.cg_iters_total,
                max_cg_iters_single_call: progress.max_cg_iters_single_call,
                stagnation_exits: progress.stagnation_exits,
                inner_loop_activations: progress.inner_loop_activations,
                collapse_events: progress.collapse_events,
                collapsed_indices_total: progress.collapsed_indices_total,
                kernel_ab_calls: kernel_perf.calc_ab_calls,
                kernel_atx_calls: kernel_perf.calc_atx_calls,
                elapsed_secs: progress.started.elapsed().as_secs_f64(),
            };
        } else if let Some(idx) = release {
            active[idx] = false;
        }
    }
}

fn worst_indices_reuse(
    x: &[f64],
    prop_kept: f64,
    negative_values: &mut Vec<f64>,
    result: &mut Vec<usize>,
) -> bool {
    if prop_kept == 0.0 {
        return false;
    }

    negative_values.clear();
    for &val in x {
        if val < 0.0 {
            negative_values.push(val);
        }
    }
    let num_neg = negative_values.len();
    if num_neg == 0 {
        result.clear();
        return false;
    }

    let nkept = ((prop_kept * num_neg as f64).ceil() as usize).max(1);
    let cutoff = {
        let kth = nkept - 1;
        let (_, pivot, _) = negative_values.select_nth_unstable_by(kth, |a, b| a.total_cmp(b));
        *pivot
    };

    result.clear();
    result.resize(nkept, 0usize);
    let mut front = 0usize;
    let mut back = nkept.saturating_sub(1);
    for (i, &val) in x.iter().enumerate() {
        if val < cutoff {
            result[front] = i;
            front += 1;
        } else if val == cutoff && back >= front {
            result[back] = i;
            if back == 0 {
                break;
            }
            back -= 1;
        }
    }
    true
}

fn run_unconstrained_ls(n: usize, d: &[f64], x: &mut [f64]) {
    let mut index = 0usize;
    for i in 0..=n - 3 {
        x[index] = (d[index] + d[index + (n - i - 2) + 1] - d[index + 1]) / 2.0;
        index += 1;
        for _j in (i + 2)..=n - 2 {
            x[index] =
                (d[index] + d[index + (n - i - 2) + 1] - d[index + 1] - d[index + (n - i - 2)])
                    / 2.0;
            index += 1;
        }
        if i == 0 {
            x[index] = (d[0] + d[n - 2] - d[2 * n - 4]) / 2.0;
        } else {
            x[index] = (d[index] + d[i] - d[i - 1] - d[index + (n - i - 2)]) / 2.0;
        }
        index += 1;
    }
    x[index] = (d[index] + d[n - 2] - d[n - 3]) / 2.0;
}

fn avg_time_us(total: Duration, calls: usize) -> f64 {
    avg_band_time_us(total, calls)
}

/// CG solver operating entirely in band-major layout; all vectors and
/// free_indices use band-major indexing.
#[allow(clippy::too_many_arguments)] // threads shared solver state through a hot path
fn circular_conjugate_grads_band_direct(
    npairs: usize,
    r: &mut [f64],
    w: &mut [f64],
    p: &mut [f64],
    b: &[f64],
    active: &[bool],
    x: &mut [f64],
    band: &BandIndex,
    buf: &mut [f64],
    row_sums: &mut [f64],
    shifted: &mut [f64],
    progress: &mut SolveProgress,
    kernel_perf: &mut KernelPerf,
    track_kernel_perf: bool,
    free_indices: &[usize],
    tol_factor: f64,
) {
    let kmax = (2 * free_indices.len()).min(band.npairs);
    progress.cg_calls += 1;
    let cg_call = progress.cg_calls;

    // r = A^T * A * x (direct band-major: AB then ATX)
    if track_kernel_perf {
        let t0 = Instant::now();
        calculate_ab_band(x, buf, band, row_sums, shifted);
        kernel_perf.calc_ab_calls += 1;
        kernel_perf.calc_ab_time += t0.elapsed();
        let t0 = Instant::now();
        calculate_atx_band(buf, r, band, row_sums, shifted);
        kernel_perf.calc_atx_calls += 1;
        kernel_perf.calc_atx_time += t0.elapsed();
    } else {
        calculate_ab_band(x, buf, band, row_sums, shifted);
        calculate_atx_band(buf, r, band, row_sums, shifted);
    }

    for k in 0..npairs {
        if !active[k] {
            r[k] = b[k] - r[k];
        } else {
            r[k] = 0.0;
        }
    }

    let mut rho: f64 = r.iter().map(|v| v * v).sum();
    let mut r_norm_sq: f64 = rho;
    let mut rho_old = 0.0;

    let b_norm_sq: f64 = b.iter().map(|v| v * v).sum();
    let e0_sq = tol_factor * b_norm_sq;
    let mut k = 0usize;
    let mut checkpoint_norm = r_norm_sq;

    while r_norm_sq > e0_sq && k < kmax {
        k += 1;
        if k == 1 {
            p.copy_from_slice(r);
        } else {
            let beta = rho / rho_old;
            for &i in free_indices {
                p[i] = r[i] + beta * p[i];
            }
        }

        // w = A^T * A * p (direct band-major: AB then ATX)
        if track_kernel_perf {
            let t0 = Instant::now();
            calculate_ab_band(p, buf, band, row_sums, shifted);
            kernel_perf.calc_ab_calls += 1;
            kernel_perf.calc_ab_time += t0.elapsed();
            let t0 = Instant::now();
            calculate_atx_band(buf, w, band, row_sums, shifted);
            kernel_perf.calc_atx_calls += 1;
            kernel_perf.calc_atx_time += t0.elapsed();
        } else {
            calculate_ab_band(p, buf, band, row_sums, shifted);
            calculate_atx_band(buf, w, band, row_sums, shifted);
        }

        let mut alpha_den = 0.0;
        for &i in free_indices {
            alpha_den += p[i] * w[i];
        }
        if alpha_den <= 0.0 || !alpha_den.is_finite() {
            break;
        }
        let alpha = rho / alpha_den;

        let mut rho_new = 0.0;
        for &i in free_indices {
            x[i] += alpha * p[i];
            r[i] -= alpha * w[i];
            rho_new += r[i] * r[i];
        }
        rho_old = rho;
        rho = rho_new;
        r_norm_sq = rho_new;

        // Stagnation detection
        if k > 0 && k.is_multiple_of(200) {
            if r_norm_sq > 0.5 * checkpoint_norm {
                if track_kernel_perf {
                    progress.stagnation_exits += 1;
                }
                break;
            }
            checkpoint_norm = r_norm_sq;
        }

        if track_kernel_perf && (k & 255) == 0 {
            let now = Instant::now();
            if now >= progress.next_inner_log {
                info!(
                    "CG solver progress: outer_iter={} cg_call={} local_iter={}/{} global_cg_iters~{} elapsed={:.1}s r_norm_sq={:.6e} threshold={:.6e}",
                    progress.outer_iter,
                    cg_call,
                    k,
                    kmax,
                    progress.cg_iters_total + k,
                    progress.started.elapsed().as_secs_f64(),
                    r_norm_sq,
                    e0_sq
                );
                progress.next_inner_log = now + Duration::from_secs(10);
            }
        }
    }
    progress.cg_iters_total += k;
    if k > progress.max_cg_iters_single_call {
        progress.max_cg_iters_single_call = k;
    }
}