semiflow 0.15.0-beta

//! Batched multi-channel evolve for graph heat kernels (ADR-0184, Issue #10).
//! `[C, N]` channel-major layout; 0-ULP vs serial (D5); hoisted GL₄/GL₆/VC samples.
#![allow(clippy::missing_errors_doc)]

#[cfg(feature = "parallel")]
mod parallel;

mod helpers;
use helpers::{
    accumulate_grad_channel, evolve_channel, evolve_magnus6_channel, evolve_magnus_channel,
    evolve_vc_channel, hoist_vc_gl4_samples, run_presampled_adj_buf,
    run_presampled_varcoef_adj_buf, validate_adj_grad_layout, validate_batched_layout,
};

use alloc::sync::Arc;

use crate::{
    chernoff::ChernoffFunction,
    error::SemiflowError,
    float::{from_f64, SemiflowFloat},
    graph::Graph,
    graph_adjoint_presampled::{PreSampledMagnusAdj, PreSampledVarCoefAdj},
    graph_sensitivity::GeneratorSensitivity,
    graph_signal::GraphSignal,
    magnus6_graph::{MagnusGraphHeat6thChernoff, GL6_C1, GL6_C2, GL6_C3},
    magnus_graph::{MagnusGraphHeatChernoff, GL4_C1_F64, GL4_C2_F64},
    scratch::ScratchPool,
    varcoef_magnus_graph::VarCoefMagnusGraphHeatChernoff,
};

// ── evolve_batched — generic ChernoffFunction<F, S = GraphSignal> ────────────

/// Evolve `n_cols` channels. `n_steps == 0` → copies src. Errors: layout violation.
#[cfg(not(feature = "parallel"))]
#[allow(clippy::cast_precision_loss)]
pub fn evolve_batched<C, F>(
    func: &C,
    graph: &Arc<Graph<F>>,
    t_final: F,
    n_steps: usize,
    src_cols: &[F],
    dst_cols: &mut [F],
) -> Result<(), SemiflowError>
where
    C: ChernoffFunction<F, S = GraphSignal<F>>,
    F: SemiflowFloat,
{
    let n = graph.n_nodes();
    let n_cols = validate_batched_layout::<F>(src_cols, dst_cols, n)?;
    if n_cols == 0 {
        return Ok(());
    }
    if n_steps == 0 {
        dst_cols.copy_from_slice(src_cols);
        return Ok(());
    }
    let tau = t_final / from_f64::<F>(n_steps as f64);
    let mut scratch = ScratchPool::<F>::new();
    let mut buf_a = GraphSignal::zeros(Arc::clone(graph));
    let mut buf_b = GraphSignal::zeros(Arc::clone(graph));
    for c in 0..n_cols {
        let src_c = &src_cols[c * n..(c + 1) * n];
        let dst_c = &mut dst_cols[c * n..(c + 1) * n];
        evolve_channel(
            func,
            tau,
            n_steps,
            src_c,
            dst_c,
            &mut buf_a,
            &mut buf_b,
            &mut scratch,
        )?;
    }
    Ok(())
}

/// Parallel build — `C: Sync` required (ADR-0184 D3). Errors: layout violation.
#[cfg(feature = "parallel")]
#[allow(clippy::cast_precision_loss)]
pub fn evolve_batched<C, F>(
    func: &C,
    graph: &Arc<Graph<F>>,
    t_final: F,
    n_steps: usize,
    src_cols: &[F],
    dst_cols: &mut [F],
) -> Result<(), SemiflowError>
where
    C: ChernoffFunction<F, S = GraphSignal<F>> + Sync,
    F: SemiflowFloat,
{
    let n = graph.n_nodes();
    let n_cols = validate_batched_layout::<F>(src_cols, dst_cols, n)?;
    if n_cols == 0 {
        return Ok(());
    }
    if n_steps == 0 {
        dst_cols.copy_from_slice(src_cols);
        return Ok(());
    }
    let tau = t_final / from_f64::<F>(n_steps as f64);
    if n_cols >= parallel::MIN_CHANNELS_PARALLEL {
        return parallel::par_evolve_batched(func, graph, tau, n_steps, src_cols, dst_cols, n);
    }
    // Serial fallback for single-channel calls.
    let mut scratch = ScratchPool::<F>::new();
    let mut buf_a = GraphSignal::zeros(Arc::clone(graph));
    let mut buf_b = GraphSignal::zeros(Arc::clone(graph));
    let src_c = &src_cols[0..n];
    let dst_c = &mut dst_cols[0..n];
    evolve_channel(
        func,
        tau,
        n_steps,
        src_c,
        dst_c,
        &mut buf_a,
        &mut buf_b,
        &mut scratch,
    )
}

// ── evolve_batched_magnus — MagnusGraphHeatChernoff K=4 (hoisted GL₄) ───────

/// Evolve `n_cols` channels with Magnus K=4, hoisting GL₄ samples once.
/// `n_steps == 0` → copies src. Errors: layout or convergence-radius.
#[allow(clippy::cast_precision_loss)]
#[allow(clippy::too_many_lines)]
pub fn evolve_batched_magnus<F: SemiflowFloat>(
    mc: &MagnusGraphHeatChernoff<F>,
    t_final: F,
    n_steps: usize,
    src_cols: &[F],
    dst_cols: &mut [F],
) -> Result<(), SemiflowError> {
    let n = mc.graph.n_nodes();
    let n_cols = validate_batched_layout::<F>(src_cols, dst_cols, n)?;
    if n_cols == 0 {
        return Ok(());
    }
    if n_steps == 0 {
        dst_cols.copy_from_slice(src_cols);
        return Ok(());
    }
    let tau = t_final / from_f64::<F>(n_steps as f64);
    crate::magnus_graph_helpers::validate_tau(tau)?;
    if mc.convergence_radius_check {
        crate::magnus_graph_helpers::validate_magnus_radius(mc.rho_bar_max, tau)?;
    }
    // Hoist: t_start = 0 every step → same l1, l2 every step.
    let c1 = from_f64::<F>(GL4_C1_F64);
    let c2 = from_f64::<F>(GL4_C2_F64);
    let l1 = mc.laplacian_at(c1 * tau);
    let l2 = mc.laplacian_at(c2 * tau);
    // Parallel dispatch (ADR-0184 D3): Laplacians are shared read-only; no
    // extra Sync bound needed (Laplacian<F>: Sync by structure).
    #[cfg(feature = "parallel")]
    if n_cols >= parallel::MIN_CHANNELS_PARALLEL {
        parallel::par_evolve_magnus(&l1, &l2, tau, n_steps, src_cols, dst_cols, n, &mc.graph);
        return Ok(());
    }
    let mut scratch = ScratchPool::<F>::new();
    let mut buf_a = GraphSignal::zeros(Arc::clone(&mc.graph));
    let mut buf_b = GraphSignal::zeros(Arc::clone(&mc.graph));
    for c in 0..n_cols {
        let src_c = &src_cols[c * n..(c + 1) * n];
        let dst_c = &mut dst_cols[c * n..(c + 1) * n];
        evolve_magnus_channel(
            &l1,
            &l2,
            tau,
            n_steps,
            src_c,
            dst_c,
            &mut buf_a,
            &mut buf_b,
            &mut scratch,
        );
    }
    Ok(())
}

// ── evolve_batched_magnus6 — MagnusGraphHeat6thChernoff K=6 (f64, hoisted) ──

/// Evolve `n_cols` channels with Magnus K=6 (f64 only, ADR-0056), hoisting GL₆ samples.
/// `n_steps == 0` → copies src. Errors: layout or convergence-radius.
#[allow(clippy::cast_precision_loss)]
#[allow(clippy::too_many_lines)]
pub fn evolve_batched_magnus6(
    mc: &MagnusGraphHeat6thChernoff<f64>,
    t_final: f64,
    n_steps: usize,
    src_cols: &[f64],
    dst_cols: &mut [f64],
) -> Result<(), SemiflowError> {
    let n = mc.graph().n_nodes();
    let n_cols = validate_batched_layout::<f64>(src_cols, dst_cols, n)?;
    if n_cols == 0 {
        return Ok(());
    }
    if n_steps == 0 {
        dst_cols.copy_from_slice(src_cols);
        return Ok(());
    }
    let tau = t_final / n_steps as f64;
    crate::magnus_graph_helpers::validate_tau(tau)?;
    if mc.convergence_radius_check {
        crate::magnus_graph_helpers::validate_magnus_radius(mc.rho_bar_max, tau)?;
    }
    // Hoist GL₆ samples.
    let l1 = mc.laplacian_at(GL6_C1 * tau);
    let l2 = mc.laplacian_at(GL6_C2 * tau);
    let l3 = mc.laplacian_at(GL6_C3 * tau);
    let graph_arc = Arc::new(Graph::<f64>::path(n));
    // Parallel dispatch (ADR-0184 D3): pre-hoisted Laplacians shared read-only.
    #[cfg(feature = "parallel")]
    if n_cols >= parallel::MIN_CHANNELS_PARALLEL {
        parallel::par_evolve_magnus6(
            &l1, &l2, &l3, tau, n_steps, src_cols, dst_cols, n, &graph_arc,
        );
        return Ok(());
    }
    let mut scratch = ScratchPool::<f64>::new();
    let mut buf_a = GraphSignal::zeros(Arc::clone(&graph_arc));
    let mut buf_b = GraphSignal::zeros(Arc::clone(&graph_arc));
    for c in 0..n_cols {
        let src_c = &src_cols[c * n..(c + 1) * n];
        let dst_c = &mut dst_cols[c * n..(c + 1) * n];
        evolve_magnus6_channel(
            &l1,
            &l2,
            &l3,
            tau,
            n_steps,
            src_c,
            dst_c,
            &mut buf_a,
            &mut buf_b,
            &mut scratch,
        );
    }
    Ok(())
}

// ── evolve_batched_varcoef_magnus — VarCoef Magnus K=4 (hoisted GL₄ + a) ───

/// Evolve `n_cols` channels with variable-coefficient Magnus K=4.
/// Hoists GL₄ + diffusion-coefficient samples. Errors: layout/a-weights/radius.
#[allow(clippy::cast_precision_loss)]
#[allow(clippy::too_many_lines)]
pub fn evolve_batched_varcoef_magnus<F: SemiflowFloat>(
    mc: &VarCoefMagnusGraphHeatChernoff<F>,
    t_final: F,
    n_steps: usize,
    src_cols: &[F],
    dst_cols: &mut [F],
) -> Result<(), SemiflowError> {
    let n = mc.n_nodes();
    let n_cols = validate_batched_layout::<F>(src_cols, dst_cols, n)?;
    if n_cols == 0 {
        return Ok(());
    }
    if n_steps == 0 {
        dst_cols.copy_from_slice(src_cols);
        return Ok(());
    }
    let tau = t_final / from_f64::<F>(n_steps as f64);
    crate::varcoef_magnus_graph::validate_tau(tau)?;
    if mc.convergence_radius_check {
        crate::varcoef_magnus_graph::validate_magnus_radius(mc.rho_bar_max, mc.a_sup_max, tau)?;
    }
    // Hoist GL₄ + a samples once (t_start = 0 every step → identical each step).
    let samples = hoist_vc_gl4_samples(mc, tau, n)?;
    // Dummy graph (only n_nodes matters for GraphSignal allocation).
    let dummy_g = Arc::new(Graph::<F>::path(n));
    // Parallel dispatch (ADR-0184 D3): hoisted Laplacians + a-slices shared read-only.
    #[cfg(feature = "parallel")]
    if n_cols >= parallel::MIN_CHANNELS_PARALLEL {
        parallel::par_evolve_vc(
            &samples.l1,
            &samples.sqrt_a1,
            &samples.l2,
            &samples.sqrt_a2,
            tau,
            n_steps,
            src_cols,
            dst_cols,
            n,
            &dummy_g,
        );
        return Ok(());
    }
    let mut buf_a = GraphSignal::zeros(Arc::clone(&dummy_g));
    let mut buf_b = GraphSignal::zeros(Arc::clone(&dummy_g));
    let mut scratch = ScratchPool::<F>::new();
    for c in 0..n_cols {
        let src_c = &src_cols[c * n..(c + 1) * n];
        let dst_c = &mut dst_cols[c * n..(c + 1) * n];
        evolve_vc_channel(
            &samples.l1,
            &samples.sqrt_a1,
            &samples.l2,
            &samples.sqrt_a2,
            tau,
            n_steps,
            src_c,
            dst_c,
            &mut buf_a,
            &mut buf_b,
            &mut scratch,
        );
    }
    Ok(())
}

// ── PreSampledMagnusAdj — batched state adjoint ───────────────────────────────

impl<F: SemiflowFloat> PreSampledMagnusAdj<F> {
    /// Backward costate sweep for `n_cols` channels (`[C, N]` flat).
    /// 0-ULP vs `C` serial calls. Errors: `n_steps` mismatch, layout, convergence.
    pub fn evolve_state_adjoint_batched_into(
        &self,
        tau: F,
        n_steps: usize,
        src_cols: &[F],
        dst_cols: &mut [F],
        scratch: &mut ScratchPool<F>,
    ) -> Result<(), SemiflowError> {
        let n = self.n_nodes;
        let n_cols = validate_batched_layout::<F>(src_cols, dst_cols, n)?;
        if n_cols == 0 {
            return Ok(());
        }
        // Delegate to single-channel method per channel (0-ULP by construction).
        // Build a dummy graph once for GraphSignal construction.
        let dummy_g: Arc<Graph<F>> = Arc::new(Graph::path(n));
        for c in 0..n_cols {
            let src_c = &src_cols[c * n..(c + 1) * n];
            let dst_c = &mut dst_cols[c * n..(c + 1) * n];
            run_presampled_adj_buf(tau, n_steps, &self.seq, &dummy_g, src_c, dst_c, scratch)?;
        }
        Ok(())
    }
}

// ── PreSampledVarCoefAdj — batched state adjoint ─────────────────────────────

impl<F: SemiflowFloat> PreSampledVarCoefAdj<F> {
    /// Backward costate sweep for `n_cols` channels (`VarCoef`). 0-ULP vs `C` serial calls.
    /// Errors: layout mismatch or `n_steps != self.n_steps()`.
    pub fn evolve_state_adjoint_batched_into(
        &self,
        tau: F,
        n_steps: usize,
        src_cols: &[F],
        dst_cols: &mut [F],
        scratch: &mut ScratchPool<F>,
    ) -> Result<(), SemiflowError> {
        let n = self.n_nodes;
        let n_cols = validate_batched_layout::<F>(src_cols, dst_cols, n)?;
        if n_cols == 0 {
            return Ok(());
        }
        let dummy_g: Arc<Graph<F>> = Arc::new(Graph::path(n));
        for c in 0..n_cols {
            let src_c = &src_cols[c * n..(c + 1) * n];
            let dst_c = &mut dst_cols[c * n..(c + 1) * n];
            run_presampled_varcoef_adj_buf(
                tau,
                n_steps,
                &self.seq,
                &self.a_seq,
                n,
                &dummy_g,
                src_c,
                dst_c,
                scratch,
            )?;
        }
        Ok(())
    }
}

// ── adjoint_state_gradient_batched — summed edge-weight gradient ─────────────

/// Summed gradient `∂J/∂θ = Σ_c ∂J_c/∂θ` (ADR-0184 D4). `grad_theta` zeroed then
/// accumulated. 0-ULP vs `C` serial calls. Errors: layout or param count mismatch.
#[cfg(not(feature = "parallel"))]
#[allow(clippy::too_many_arguments)]
pub fn adjoint_state_gradient_batched<F, P>(
    mc: &MagnusGraphHeatChernoff<F>,
    u0_cols: &[F],
    dj_cols: &[F],
    n_steps: usize,
    tau: F,
    param_deriv: &P,
    grad_theta: &mut [F],
    scratch: &mut ScratchPool<F>,
) -> Result<(), SemiflowError>
where
    F: SemiflowFloat,
    P: GeneratorSensitivity<F>,
{
    let n = mc.graph.n_nodes();
    let n_cols = validate_adj_grad_layout(
        u0_cols,
        dj_cols,
        n,
        grad_theta.len(),
        param_deriv.n_params(),
    )?;
    grad_theta.fill(F::zero());
    if n_steps == 0 || n_cols == 0 {
        return Ok(());
    }
    let graph_arc = Arc::clone(&mc.graph);
    let n_params = param_deriv.n_params();
    for c in 0..n_cols {
        let u0_c = &u0_cols[c * n..(c + 1) * n];
        let dj_c = &dj_cols[c * n..(c + 1) * n];
        accumulate_grad_channel(
            mc,
            u0_c,
            dj_c,
            &graph_arc,
            n_steps,
            tau,
            param_deriv,
            n_params,
            grad_theta,
            scratch,
        )?;
    }
    Ok(())
}

/// Parallel build — `P: Sync` required (ADR-0184 D3/D4). `_scratch` unused (API compat).
/// Errors: layout or param count mismatch.
#[cfg(feature = "parallel")]
#[allow(clippy::too_many_arguments)]
#[allow(clippy::too_many_lines)]
pub fn adjoint_state_gradient_batched<F, P>(
    mc: &MagnusGraphHeatChernoff<F>,
    u0_cols: &[F],
    dj_cols: &[F],
    n_steps: usize,
    tau: F,
    param_deriv: &P,
    grad_theta: &mut [F],
    _scratch: &mut ScratchPool<F>,
) -> Result<(), SemiflowError>
where
    F: SemiflowFloat,
    P: GeneratorSensitivity<F> + Sync,
{
    let n = mc.graph.n_nodes();
    let n_cols = validate_adj_grad_layout(
        u0_cols,
        dj_cols,
        n,
        grad_theta.len(),
        param_deriv.n_params(),
    )?;
    grad_theta.fill(F::zero());
    if n_steps == 0 || n_cols == 0 {
        return Ok(());
    }
    if n_cols >= parallel::MIN_CHANNELS_PARALLEL {
        return parallel::par_grad_batched(
            mc,
            u0_cols,
            dj_cols,
            n_steps,
            tau,
            param_deriv,
            grad_theta,
            n_cols,
            n,
        );
    }
    // Serial fallback for single-channel gradient.
    let graph_arc = Arc::clone(&mc.graph);
    let n_params = param_deriv.n_params();
    let mut scr = ScratchPool::<F>::new();
    let u0_c = &u0_cols[0..n];
    let dj_c = &dj_cols[0..n];
    accumulate_grad_channel(
        mc,
        u0_c,
        dj_c,
        &graph_arc,
        n_steps,
        tau,
        param_deriv,
        n_params,
        grad_theta,
        &mut scr,
    )
}