cobre-sddp 0.3.2

//! Backward pass execution for the SDDP training loop.
//!
//! [`run_backward_pass`] sweeps stages in reverse order (`T-2` down to `0`),
//! evaluating the cost-to-go at each trial point **assigned to this rank**
//! during the forward pass. For each trial point, the backward pass iterates
//! over every opening from the fixed opening tree, extracts LP duals to form
//! Benders cut coefficients, and aggregates per-opening outcomes via
//! [`RiskMeasure::aggregate_cut`] to produce one cut per trial point per
//! stage. Each aggregated cut is inserted into the [`FutureCostFunction`].
//!
//! Although [`ExchangeBuffers`] contains trial points from all ranks (after
//! `allgatherv`), each rank only processes its own forward pass assignments
//! to avoid generating duplicate cuts. Cut synchronization (`allgatherv`)
//! distributes the generated cuts to all ranks after the backward pass.
//!
//! ## Stage indexing convention
//!
//! The backward pass generates a cut **at stage `t`** by solving the LP
//! **at stage `t + 1`** (the successor) under each opening noise vector from
//! that successor stage. The opening tree provides noise at `t + 1`.
//!
//! ## Cut coefficient formula
//!
//! For a solve at stage `t + 1` with trial point state `x_hat`:
//!
//! ```text
//! pi[i]  = dual[i] * row_scale[i]           for i in 0..n_state
//! alpha  = Q - sum_i(pi[i] * x_hat[i])      (intercept)
//! ```
//!
//! where `Q` is the LP objective and `dual[0..n_state]` are the duals of the
//! fixing constraints (storage-fixing and lag-fixing rows).
//!
//! The coefficients stored in the [`FutureCostFunction`] are the raw (unscaled)
//! duals of the state-fixing rows. Negation is applied later when building the
//! LP cut row in `build_cut_row_batch_into` (forward.rs):
//! `-coeff * x + theta >= intercept`.
//! See the project convention: "coefficients = dual (NOT -dual)".
//!
//! ## Cut activity tracking
//!
//! After each backward solve, the duals of the appended cut rows are inspected
//! to determine which existing cuts at the successor stage are binding. The
//! metadata of binding cuts is updated in-place so that cut selection
//! strategies have accurate activity counts at the end of the iteration.
//!
//! ## Thread-level parallelism
//!
//! Within a rank, the outer per-stage loop remains sequential (stage `t`
//! depends on cuts generated at stage `t+1`). The inner trial-point loop is
//! parallelised across [`SolverWorkspace`] instances using rayon's
//! `par_iter_mut`. Trial points are statically partitioned across workspaces
//! (not rayon default work-stealing) to ensure deterministic assignment.
//!
//! Each worker thread generates cuts into a thread-local `StagedCut` buffer
//! rather than directly into the FCF. After the parallel region completes for
//! a stage, the staged cuts are sorted by `trial_point_idx` and inserted into
//! the FCF in that deterministic order. This ensures bit-for-bit identical
//! results regardless of thread count or thread completion order.
//!
//! ## Hot-path allocation discipline
//!
//! Allocations are limited to:
//! - One `Vec<f64>` for opening probabilities per stage (outside the trial
//!   point loop).
//! - One `Vec<BackwardOutcome>` per worker thread, allocated once per stage
//!   in the parallel region and reused via `clear()` per trial point.
//! - One `RowBatch` per stage built by `build_cut_row_batch` (outside the
//!   trial point loop, before the parallel region).
//! - One `Vec<StagedCut>` per stage for the merge phase (bounded by
//!   `local_work` entries, each holding one cut and its binding slot list).
//!
//! The `binding_slots` vector inside each `StagedCut` is allocated per
//! trial point — a flat buffer optimization is deferred to profiling.

use std::time::Instant;

use cobre_comm::Communicator;
use cobre_solver::{RowBatch, SolverError, SolverInterface};
use rayon::iter::{IndexedParallelIterator, IntoParallelRefMutIterator, ParallelIterator};

use crate::{
    FutureCostFunction, SddpError, TrajectoryRecord,
    context::{StageContext, TrainingContext},
    cut_sync::CutSyncBuffers,
    forward::{build_cut_row_batch_into, partition},
    noise::{transform_inflow_noise, transform_load_noise, transform_ncs_noise},
    risk_measure::BackwardOutcome,
    risk_measure::RiskMeasure,
    solver_stats::{SolverStatsDelta, aggregate_solver_statistics},
    state_exchange::ExchangeBuffers,
    workspace::{BasisStore, SolverWorkspace},
};

/// Result produced by the backward pass on a single rank.
#[derive(Debug, Clone)]
#[must_use]
pub struct BackwardResult {
    /// Total number of cuts generated by this rank during the backward pass.
    pub cuts_generated: usize,

    /// Wall-clock time in milliseconds for this rank's backward pass.
    pub elapsed_ms: u64,

    /// Number of LP solves performed during this backward pass.
    pub lp_solves: u64,

    /// Per-stage solver statistics deltas.
    ///
    /// Each entry is `(successor_stage_index, delta)` where the stage index
    /// identifies which stage's LP was solved (the successor), not the stage
    /// where cuts are added. Entries are in reverse stage order (matching
    /// the backward iteration direction).
    pub stage_stats: Vec<(usize, SolverStatsDelta)>,

    /// Wall-clock time for state exchange (`allgatherv`) accumulated across
    /// all stages, in milliseconds.
    pub state_exchange_time_ms: u64,

    /// Wall-clock time for `build_cut_row_batch_into` accumulated across
    /// all stages, in milliseconds.
    pub cut_batch_build_time_ms: u64,

    /// Estimated rayon barrier + scheduling overhead accumulated across all
    /// stages, in milliseconds. Computed per-stage as
    /// `process_stage_wall_ms - (solve_time_ms / n_workers)`.
    pub rayon_overhead_time_ms: u64,

    /// Wall-clock time for per-stage cut synchronization (`allgatherv`)
    /// accumulated across all stages, in milliseconds.
    pub cut_sync_time_ms: u64,
}

/// Per-thread staging buffer for one aggregated cut produced at a single trial
/// point during the parallel backward sweep.
///
/// Each worker thread populates one `StagedCut` per trial point instead of
/// writing directly into the [`FutureCostFunction`]. After the parallel region,
/// staged cuts are sorted by `trial_point_idx` and merged into the FCF in
/// deterministic order regardless of thread completion order.
struct StagedCut {
    /// Local trial-point index within `0..local_work`. Used for deterministic
    /// merge ordering after the parallel region.
    trial_point_idx: usize,

    /// Aggregated cut intercept (result of `RiskMeasure::aggregate_cut`).
    intercept: f64,

    /// Aggregated cut coefficients (length = `n_state`).
    coefficients: Vec<f64>,

    /// Global forward-pass index (`fwd_offset + m`), stored as `u32` for the
    /// FCF slot formula.
    forward_pass_index: u32,

    /// Slots in the successor pool that were binding for at least one opening
    /// during this trial point. These are accumulated across openings within
    /// the worker and applied to FCF metadata in the sequential merge phase.
    ///
    /// Each entry is `(slot_index, active_count_increment)` where
    /// `active_count_increment` is the number of openings in which that cut
    /// was binding. This allows the merge phase to correctly accumulate
    /// activity counts without double-counting.
    binding_increments: Vec<(usize, u64)>,
}

/// Inputs bundled from the training loop for the backward pass.
///
/// Groups the iteration-varying scalars and slices that would otherwise push
/// the argument count of [`run_backward_pass`] beyond seven.
pub struct BackwardPassSpec<'a> {
    /// Exchange buffers for gathering trial-point states via `allgatherv`.
    ///
    /// When `records` is non-empty, `run_backward_pass` calls
    /// `exchange.exchange(records, t, ...)` once per stage before processing
    /// trial points at that stage. When `records` is empty (test path), the
    /// caller is responsible for pre-populating the exchange buffers before
    /// calling `run_backward_pass`.
    pub exchange: &'a mut ExchangeBuffers,

    /// Forward-pass trajectory records used to populate `exchange` per stage.
    ///
    /// Length must be `local_work * num_stages` when non-empty. Pass `&[]` in
    /// tests that pre-populate `exchange` directly.
    pub records: &'a [TrajectoryRecord],

    /// Current training iteration index (1-based), used for cut metadata.
    pub iteration: u64,

    /// Number of trial points assigned to this rank for the backward pass.
    pub local_work: usize,

    /// Global offset for this rank's trial points (`rank * fwd_per_rank`).
    pub fwd_offset: usize,

    /// Per-stage risk measures (length = `num_stages`). `risk_measures[t]`
    /// is applied when aggregating cut outcomes at stage `t`.
    pub risk_measures: &'a [RiskMeasure],

    /// Minimum dual multiplier for a cut to count as binding (BUG-2 fix).
    ///
    /// Cuts with duals above this threshold are marked as binding during the
    /// backward pass activity tracking. A value of `0.0` preserves the
    /// original strict-positive behavior; a value like `1e-8` filters out
    /// numerical noise from near-zero positive duals.
    pub cut_activity_tolerance: f64,

    /// Pre-allocated cut synchronization buffers for per-stage `allgatherv`.
    ///
    /// After local cuts are inserted into the FCF at each stage, `sync_cuts`
    /// is called to distribute cuts across all ranks. For single-rank runs
    /// the `allgatherv` is a no-op and completes immediately.
    pub cut_sync_bufs: &'a mut CutSyncBuffers,

    /// Pre-allocated buffer for uniform opening probabilities. Reused per stage
    /// via `clear()` + `resize()` to avoid per-stage allocation.
    ///
    /// Passed as a mutable reference to avoid `mem::take` capacity loss when
    /// the training loop encounters an error before buffer recovery.
    pub probabilities_buf: &'a mut Vec<f64>,

    /// Pre-allocated buffer for successor active cut slot indices. Reused per
    /// stage via `clear()` + `extend()` to avoid per-stage allocation.
    ///
    /// Passed as a mutable reference (same rationale as `probabilities_buf`).
    pub successor_active_slots_buf: &'a mut Vec<usize>,

    /// Optional visited-states archive for dominated cut selection.
    ///
    /// When `Some`, gathered states are archived after each per-stage exchange
    /// so that the domination test has access to all forward-pass trial points.
    /// Pass `None` when using `Level1`, `Lml1`, or no cut selection.
    pub visited_archive: Option<&'a mut crate::visited_states::VisitedStatesArchive>,
}

/// Per-successor data bundled for `process_stage_backward` and the trial-point helper.
///
/// Groups the successor-specific arguments — including the stage index `t`,
/// opening probabilities, pre-built cut batch, cut activity metadata, and the
/// basis store — to keep per-function argument counts at or below seven.
struct SuccessorSpec<'a> {
    /// Stage index being cut (the stage whose cost-to-go we are computing).
    t: usize,
    /// Successor stage index (`t + 1`), where the LP is actually solved.
    successor: usize,
    /// This rank's MPI rank index (used to address exchange buffer state).
    my_rank: usize,
    /// Uniform opening probabilities for the successor stage.
    probabilities: &'a [f64],
    /// Pre-built cut rows to append to each successor LP.
    cut_batch: &'a RowBatch,
    /// Number of active cuts at the successor stage (= `cut_batch.num_rows`).
    num_cuts_at_successor: usize,
    /// Base row count of the successor template (excludes appended cuts).
    template_num_rows: usize,
    /// Ordered slot indices of the active cuts at the successor stage.
    successor_active_slots: &'a [usize],
    /// Basis store for warm-starting each worker's first opening solve.
    basis_store: &'a BasisStore,
    /// Minimum dual multiplier for a cut to count as binding.
    cut_activity_tolerance: f64,
    /// Populated count of the successor's cut pool. Used to size the
    /// `slot_increments` Vec for O(1) indexed binding tracking.
    successor_populated_count: usize,
}

/// Per-trial-point mutable accumulators reused across openings.
struct TrialAccumulators {
    /// Collected per-opening outcomes (intercept + coefficients + objective).
    outcomes: Vec<BackwardOutcome>,
    /// Binding count per cut pool slot. Indexed by slot index, pre-allocated
    /// to `pool.populated_count` and zeroed per trial point via `fill(0)`.
    /// Replaces the previous `HashMap<usize, u64>` to eliminate hashing
    /// overhead and per-stage allocation.
    slot_increments: Vec<u64>,
}

/// Load the stage LP template and append cuts once per trial point.
///
/// Called once before the opening loop in [`process_trial_point_backward`].
/// The LP structure (template + cuts) is identical across all openings for
/// the same trial point — only the noise-dependent bounds change.
///
/// # Future: incremental cut injection
///
/// A future incremental path could skip `load_model` on iterations 2+
/// and only append new cuts via `add_rows`. Implementing this requires
/// adapting the dual extraction logic in [`process_trial_point_backward`]
/// to use [`CutRowMap`] for LP row lookups instead of assuming sequential
/// row ordering (the current
/// `view.dual[template_num_rows..template_num_rows + num_cuts]` pattern).
///
/// [`CutRowMap`]: crate::cut::CutRowMap
fn load_backward_lp<S: SolverInterface + Send>(
    ws: &mut SolverWorkspace<S>,
    ctx: &StageContext<'_>,
    cut_batch: &RowBatch,
    s: usize,
) {
    ws.solver.load_model(&ctx.templates[s]);
    ws.solver.add_rows(cut_batch);
}

/// Transform opening noise and patch LP bounds for one backward opening.
///
/// Called once per opening inside [`process_trial_point_backward`].  The LP
/// structure is already loaded by [`load_backward_lp`]; this function only
/// updates noise-dependent row and column bounds via `set_row_bounds` /
/// `set_col_bounds`.
fn patch_opening_bounds<S: SolverInterface + Send>(
    ws: &mut SolverWorkspace<S>,
    ctx: &StageContext<'_>,
    training_ctx: &TrainingContext<'_>,
    raw_noise: &[f64],
    x_hat: &[f64],
    s: usize,
) {
    let n_blks = if ctx.n_load_buses > 0 {
        ctx.block_counts_per_stage[s]
    } else {
        0
    };
    transform_inflow_noise(raw_noise, s, x_hat, ctx, training_ctx, &mut ws.scratch);
    transform_load_noise(
        raw_noise,
        ctx.n_hydros,
        ctx.n_load_buses,
        training_ctx.stochastic,
        s,
        n_blks,
        &mut ws.scratch.load_rhs_buf,
    );
    let n_stochastic_ncs = training_ctx.stochastic.n_stochastic_ncs();
    if n_stochastic_ncs > 0 {
        transform_ncs_noise(
            raw_noise,
            ctx.n_hydros,
            ctx.n_load_buses,
            training_ctx.stochastic,
            s,
            ctx.block_counts_per_stage[s],
            ctx.ncs_max_gen,
            &mut ws.scratch.ncs_col_upper_buf,
        );
    }
    ws.patch_buf.fill_forward_patches(
        training_ctx.indexer,
        x_hat,
        &ws.scratch.noise_buf,
        ctx.base_rows[s],
        &ctx.templates[s].row_scale,
    );
    if ctx.n_load_buses > 0 {
        ws.patch_buf.fill_load_patches(
            ctx.load_balance_row_starts[s],
            n_blks,
            &ws.scratch.load_rhs_buf,
            ctx.load_bus_indices,
            &ctx.templates[s].row_scale,
        );
    }
    ws.patch_buf.fill_z_inflow_patches(
        training_ctx.indexer.z_inflow_row_start,
        &ws.scratch.z_inflow_rhs_buf,
        &ctx.templates[s].row_scale,
    );
    let pc = ws.patch_buf.forward_patch_count();
    ws.solver.set_row_bounds(
        &ws.patch_buf.indices[..pc],
        &ws.patch_buf.lower[..pc],
        &ws.patch_buf.upper[..pc],
    );
    if n_stochastic_ncs > 0 && !training_ctx.indexer.ncs_generation.is_empty() {
        let n_blks_stage = ctx.block_counts_per_stage[s];
        let expected_len = n_stochastic_ncs * n_blks_stage;
        if ws.scratch.ncs_col_indices_buf.len() != expected_len {
            ws.scratch.ncs_col_indices_buf.clear();
            ws.scratch.ncs_col_lower_buf.clear();
            for ncs_idx in 0..n_stochastic_ncs {
                for blk in 0..n_blks_stage {
                    ws.scratch.ncs_col_indices_buf.push(
                        training_ctx.indexer.ncs_generation.start + ncs_idx * n_blks_stage + blk,
                    );
                    ws.scratch.ncs_col_lower_buf.push(0.0);
                }
            }
        }
        ws.solver.set_col_bounds(
            &ws.scratch.ncs_col_indices_buf,
            &ws.scratch.ncs_col_lower_buf,
            &ws.scratch.ncs_col_upper_buf,
        );
    }
}

/// Process one trial point `m` in the backward pass, iterating over all openings.
///
/// Called once per trial point inside the parallel worker closure of
/// `process_stage_backward`. The caller pre-allocates and clears `accum`
/// before each call. Returns a single aggregated [`StagedCut`].
fn process_trial_point_backward<S: SolverInterface + Send>(
    ws: &mut SolverWorkspace<S>,
    ctx: &StageContext<'_>,
    training_ctx: &TrainingContext<'_>,
    spec: &BackwardPassSpec<'_>,
    succ: &SuccessorSpec<'_>,
    m: usize,
    accum: &mut TrialAccumulators,
) -> Result<StagedCut, SddpError> {
    let indexer = training_ctx.indexer;
    let tree_view = training_ctx.stochastic.tree_view();
    let x_hat = spec.exchange.state_at(succ.my_rank, m);
    let scenario = spec.fwd_offset + m;
    let s = succ.successor;

    for omega in 0..succ.probabilities.len() {
        let raw_noise = tree_view.opening(s, omega);
        patch_opening_bounds(ws, ctx, training_ctx, raw_noise, x_hat, s);

        // Opening 0: use forward-pass basis from BasisStore for warm start.
        // Openings 1+: HiGHS retains internal factorization from the previous
        // solve — only bounds/RHS changed via `patch_opening_bounds`, so
        // `solve()` hot-starts without redundant refactorization (P3b).
        let view = if omega == 0 {
            if let Some(rb) = succ.basis_store.get(m, s) {
                ws.solver.solve_with_basis(rb)
            } else {
                ws.solver.solve()
            }
        } else {
            ws.solver.solve()
        }
        .map_err(|e| match e {
            SolverError::Infeasible => SddpError::Infeasible {
                stage: succ.t,
                iteration: spec.iteration,
                scenario,
            },
            other => SddpError::Solver(other),
        })?;

        let objective = view.objective;
        // Unscale duals from scaled units back to original units.
        // `dual_original[i] = row_scale[i] * dual_scaled[i]`.
        // Only the state-fixing rows [0, n_state) are needed for cut coefficients.
        // Cut rows have implicit row_scale = 1.0 and require no adjustment.
        let row_scale = &ctx.templates[s].row_scale;
        let out = &mut accum.outcomes[omega];
        if row_scale.is_empty() {
            out.coefficients
                .copy_from_slice(&view.dual[..indexer.n_state]);
        } else {
            for (i, &d) in view.dual[..indexer.n_state].iter().enumerate() {
                out.coefficients[i] = d * row_scale[i];
            }
        }
        out.objective_value = objective;
        // Intercept: alpha = Q_scaled - pi' * x_hat.
        //
        // `objective` is the LP objective in scaled cost units (divided by
        // COST_SCALE_FACTOR). `coefficients` are unscaled duals (dual_i *
        // row_scale_i). The product `pi' * x_hat` is also in scaled cost
        // units because the LP duals inherit the cost scaling. The cut
        // `theta >= alpha + pi'(x - x_hat)` is evaluated in the parent
        // LP, where theta is also in scaled units, so all terms are
        // consistently scaled.
        out.intercept = objective
            - out
                .coefficients
                .iter()
                .zip(x_hat)
                .map(|(pi, x)| pi * x)
                .sum::<f64>();

        if succ.num_cuts_at_successor > 0 {
            let cut_duals = &view.dual
                [succ.template_num_rows..succ.template_num_rows + succ.num_cuts_at_successor];
            for (cut_idx, &slot) in succ.successor_active_slots.iter().enumerate() {
                if cut_duals[cut_idx] > succ.cut_activity_tolerance {
                    accum.slot_increments[slot] += 1;
                }
            }
        }
    }

    let (agg_intercept, agg_coefficients) =
        spec.risk_measures[succ.t].aggregate_cut(&accum.outcomes, succ.probabilities);
    debug_assert!(
        u32::try_from(scenario).is_ok(),
        "global scenario index overflows u32"
    );
    #[allow(clippy::cast_possible_truncation)]
    let forward_pass_index = scenario as u32;
    let binding_increments: Vec<(usize, u64)> = accum
        .slot_increments
        .iter()
        .enumerate()
        .filter(|&(_, c)| *c > 0)
        .map(|(s, c)| (s, *c))
        .collect();
    Ok(StagedCut {
        trial_point_idx: m,
        intercept: agg_intercept,
        coefficients: agg_coefficients,
        forward_pass_index,
        binding_increments,
    })
}

/// Evaluate all trial points for a single backward stage, returning staged cuts.
///
/// Runs the parallel trial-point loop over `0..spec.local_work` for the
/// successor of stage `t`. Workers are statically partitioned across `workspaces`.
/// Returns `Vec<StagedCut>` in completion order (unsorted); caller sorts by
/// `trial_point_idx` before inserting into the FCF.
fn process_stage_backward<S: SolverInterface + Send>(
    workspaces: &mut [SolverWorkspace<S>],
    ctx: &StageContext<'_>,
    training_ctx: &TrainingContext<'_>,
    spec: &BackwardPassSpec<'_>,
    succ: &SuccessorSpec<'_>,
) -> Vec<Result<Vec<StagedCut>, SddpError>> {
    let n_openings = succ.probabilities.len();
    let n_workers = workspaces.len().max(1);
    let local_work = spec.local_work;

    workspaces
        .par_iter_mut()
        .enumerate()
        .map(|(worker_id, ws)| {
            let (start_m, end_m) = partition(local_work, n_workers, worker_id);

            // Load the LP structure once per worker per stage; trial points
            // only patch bounds. Skip for workers with no assigned work.
            if start_m < end_m {
                load_backward_lp(ws, ctx, succ.cut_batch, succ.successor);
            }

            let mut staged: Vec<StagedCut> = Vec::with_capacity(end_m - start_m);
            let n_state = training_ctx.indexer.n_state;
            let mut accum = TrialAccumulators {
                outcomes: (0..n_openings)
                    .map(|_| BackwardOutcome {
                        intercept: 0.0,
                        coefficients: vec![0.0; n_state],
                        objective_value: 0.0,
                    })
                    .collect(),
                slot_increments: vec![0u64; succ.successor_populated_count],
            };

            for m in start_m..end_m {
                accum.slot_increments.fill(0);
                staged.push(process_trial_point_backward(
                    ws,
                    ctx,
                    training_ctx,
                    spec,
                    succ,
                    m,
                    &mut accum,
                )?);
            }
            Ok(staged)
        })
        .collect()
}

/// Execute the backward pass for one training iteration on this rank.
///
/// Sweeps stages from `num_stages - 2` down to `0` (the last stage `T-1` has
/// no successor stage and therefore produces no cuts). For each stage, the
/// trial-point loop is parallelised across the provided [`SolverWorkspace`]
/// instances using static work partitioning. Each worker generates cut data
/// into a thread-local `StagedCut` buffer. After the parallel region, staged
/// cuts are sorted by trial-point index and inserted into the FCF in
/// deterministic order.
///
/// The outer per-stage loop remains sequential: stage `t` depends on cuts
/// inserted at stage `t+1` in the previous iteration of the stage loop.
///
/// ## Error handling
///
/// On `SolverError::Infeasible`, returns `SddpError::Infeasible` with the
/// backward stage (0-based), iteration, and global trial point index.
/// On any other `SolverError`, returns `SddpError::Solver`. On error, the
/// FCF may be partially populated.
///
/// # Errors
///
/// Returns `Err(SddpError::Infeasible { .. })` when a stage LP has no
/// feasible solution during the backward sweep. Returns
/// `Err(SddpError::Solver(_))` for all other terminal LP solver failures.
///
/// # Panics (debug builds only)
///
/// Panics if any of the following debug preconditions are violated:
///
/// - `ctx.templates.len() != num_stages`
/// - `ctx.base_rows.len() != num_stages`
/// - `spec.risk_measures.len() != num_stages`
#[allow(clippy::too_many_arguments, clippy::too_many_lines)]
pub fn run_backward_pass<S: SolverInterface + Send, C: Communicator>(
    workspaces: &mut [SolverWorkspace<S>],
    basis_store: &BasisStore,
    ctx: &StageContext<'_>,
    fcf: &mut FutureCostFunction,
    cut_batches: &mut [RowBatch],
    training_ctx: &TrainingContext<'_>,
    spec: &mut BackwardPassSpec<'_>,
    comm: &C,
) -> Result<BackwardResult, SddpError> {
    let TrainingContext {
        horizon,
        indexer,
        stochastic,
        ..
    } = training_ctx;
    let num_stages = horizon.num_stages();
    let my_rank = comm.rank();

    debug_assert_eq!(ctx.templates.len(), num_stages);
    debug_assert_eq!(ctx.base_rows.len(), num_stages);
    debug_assert_eq!(spec.risk_measures.len(), num_stages);

    let start = Instant::now();
    let solves_before: u64 = workspaces
        .iter()
        .map(|ws| ws.solver.statistics().solve_count)
        .sum();
    let mut cuts_generated: usize = 0;
    let mut stage_stats: Vec<(usize, SolverStatsDelta)> = Vec::new();
    let mut state_exchange_ms: u64 = 0;
    let mut cut_batch_build_ms: u64 = 0;
    let mut rayon_overhead_ms: u64 = 0;
    let mut cut_sync_ms: u64 = 0;
    #[allow(clippy::cast_precision_loss)]
    let n_workers = workspaces.len() as f64;
    let tree_view = stochastic.tree_view();
    let mut staged_cuts_buf: Vec<StagedCut> = Vec::new();

    for t in (0..num_stages.saturating_sub(1)).rev() {
        // When the caller supplies forward-pass records, perform the per-stage
        // allgatherv here so that `exchange.state_at(rank, m)` returns the
        // state for stage `t` (the state entering stage `t+1`). When records
        // is empty the caller has pre-populated the exchange buffers (test path).
        if !spec.records.is_empty() {
            let exch_start = Instant::now();
            spec.exchange.exchange(spec.records, t, num_stages, comm)?;
            #[allow(clippy::cast_possible_truncation)]
            {
                state_exchange_ms += exch_start.elapsed().as_millis() as u64;
            }
        }

        // Archive gathered states for dominated cut selection (if active).
        if let Some(ref mut archive) = spec.visited_archive {
            let total_fwd = spec.exchange.total_scenarios();
            archive.archive_gathered_states(t, spec.exchange.gathered_states(), total_fwd);
        }

        let successor = t + 1;

        // Snapshot pool stats before this stage's solves.
        let stage_stats_before = {
            let pool_stats: Vec<_> = workspaces.iter().map(|w| w.solver.statistics()).collect();
            aggregate_solver_statistics(&pool_stats)
        };

        let n_openings = tree_view.n_openings(successor);
        spec.probabilities_buf.clear();
        #[allow(clippy::cast_precision_loss)]
        spec.probabilities_buf
            .resize(n_openings, 1.0_f64 / n_openings as f64);

        let batch_start = Instant::now();
        build_cut_row_batch_into(
            &mut cut_batches[successor],
            fcf,
            successor,
            indexer,
            &ctx.templates[successor].col_scale,
        );
        #[allow(clippy::cast_possible_truncation)]
        {
            cut_batch_build_ms += batch_start.elapsed().as_millis() as u64;
        }
        let num_cuts_at_successor = cut_batches[successor].num_rows;
        let template_num_rows = ctx.templates[successor].num_rows;
        spec.successor_active_slots_buf.clear();
        spec.successor_active_slots_buf
            .extend(fcf.active_cuts(successor).map(|(slot, _, _)| slot));

        let succ_spec = SuccessorSpec {
            t,
            successor,
            my_rank,
            probabilities: spec.probabilities_buf,
            cut_batch: &cut_batches[successor],
            num_cuts_at_successor,
            template_num_rows,
            successor_active_slots: spec.successor_active_slots_buf,
            basis_store,
            cut_activity_tolerance: spec.cut_activity_tolerance,
            successor_populated_count: fcf.pools[successor].populated_count,
        };
        let process_start = Instant::now();
        let worker_staged = process_stage_backward(workspaces, ctx, training_ctx, spec, &succ_spec);

        staged_cuts_buf.clear();
        for worker_result in worker_staged {
            staged_cuts_buf.extend(worker_result?);
        }
        // Ordering invariant: each worker processes a contiguous index range
        // [start_m, end_m) in ascending order, and rayon's indexed collect()
        // preserves worker order, so concatenation is globally sorted.
        debug_assert!(
            staged_cuts_buf
                .windows(2)
                .all(|w| w[0].trial_point_idx <= w[1].trial_point_idx),
            "backward pass cuts must be sorted by trial_point_idx after worker concatenation"
        );

        for cut in &staged_cuts_buf {
            fcf.add_cut(
                t,
                spec.iteration,
                cut.forward_pass_index,
                cut.intercept,
                &cut.coefficients,
            );
            cuts_generated += 1;
            for &(slot, increment) in &cut.binding_increments {
                fcf.pools[successor].metadata[slot].active_count += increment;
                fcf.pools[successor].metadata[slot].last_active_iter = spec.iteration;
            }
        }

        // Per-stage cut sync: allgatherv distributes this stage's local cuts
        // to all ranks so every rank sees the same FCF before the next stage.
        let sync_start = Instant::now();
        let n_local = spec.cut_sync_bufs.pack_local_cuts(fcf, t, spec.iteration);
        spec.cut_sync_bufs.sync_packed_cuts(t, n_local, fcf, comm)?;
        #[allow(clippy::cast_possible_truncation)]
        {
            cut_sync_ms += sync_start.elapsed().as_millis() as u64;
        }

        // Snapshot pool stats after this stage's solves and compute delta.
        let stage_stats_after = {
            let pool_stats: Vec<_> = workspaces.iter().map(|w| w.solver.statistics()).collect();
            aggregate_solver_statistics(&pool_stats)
        };
        let stage_delta = SolverStatsDelta::from_snapshots(&stage_stats_before, &stage_stats_after);

        // Rayon overhead: wall-clock of the parallel region minus the average
        // per-worker solve time. This is an upper bound on barrier + scheduling.
        #[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
        {
            let process_wall_ms = process_start.elapsed().as_millis() as u64;
            let avg_solve_ms = (stage_delta.solve_time_ms / n_workers) as u64;
            rayon_overhead_ms += process_wall_ms.saturating_sub(avg_solve_ms);
        }

        stage_stats.push((successor, stage_delta));
    }

    #[allow(clippy::cast_possible_truncation)]
    let elapsed_ms = start.elapsed().as_millis() as u64;
    let solves_after: u64 = workspaces
        .iter()
        .map(|ws| ws.solver.statistics().solve_count)
        .sum();

    Ok(BackwardResult {
        cuts_generated,
        elapsed_ms,
        lp_solves: solves_after - solves_before,
        stage_stats,
        state_exchange_time_ms: state_exchange_ms,
        cut_batch_build_time_ms: cut_batch_build_ms,
        rayon_overhead_time_ms: rayon_overhead_ms,
        cut_sync_time_ms: cut_sync_ms,
    })
}

#[cfg(test)]
mod tests {
    use cobre_comm::{CommData, CommError, Communicator, ReduceOp};
    use cobre_solver::{
        Basis, LpSolution, RowBatch, SolverError, SolverInterface, SolverStatistics, StageTemplate,
    };

    use super::{BackwardPassSpec, BackwardResult, run_backward_pass};
    use crate::{
        ExchangeBuffers, FutureCostFunction, HorizonMode, InflowNonNegativityMethod, RiskMeasure,
        StageIndexer, TrajectoryRecord,
        context::{StageContext, TrainingContext},
        cut_sync::CutSyncBuffers,
        workspace::{BasisStore, SolverWorkspace},
    };

    fn empty_cut_batches(n_stages: usize) -> Vec<RowBatch> {
        (0..n_stages)
            .map(|_| RowBatch {
                num_rows: 0,
                row_starts: Vec::new(),
                col_indices: Vec::new(),
                values: Vec::new(),
                row_lower: Vec::new(),
                row_upper: Vec::new(),
            })
            .collect()
    }

    /// Stub communicator for tests (single-rank).
    struct StubComm;

    impl Communicator for StubComm {
        fn allgatherv<T: CommData>(
            &self,
            send: &[T],
            recv: &mut [T],
            _counts: &[usize],
            _displs: &[usize],
        ) -> Result<(), CommError> {
            // Single-rank: copy send to recv (mirrors LocalBackend behavior).
            recv[..send.len()].copy_from_slice(send);
            Ok(())
        }

        fn allreduce<T: CommData>(
            &self,
            _send: &[T],
            _recv: &mut [T],
            _op: ReduceOp,
        ) -> Result<(), CommError> {
            unreachable!("StubComm allreduce not used in backward pass tests")
        }

        fn broadcast<T: CommData>(&self, _buf: &mut [T], _root: usize) -> Result<(), CommError> {
            unreachable!("StubComm broadcast not used in backward pass tests")
        }

        fn barrier(&self) -> Result<(), CommError> {
            Ok(())
        }

        fn rank(&self) -> usize {
            0
        }

        fn size(&self) -> usize {
            1
        }
    }

    /// Mock solver for testing: returns fixed solution or infeasible error on demand.
    ///
    /// Buffer fields (`buf_primal`, `buf_dual`, `buf_reduced_costs`) store the
    /// solution data that [`SolutionView`] borrows from. They are filled in
    /// `solve` before the borrow is established.
    struct MockSolver {
        solution: LpSolution,
        infeasible_at: Option<usize>,
        call_count: usize,
        /// Tracks the current number of rows (template + appended cuts).
        current_num_rows: usize,
        /// Number of times `solve_with_basis` was called.
        warm_start_calls: usize,
        buf_primal: Vec<f64>,
        buf_dual: Vec<f64>,
        buf_reduced_costs: Vec<f64>,
    }

    impl MockSolver {
        fn always_ok(solution: LpSolution) -> Self {
            let base_rows = solution.dual.len();
            let buf_primal = solution.primal.clone();
            let buf_dual = solution.dual.clone();
            let buf_reduced_costs = solution.reduced_costs.clone();
            Self {
                solution,
                infeasible_at: None,
                call_count: 0,
                current_num_rows: base_rows,
                warm_start_calls: 0,
                buf_primal,
                buf_dual,
                buf_reduced_costs,
            }
        }

        fn infeasible_on(solution: LpSolution, n: usize) -> Self {
            let base_rows = solution.dual.len();
            let buf_primal = solution.primal.clone();
            let buf_dual = solution.dual.clone();
            let buf_reduced_costs = solution.reduced_costs.clone();
            Self {
                solution,
                infeasible_at: Some(n),
                call_count: 0,
                current_num_rows: base_rows,
                warm_start_calls: 0,
                buf_primal,
                buf_dual,
                buf_reduced_costs,
            }
        }
    }

    impl SolverInterface for MockSolver {
        fn load_model(&mut self, template: &StageTemplate) {
            self.current_num_rows = template.num_rows;
        }

        fn add_rows(&mut self, cuts: &RowBatch) {
            self.current_num_rows += cuts.num_rows;
        }

        fn set_row_bounds(&mut self, _indices: &[usize], _lower: &[f64], _upper: &[f64]) {}
        fn set_col_bounds(&mut self, _indices: &[usize], _lower: &[f64], _upper: &[f64]) {}

        fn solve(&mut self) -> Result<cobre_solver::SolutionView<'_>, SolverError> {
            let call = self.call_count;
            self.call_count += 1;
            if self.infeasible_at == Some(call) {
                return Err(SolverError::Infeasible);
            }
            // Fill internal buffers, resizing dual to match current LP row count.
            self.buf_primal.clone_from(&self.solution.primal);
            self.buf_dual.clone_from(&self.solution.dual);
            self.buf_dual.resize(self.current_num_rows, 0.0);
            self.buf_reduced_costs
                .clone_from(&self.solution.reduced_costs);
            Ok(cobre_solver::SolutionView {
                objective: self.solution.objective,
                primal: &self.buf_primal,
                dual: &self.buf_dual,
                reduced_costs: &self.buf_reduced_costs,
                iterations: self.solution.iterations,
                solve_time_seconds: self.solution.solve_time_seconds,
            })
        }

        fn reset(&mut self) {}

        fn get_basis(&mut self, _out: &mut Basis) {}

        fn solve_with_basis(
            &mut self,
            _basis: &Basis,
        ) -> Result<cobre_solver::SolutionView<'_>, SolverError> {
            self.warm_start_calls += 1;
            self.solve()
        }

        fn statistics(&self) -> SolverStatistics {
            SolverStatistics::default()
        }

        fn name(&self) -> &'static str {
            "Mock"
        }
    }

    fn minimal_template_1_0() -> StageTemplate {
        StageTemplate {
            num_cols: 3,
            num_rows: 1,
            num_nz: 1,
            col_starts: vec![0_i32, 0, 1, 1],
            row_indices: vec![0_i32],
            values: vec![1.0],
            col_lower: vec![0.0, 0.0, 0.0],
            col_upper: vec![f64::INFINITY, f64::INFINITY, f64::INFINITY],
            objective: vec![0.0, 0.0, 1.0],
            row_lower: vec![0.0],
            row_upper: vec![0.0],
            n_state: 1,
            n_transfer: 0,
            n_dual_relevant: 1,
            n_hydro: 1,
            max_par_order: 0,
            col_scale: Vec::new(),
            row_scale: Vec::new(),
        }
    }

    fn solution_1_0(objective: f64, dual_storage: f64) -> LpSolution {
        LpSolution {
            objective,
            primal: vec![0.0, 0.0, 0.0],
            dual: vec![dual_storage],
            reduced_costs: vec![0.0; 3],
            iterations: 0,
            solve_time_seconds: 0.0,
        }
    }

    /// Wrap a `MockSolver` into a single-element `Vec<SolverWorkspace<MockSolver>>`
    /// for tests that exercise the workspace-based backward-pass API.
    ///
    /// The workspace is sized for `n_hydro=1`, `max_par_order=0`, and `n_state`
    /// state dimensions.
    fn single_workspace(solver: MockSolver, n_state: usize) -> Vec<SolverWorkspace<MockSolver>> {
        use crate::lp_builder::PatchBuffer;
        vec![SolverWorkspace {
            solver,
            patch_buf: PatchBuffer::new(1, 0, 0, 0),
            current_state: Vec::with_capacity(n_state),
            scratch: crate::workspace::ScratchBuffers {
                noise_buf: Vec::new(),
                inflow_m3s_buf: Vec::new(),
                lag_matrix_buf: Vec::new(),
                par_inflow_buf: Vec::new(),
                eta_floor_buf: Vec::new(),
                zero_targets_buf: Vec::new(),
                ncs_col_upper_buf: Vec::new(),
                ncs_col_lower_buf: Vec::new(),
                ncs_col_indices_buf: Vec::new(),
                load_rhs_buf: Vec::new(),
                row_lower_buf: Vec::new(),
                z_inflow_rhs_buf: Vec::new(),
                effective_eta_buf: Vec::new(),
                unscaled_primal: Vec::new(),
                unscaled_dual: Vec::new(),
            },
        }]
    }

    /// Create an empty `BasisStore` for `num_scenarios` scenarios and
    /// `num_stages` stages (all slots `None`).
    fn empty_basis_store(num_scenarios: usize, num_stages: usize) -> BasisStore {
        BasisStore::new(num_scenarios, num_stages)
    }

    /// Create a `BasisStore` with one slot pre-populated at
    /// `[scenario][stage]` with the given `Basis`.
    fn basis_store_with_one(
        num_scenarios: usize,
        num_stages: usize,
        scenario: usize,
        stage: usize,
        basis: Basis,
    ) -> BasisStore {
        let mut store = BasisStore::new(num_scenarios, num_stages);
        *store.get_mut(scenario, stage) = Some(basis);
        store
    }

    fn exchange_with_states(n_state: usize, states: Vec<Vec<f64>>) -> ExchangeBuffers {
        use cobre_comm::LocalBackend;

        let local_count = states.len();
        let mut bufs = ExchangeBuffers::new(n_state, local_count, 1);
        let records: Vec<TrajectoryRecord> = states
            .into_iter()
            .map(|state| TrajectoryRecord {
                primal: vec![],
                dual: vec![],
                stage_cost: 0.0,
                state,
            })
            .collect();

        let comm = LocalBackend;
        bufs.exchange(&records, 0, 1, &comm).unwrap();
        bufs
    }

    #[allow(clippy::too_many_lines)]
    fn make_stochastic_context(
        n_stages: usize,
        branching_factor: usize,
    ) -> cobre_stochastic::StochasticContext {
        use chrono::NaiveDate;
        use cobre_core::entities::hydro::{Hydro, HydroGenerationModel, HydroPenalties};
        use cobre_core::{
            Bus, DeficitSegment, EntityId, SystemBuilder,
            scenario::{
                CorrelationEntity, CorrelationGroup, CorrelationModel, CorrelationProfile,
                InflowModel,
            },
            temporal::{
                Block, BlockMode, NoiseMethod, ScenarioSourceConfig, Stage, StageRiskConfig,
                StageStateConfig,
            },
        };
        use cobre_stochastic::context::build_stochastic_context;
        use std::collections::BTreeMap;

        let bus = Bus {
            id: EntityId(0),
            name: "B0".to_string(),
            deficit_segments: vec![DeficitSegment {
                depth_mw: None,
                cost_per_mwh: 1000.0,
            }],
            excess_cost: 0.0,
        };
        let hydro = Hydro {
            id: EntityId(1),
            name: "H1".to_string(),
            bus_id: EntityId(0),
            downstream_id: None,
            entry_stage_id: None,
            exit_stage_id: None,
            min_storage_hm3: 0.0,
            max_storage_hm3: 100.0,
            min_outflow_m3s: 0.0,
            max_outflow_m3s: None,
            generation_model: HydroGenerationModel::ConstantProductivity {
                productivity_mw_per_m3s: 1.0,
            },
            min_turbined_m3s: 0.0,
            max_turbined_m3s: 100.0,
            min_generation_mw: 0.0,
            max_generation_mw: 100.0,
            tailrace: None,
            hydraulic_losses: None,
            efficiency: None,
            evaporation_coefficients_mm: None,
            evaporation_reference_volumes_hm3: None,
            diversion: None,
            filling: None,
            penalties: HydroPenalties {
                spillage_cost: 0.0,
                diversion_cost: 0.0,
                fpha_turbined_cost: 0.0,
                storage_violation_below_cost: 0.0,
                filling_target_violation_cost: 0.0,
                turbined_violation_below_cost: 0.0,
                outflow_violation_below_cost: 0.0,
                outflow_violation_above_cost: 0.0,
                generation_violation_below_cost: 0.0,
                evaporation_violation_cost: 0.0,
                water_withdrawal_violation_cost: 0.0,
                water_withdrawal_violation_pos_cost: 0.0,
                water_withdrawal_violation_neg_cost: 0.0,
                evaporation_violation_pos_cost: 0.0,
                evaporation_violation_neg_cost: 0.0,
                inflow_nonnegativity_cost: 1000.0,
            },
        };

        #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
        let make_stage = |idx: usize| Stage {
            index: idx,
            id: idx as i32,
            start_date: NaiveDate::from_ymd_opt(2024, 1, 1).unwrap(),
            end_date: NaiveDate::from_ymd_opt(2024, 2, 1).unwrap(),
            season_id: Some(0),
            blocks: vec![Block {
                index: 0,
                name: "S".to_string(),
                duration_hours: 744.0,
            }],
            block_mode: BlockMode::Parallel,
            state_config: StageStateConfig {
                storage: true,
                inflow_lags: false,
            },
            risk_config: StageRiskConfig::Expectation,
            scenario_config: ScenarioSourceConfig {
                branching_factor,
                noise_method: NoiseMethod::Saa,
            },
        };

        let stages: Vec<Stage> = (0..n_stages).map(make_stage).collect();

        #[allow(clippy::cast_possible_truncation)]
        let inflow = |stage_idx: usize| InflowModel {
            hydro_id: EntityId(1),
            #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
            stage_id: stage_idx as i32,
            mean_m3s: 100.0,
            std_m3s: 30.0,
            ar_coefficients: vec![],
            residual_std_ratio: 1.0,
        };

        let inflow_models: Vec<InflowModel> = (0..n_stages).map(inflow).collect();

        let mut profiles = BTreeMap::new();
        profiles.insert(
            "default".to_string(),
            CorrelationProfile {
                groups: vec![CorrelationGroup {
                    name: "g1".to_string(),
                    entities: vec![CorrelationEntity {
                        entity_type: "inflow".to_string(),
                        id: EntityId(1),
                    }],
                    matrix: vec![vec![1.0]],
                }],
            },
        );
        let correlation = CorrelationModel {
            method: "cholesky".to_string(),
            profiles,
            schedule: vec![],
        };

        let system = SystemBuilder::new()
            .buses(vec![bus])
            .hydros(vec![hydro])
            .stages(stages)
            .inflow_models(inflow_models)
            .correlation(correlation)
            .build()
            .unwrap();

        build_stochastic_context(&system, 42, &[], &[], None).unwrap()
    }

    // ── Unit tests ────────────────────────────────────────────────────────────

    #[test]
    fn backward_result_fields_accessible() {
        let r = BackwardResult {
            cuts_generated: 6,
            elapsed_ms: 42,
            lp_solves: 0,
            stage_stats: Vec::new(),
            state_exchange_time_ms: 0,
            cut_batch_build_time_ms: 0,
            rayon_overhead_time_ms: 0,
            cut_sync_time_ms: 0,
        };
        assert_eq!(r.cuts_generated, 6);
        assert_eq!(r.elapsed_ms, 42);
        assert!(r.stage_stats.is_empty());
        assert_eq!(r.state_exchange_time_ms, 0);
        assert_eq!(r.cut_batch_build_time_ms, 0);
        assert_eq!(r.rayon_overhead_time_ms, 0);
        assert_eq!(r.cut_sync_time_ms, 0);
    }

    #[test]
    fn backward_result_clone_and_debug() {
        let r = BackwardResult {
            cuts_generated: 3,
            elapsed_ms: 100,
            lp_solves: 0,
            stage_stats: Vec::new(),
            state_exchange_time_ms: 0,
            cut_batch_build_time_ms: 0,
            rayon_overhead_time_ms: 0,
            cut_sync_time_ms: 0,
        };
        let c = r.clone();
        assert_eq!(c.cuts_generated, 3);
        let s = format!("{r:?}");
        assert!(s.contains("BackwardResult"));
    }

    #[test]
    fn dual_extraction_formula_coefficients_are_negated_duals() {
        // Given known dual values [d0, d1], coefficients must be [-d0, -d1].
        let d0 = 3.5_f64;
        let d1 = -1.2_f64;
        let dual = [d0, d1];

        let coefficients: Vec<f64> = dual.iter().map(|&d| -d).collect();

        assert!((coefficients[0] - (-d0)).abs() < f64::EPSILON);
        assert!((coefficients[1] - (-d1)).abs() < f64::EPSILON);
    }

    #[test]
    fn intercept_formula_matches_spec() {
        // alpha = Q - pi^T * x_hat
        // Given: objective=50.0, pi=[2.0, -1.0], x_hat=[10.0, 5.0]
        // Expected: alpha = 50.0 - (2.0*10.0 + (-1.0)*5.0) = 50.0 - 15.0 = 35.0
        let objective = 50.0_f64;
        let coefficients = [2.0_f64, -1.0_f64];
        let x_hat = [10.0_f64, 5.0_f64];
        let pi_dot_x: f64 = coefficients
            .iter()
            .zip(x_hat.iter())
            .map(|(p, x)| p * x)
            .sum();
        let intercept = objective - pi_dot_x;
        assert!((intercept - 35.0).abs() < f64::EPSILON);
    }

    #[test]
    fn single_stage_system_produces_no_cuts() {
        // A 1-stage system has no stages with a successor, so the backward
        // sweep (0..0) is empty — zero cuts are generated.
        let stochastic = make_stochastic_context(1, 2);
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template_1_0()];
        let base_rows = vec![1_usize];

        let n_state = indexer.n_state;
        let n_stages = 1_usize;
        let forward_passes = 2_u32;
        let mut fcf = FutureCostFunction::new(n_stages, n_state, forward_passes, 10, 0);
        let mut exchange = exchange_with_states(n_state, vec![vec![10.0], vec![20.0]]);

        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let risk_measures = vec![RiskMeasure::Expectation];

        let solution = solution_1_0(100.0, -5.0);
        let solver = MockSolver::always_ok(solution);
        let comm = StubComm;
        let mut workspaces = single_workspace(solver, n_state);
        let basis_store = empty_basis_store(exchange.local_count(), n_stages);

        let mut csb = CutSyncBuffers::new(n_state, 64, 1);
        let result = run_backward_pass(
            &mut workspaces,
            &basis_store,
            &StageContext {
                templates: &templates,
                base_rows: &base_rows,
                noise_scale: &[],
                n_hydros: 0,
                n_load_buses: 0,
                load_balance_row_starts: &[],
                load_bus_indices: &[],
                block_counts_per_stage: &[],
                ncs_max_gen: &[],
                discount_factors: &[],
                cumulative_discount_factors: &[],
            },
            &mut fcf,
            &mut empty_cut_batches(templates.len()),
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &[],
            },
            &mut BackwardPassSpec {
                records: &[],
                iteration: 0,
                local_work: exchange.local_count(),
                fwd_offset: 0,
                risk_measures: &risk_measures,
                exchange: &mut exchange,
                cut_activity_tolerance: 0.0,
                cut_sync_bufs: &mut csb,
                probabilities_buf: &mut Vec::new(),
                successor_active_slots_buf: &mut Vec::new(),
                visited_archive: None,
            },
            &comm,
        )
        .unwrap();

        assert_eq!(result.cuts_generated, 0);
        assert_eq!(fcf.total_active_cuts(), 0);
    }

    #[test]
    fn two_stage_system_two_trial_points_generates_two_cuts_at_stage_0() {
        // Acceptance criterion: 3-stage system, 1 hydro (n_state=1), 2 openings,
        // 2 trial points → 2 cuts at stage 0. This is the 2-stage version
        // (stages 0 and 1); cuts should exist only at stage 0.
        let n_stages = 2_usize;
        let n_openings = 2_usize;
        let stochastic = make_stochastic_context(n_stages, n_openings);
        let indexer = StageIndexer::new(1, 0); // N=1, L=0
        let templates = vec![minimal_template_1_0(); n_stages];
        let base_rows = vec![1_usize; n_stages];

        let n_state = indexer.n_state; // 1
        let forward_passes = 2_u32;
        let mut fcf = FutureCostFunction::new(n_stages, n_state, forward_passes, 10, 0);

        // Two trial points with states [10.0] and [20.0] at stage 0.
        let mut exchange = exchange_with_states(n_state, vec![vec![10.0], vec![20.0]]);

        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let risk_measures = vec![RiskMeasure::Expectation; n_stages];

        // MockSolver returns objective=100.0, dual[0]=-5.0 for every solve.
        // With x_hat=[10.0]: pi=[5.0], alpha = 100 - 5*10 = 50.
        // With x_hat=[20.0]: pi=[5.0], alpha = 100 - 5*20 = 0 (could be negative).
        let solution = solution_1_0(100.0, -5.0);
        let solver = MockSolver::always_ok(solution);
        let comm = StubComm;
        let mut workspaces = single_workspace(solver, n_state);
        let basis_store = empty_basis_store(exchange.local_count(), n_stages);

        let mut csb = CutSyncBuffers::new(n_state, 64, 1);
        let result = run_backward_pass(
            &mut workspaces,
            &basis_store,
            &StageContext {
                templates: &templates,
                base_rows: &base_rows,
                noise_scale: &[],
                n_hydros: 0,
                n_load_buses: 0,
                load_balance_row_starts: &[],
                load_bus_indices: &[],
                block_counts_per_stage: &[],
                ncs_max_gen: &[],
                discount_factors: &[],
                cumulative_discount_factors: &[],
            },
            &mut fcf,
            &mut empty_cut_batches(templates.len()),
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &[],
            },
            &mut BackwardPassSpec {
                records: &[],
                iteration: 0,
                local_work: exchange.local_count(),
                fwd_offset: 0,
                risk_measures: &risk_measures,
                exchange: &mut exchange,
                cut_activity_tolerance: 0.0,
                cut_sync_bufs: &mut csb,
                probabilities_buf: &mut Vec::new(),
                successor_active_slots_buf: &mut Vec::new(),
                visited_archive: None,
            },
            &comm,
        )
        .unwrap();

        // 2 trial points × 1 stage with a successor = 2 cuts at stage 0.
        assert_eq!(result.cuts_generated, 2);
        assert_eq!(fcf.active_cuts(0).count(), 2);
        // Stage 1 (the last stage) gets no cuts.
        assert_eq!(fcf.active_cuts(1).count(), 0);
    }

    #[test]
    fn cut_inserted_with_correct_stage_iteration_and_forward_pass_index() {
        // Acceptance criterion: iteration=2, forward_passes=3, global
        // trial point m=1 → fcf.add_cut(stage=0, iteration=2, fpi=1, ...).
        // slot = warm_start + 2*3 + 1 = 7.
        let n_stages = 2_usize;
        let n_openings = 2_usize;
        let stochastic = make_stochastic_context(n_stages, n_openings);
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template_1_0(); n_stages];
        let base_rows = vec![1_usize; n_stages];

        let n_state = indexer.n_state;
        let forward_passes = 3_u32;
        let mut fcf = FutureCostFunction::new(n_stages, n_state, forward_passes, 20, 0);

        // 3 trial points (forward_passes=3 on a single rank).
        let mut exchange = exchange_with_states(n_state, vec![vec![5.0], vec![10.0], vec![15.0]]);

        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let risk_measures = vec![RiskMeasure::Expectation; n_stages];

        let solution = solution_1_0(50.0, 0.0);
        let solver = MockSolver::always_ok(solution);
        let comm = StubComm;
        let mut workspaces = single_workspace(solver, n_state);
        let basis_store = empty_basis_store(exchange.local_count(), n_stages);

        let mut csb = CutSyncBuffers::new(n_state, 64, 1);
        let _ = run_backward_pass(
            &mut workspaces,
            &basis_store,
            &StageContext {
                templates: &templates,
                base_rows: &base_rows,
                noise_scale: &[],
                n_hydros: 0,
                n_load_buses: 0,
                load_balance_row_starts: &[],
                load_bus_indices: &[],
                block_counts_per_stage: &[],
                ncs_max_gen: &[],
                discount_factors: &[],
                cumulative_discount_factors: &[],
            },
            &mut fcf,
            &mut empty_cut_batches(templates.len()),
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &[],
            },
            &mut BackwardPassSpec {
                records: &[],
                iteration: 2,
                local_work: exchange.local_count(),
                fwd_offset: 0,
                risk_measures: &risk_measures,
                exchange: &mut exchange,
                cut_activity_tolerance: 0.0,
                cut_sync_bufs: &mut csb,
                probabilities_buf: &mut Vec::new(),
                successor_active_slots_buf: &mut Vec::new(),
                visited_archive: None,
            },
            &comm,
        )
        .unwrap();

        // Trial point m=1: slot = 0 + 2*3 + 1 = 7
        // Verify that pool[0].metadata[7] has the correct iteration and fpi.
        let meta = &fcf.pools[0].metadata[7];
        assert_eq!(meta.iteration_generated, 2);
        assert_eq!(meta.forward_pass_index, 1);
    }

    #[test]
    fn no_cuts_generated_at_last_stage() {
        // Acceptance criterion: 5-stage system → cuts at stages 0..3, not at 4.
        let n_stages = 5_usize;
        let n_openings = 2_usize;
        let stochastic = make_stochastic_context(n_stages, n_openings);
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template_1_0(); n_stages];
        let base_rows = vec![1_usize; n_stages];

        let n_state = indexer.n_state;
        let forward_passes = 1_u32;
        let mut fcf = FutureCostFunction::new(n_stages, n_state, forward_passes, 10, 0);
        let mut exchange = exchange_with_states(n_state, vec![vec![10.0]]);

        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let risk_measures = vec![RiskMeasure::Expectation; n_stages];

        let solution = solution_1_0(100.0, -3.0);
        let solver = MockSolver::always_ok(solution);
        let comm = StubComm;
        let mut workspaces = single_workspace(solver, n_state);
        let basis_store = empty_basis_store(exchange.local_count(), n_stages);

        let mut csb = CutSyncBuffers::new(n_state, 64, 1);
        let result = run_backward_pass(
            &mut workspaces,
            &basis_store,
            &StageContext {
                templates: &templates,
                base_rows: &base_rows,
                noise_scale: &[],
                n_hydros: 0,
                n_load_buses: 0,
                load_balance_row_starts: &[],
                load_bus_indices: &[],
                block_counts_per_stage: &[],
                ncs_max_gen: &[],
                discount_factors: &[],
                cumulative_discount_factors: &[],
            },
            &mut fcf,
            &mut empty_cut_batches(templates.len()),
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &[],
            },
            &mut BackwardPassSpec {
                records: &[],
                iteration: 0,
                local_work: exchange.local_count(),
                fwd_offset: 0,
                risk_measures: &risk_measures,
                exchange: &mut exchange,
                cut_activity_tolerance: 0.0,
                cut_sync_bufs: &mut csb,
                probabilities_buf: &mut Vec::new(),
                successor_active_slots_buf: &mut Vec::new(),
                visited_archive: None,
            },
            &comm,
        )
        .unwrap();

        // 1 trial point × 4 stages with successors = 4 cuts total.
        assert_eq!(result.cuts_generated, 4);
        for t in 0..4 {
            assert_eq!(fcf.active_cuts(t).count(), 1, "stage {t} should have 1 cut");
        }
        // The last stage (4) must have no cuts.
        assert_eq!(fcf.active_cuts(4).count(), 0, "stage 4 must have no cuts");
    }

    #[test]
    fn elapsed_ms_is_non_negative() {
        let n_stages = 2_usize;
        let stochastic = make_stochastic_context(n_stages, 2);
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template_1_0(); n_stages];
        let base_rows = vec![1_usize; n_stages];

        let n_state = indexer.n_state;
        let forward_passes = 1_u32;
        let mut fcf = FutureCostFunction::new(n_stages, n_state, forward_passes, 10, 0);
        let mut exchange = exchange_with_states(n_state, vec![vec![5.0]]);

        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let risk_measures = vec![RiskMeasure::Expectation; n_stages];

        let solution = solution_1_0(10.0, 0.0);
        let solver = MockSolver::always_ok(solution);
        let comm = StubComm;
        let mut workspaces = single_workspace(solver, n_state);
        let basis_store = empty_basis_store(exchange.local_count(), n_stages);

        let mut csb = CutSyncBuffers::new(n_state, 64, 1);
        let result = run_backward_pass(
            &mut workspaces,
            &basis_store,
            &StageContext {
                templates: &templates,
                base_rows: &base_rows,
                noise_scale: &[],
                n_hydros: 0,
                n_load_buses: 0,
                load_balance_row_starts: &[],
                load_bus_indices: &[],
                block_counts_per_stage: &[],
                ncs_max_gen: &[],
                discount_factors: &[],
                cumulative_discount_factors: &[],
            },
            &mut fcf,
            &mut empty_cut_batches(templates.len()),
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &[],
            },
            &mut BackwardPassSpec {
                records: &[],
                iteration: 0,
                local_work: exchange.local_count(),
                fwd_offset: 0,
                risk_measures: &risk_measures,
                exchange: &mut exchange,
                cut_activity_tolerance: 0.0,
                cut_sync_bufs: &mut csb,
                probabilities_buf: &mut Vec::new(),
                successor_active_slots_buf: &mut Vec::new(),
                visited_archive: None,
            },
            &comm,
        )
        .unwrap();

        // elapsed_ms is u64, so it is always >= 0.
        let _ = result.elapsed_ms;
    }

    #[test]
    fn infeasible_solver_returns_sddp_infeasible_error() {
        // Acceptance criterion: MockSolver::infeasible_on(0) for the first
        // backward solve → SddpError::Infeasible is returned.
        let n_stages = 2_usize;
        let stochastic = make_stochastic_context(n_stages, 1);
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template_1_0(); n_stages];
        let base_rows = vec![1_usize; n_stages];

        let n_state = indexer.n_state;
        let forward_passes = 1_u32;
        let mut fcf = FutureCostFunction::new(n_stages, n_state, forward_passes, 10, 0);
        let mut exchange = exchange_with_states(n_state, vec![vec![10.0]]);

        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let risk_measures = vec![RiskMeasure::Expectation; n_stages];

        let solution = solution_1_0(0.0, 0.0);
        // First solve call returns infeasible.
        let solver = MockSolver::infeasible_on(solution, 0);
        let comm = StubComm;
        let mut workspaces = single_workspace(solver, n_state);
        let basis_store = empty_basis_store(exchange.local_count(), n_stages);

        let mut csb = CutSyncBuffers::new(n_state, 64, 1);
        let result = run_backward_pass(
            &mut workspaces,
            &basis_store,
            &StageContext {
                templates: &templates,
                base_rows: &base_rows,
                noise_scale: &[],
                n_hydros: 0,
                n_load_buses: 0,
                load_balance_row_starts: &[],
                load_bus_indices: &[],
                block_counts_per_stage: &[],
                ncs_max_gen: &[],
                discount_factors: &[],
                cumulative_discount_factors: &[],
            },
            &mut fcf,
            &mut empty_cut_batches(templates.len()),
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &[],
            },
            &mut BackwardPassSpec {
                records: &[],
                iteration: 0,
                local_work: exchange.local_count(),
                fwd_offset: 0,
                risk_measures: &risk_measures,
                exchange: &mut exchange,
                cut_activity_tolerance: 0.0,
                cut_sync_bufs: &mut csb,
                probabilities_buf: &mut Vec::new(),
                successor_active_slots_buf: &mut Vec::new(),
                visited_archive: None,
            },
            &comm,
        );

        assert!(
            matches!(result, Err(crate::SddpError::Infeasible { .. })),
            "expected SddpError::Infeasible, got: {result:?}",
        );
    }

    #[test]
    fn expectation_aggregation_mean_of_per_opening_intercepts() {
        // Given 3 openings with uniform probability 1/3 and per-opening
        // intercepts [10.0, 20.0, 30.0], the aggregated intercept must be 20.0.
        use crate::risk_measure::BackwardOutcome as BO;

        let outcomes = vec![
            BO {
                intercept: 10.0,
                coefficients: vec![],
                objective_value: 10.0,
            },
            BO {
                intercept: 20.0,
                coefficients: vec![],
                objective_value: 20.0,
            },
            BO {
                intercept: 30.0,
                coefficients: vec![],
                objective_value: 30.0,
            },
        ];
        let probs = vec![1.0 / 3.0; 3];
        let (intercept, _) = RiskMeasure::Expectation.aggregate_cut(&outcomes, &probs);
        assert!(
            (intercept - 20.0).abs() < 1e-10,
            "expected 20.0, got {intercept}"
        );
    }

    // ── Integration tests ─────────────────────────────────────────────────────

    #[test]
    fn cut_coefficients_and_intercept_match_dual_extraction_formula() {
        // Integration test: verify that the backward pass uses the correct
        // dual extraction formula by checking cuts in the FCF.
        //
        // Setup: 2-stage, N=1, L=0, 1 opening, 1 trial point.
        //   dual[0] = -3.0 (storage-fixing dual from MockSolver)
        //   objective = 80.0
        //   x_hat = [10.0]
        //
        // Expected (coefficients = dual, not -dual):
        //   pi[0] = dual[0] = -3.0
        //   intercept = 80.0 - (-3.0) * 10.0 = 110.0
        //   coefficients = [-3.0]
        let n_stages = 2_usize;
        let stochastic = make_stochastic_context(n_stages, 1);
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template_1_0(); n_stages];
        let base_rows = vec![1_usize; n_stages];

        let n_state = indexer.n_state;
        let forward_passes = 1_u32;
        let mut fcf = FutureCostFunction::new(n_stages, n_state, forward_passes, 10, 0);
        let mut exchange = exchange_with_states(n_state, vec![vec![10.0]]);

        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let risk_measures = vec![RiskMeasure::Expectation; n_stages];

        // dual[0] = -3.0, objective = 80.0
        let solution = solution_1_0(80.0, -3.0);
        let solver = MockSolver::always_ok(solution);
        let comm = StubComm;
        let mut workspaces = single_workspace(solver, n_state);
        let basis_store = empty_basis_store(exchange.local_count(), n_stages);

        let mut csb = CutSyncBuffers::new(n_state, 64, 1);
        let _ = run_backward_pass(
            &mut workspaces,
            &basis_store,
            &StageContext {
                templates: &templates,
                base_rows: &base_rows,
                noise_scale: &[],
                n_hydros: 0,
                n_load_buses: 0,
                load_balance_row_starts: &[],
                load_bus_indices: &[],
                block_counts_per_stage: &[],
                ncs_max_gen: &[],
                discount_factors: &[],
                cumulative_discount_factors: &[],
            },
            &mut fcf,
            &mut empty_cut_batches(templates.len()),
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &[],
            },
            &mut BackwardPassSpec {
                records: &[],
                iteration: 0,
                local_work: exchange.local_count(),
                fwd_offset: 0,
                risk_measures: &risk_measures,
                exchange: &mut exchange,
                cut_activity_tolerance: 0.0,
                cut_sync_bufs: &mut csb,
                probabilities_buf: &mut Vec::new(),
                successor_active_slots_buf: &mut Vec::new(),
                visited_archive: None,
            },
            &comm,
        )
        .unwrap();

        let cuts: Vec<_> = fcf.active_cuts(0).collect();
        assert_eq!(cuts.len(), 1);
        let (_, intercept, coefficients) = &cuts[0];

        assert!(
            (intercept - 110.0).abs() < 1e-10,
            "expected intercept=110.0, got {intercept}"
        );
        assert_eq!(coefficients.len(), 1);
        assert!(
            (coefficients[0] - (-3.0)).abs() < 1e-10,
            "expected coefficient=-3.0, got {}",
            coefficients[0]
        );
    }

    #[test]
    fn cut_gradient_sign_physically_correct() {
        // Regression test for the Benders cut sign bug.
        //
        // Physical invariant: more initial storage → lower future cost.
        // The storage-fixing dual π is negative (shadow price of relaxing
        // the fixing constraint increases cost when storage decreases).
        //
        // Correct: coefficient = π < 0, so the cut slope is negative
        //   (more storage → lower cut value → lower theta → lower total cost).
        //
        // Old bug: coefficient = -π > 0, so the cut slope was positive
        //   (more storage → higher cut value → wrong incentive).
        let n_stages = 2_usize;
        let stochastic = make_stochastic_context(n_stages, 1);
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template_1_0(); n_stages];
        let base_rows = vec![1_usize; n_stages];

        let n_state = indexer.n_state;
        let forward_passes = 1_u32;
        let mut fcf = FutureCostFunction::new(n_stages, n_state, forward_passes, 10, 0);
        let mut exchange = exchange_with_states(n_state, vec![vec![50.0]]);

        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let risk_measures = vec![RiskMeasure::Expectation; n_stages];

        // dual[0] = -2.0 (negative: more storage → less cost)
        // objective = 100.0, x_hat = 50.0
        let solution = solution_1_0(100.0, -2.0);
        let solver = MockSolver::always_ok(solution);
        let comm = StubComm;
        let mut workspaces = single_workspace(solver, n_state);
        let basis_store = empty_basis_store(exchange.local_count(), n_stages);

        let mut csb = CutSyncBuffers::new(n_state, 64, 1);
        let _ = run_backward_pass(
            &mut workspaces,
            &basis_store,
            &StageContext {
                templates: &templates,
                base_rows: &base_rows,
                noise_scale: &[],
                n_hydros: 0,
                n_load_buses: 0,
                load_balance_row_starts: &[],
                load_bus_indices: &[],
                block_counts_per_stage: &[],
                ncs_max_gen: &[],
                discount_factors: &[],
                cumulative_discount_factors: &[],
            },
            &mut fcf,
            &mut empty_cut_batches(templates.len()),
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &[],
            },
            &mut BackwardPassSpec {
                records: &[],
                iteration: 0,
                local_work: exchange.local_count(),
                fwd_offset: 0,
                risk_measures: &risk_measures,
                exchange: &mut exchange,
                cut_activity_tolerance: 0.0,
                cut_sync_bufs: &mut csb,
                probabilities_buf: &mut Vec::new(),
                successor_active_slots_buf: &mut Vec::new(),
                visited_archive: None,
            },
            &comm,
        )
        .unwrap();

        let cuts: Vec<_> = fcf.active_cuts(0).collect();
        assert_eq!(cuts.len(), 1, "expected exactly one cut");
        let (_, _intercept, coefficients) = &cuts[0];

        // The coefficient must be negative (same sign as the dual).
        // The old bug would produce +2.0 here instead of -2.0.
        assert!(
            coefficients[0] < 0.0,
            "cut coefficient must be negative (more storage → less future cost), \
             got {} — likely the Benders cut sign bug has been reintroduced",
            coefficients[0]
        );
        assert!(
            (coefficients[0] - (-2.0)).abs() < 1e-10,
            "expected coefficient=-2.0, got {}",
            coefficients[0]
        );
    }

    #[test]
    fn cut_is_tight_at_trial_point() {
        // Regression test: a Benders cut must be tight (exact) at the trial
        // point x̂ where it was generated. That is:
        //   intercept + coefficient * x̂ = Q(x̂)
        // where Q(x̂) = objective value of the subproblem at x̂.
        //
        // The cut equation is: θ ≥ intercept + coefficient * x
        // At x = x̂: θ ≥ Q(x̂) + π'(x̂ - x̂) = Q(x̂)
        //
        // If the sign is wrong (coefficient = -π instead of π), then:
        //   intercept + (-π) * x̂ ≠ Q(x̂) in general
        //
        // This test verifies the tightness property.
        let n_stages = 2_usize;
        let stochastic = make_stochastic_context(n_stages, 1);
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template_1_0(); n_stages];
        let base_rows = vec![1_usize; n_stages];

        let n_state = indexer.n_state;
        let forward_passes = 1_u32;
        let mut fcf = FutureCostFunction::new(n_stages, n_state, forward_passes, 10, 0);
        let x_hat = 30.0_f64;
        let mut exchange = exchange_with_states(n_state, vec![vec![x_hat]]);

        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let risk_measures = vec![RiskMeasure::Expectation; n_stages];

        let q_xhat = 200.0_f64; // subproblem objective at x̂
        let dual_storage = -4.0_f64;
        let solution = solution_1_0(q_xhat, dual_storage);
        let solver = MockSolver::always_ok(solution);
        let comm = StubComm;
        let mut workspaces = single_workspace(solver, n_state);
        let basis_store = empty_basis_store(exchange.local_count(), n_stages);

        let mut csb = CutSyncBuffers::new(n_state, 64, 1);
        let _ = run_backward_pass(
            &mut workspaces,
            &basis_store,
            &StageContext {
                templates: &templates,
                base_rows: &base_rows,
                noise_scale: &[],
                n_hydros: 0,
                n_load_buses: 0,
                load_balance_row_starts: &[],
                load_bus_indices: &[],
                block_counts_per_stage: &[],
                ncs_max_gen: &[],
                discount_factors: &[],
                cumulative_discount_factors: &[],
            },
            &mut fcf,
            &mut empty_cut_batches(templates.len()),
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &[],
            },
            &mut BackwardPassSpec {
                records: &[],
                iteration: 0,
                local_work: exchange.local_count(),
                fwd_offset: 0,
                risk_measures: &risk_measures,
                exchange: &mut exchange,
                cut_activity_tolerance: 0.0,
                cut_sync_bufs: &mut csb,
                probabilities_buf: &mut Vec::new(),
                successor_active_slots_buf: &mut Vec::new(),
                visited_archive: None,
            },
            &comm,
        )
        .unwrap();

        let cuts: Vec<_> = fcf.active_cuts(0).collect();
        assert_eq!(cuts.len(), 1);
        let (_, intercept, coefficients) = &cuts[0];

        // Evaluate the cut at x̂: cut_value = intercept + coeff * x̂
        let cut_at_xhat = intercept + coefficients[0] * x_hat;

        // Must equal Q(x̂) (tightness property)
        assert!(
            (cut_at_xhat - q_xhat).abs() < 1e-10,
            "cut must be tight at trial point: \
             cut_value={cut_at_xhat}, Q(x̂)={q_xhat}, \
             intercept={intercept}, coeff={}, x̂={x_hat}",
            coefficients[0]
        );
    }

    #[test]
    fn single_rank_backward_pass_with_local_backend_produces_correct_fcf() {
        // Integration test with LocalBackend communicator (exercises single-rank path).
        use cobre_comm::LocalBackend;

        let n_stages = 3_usize;
        let n_openings = 2_usize;
        let stochastic = make_stochastic_context(n_stages, n_openings);
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template_1_0(); n_stages];
        let base_rows = vec![1_usize; n_stages];

        let n_state = indexer.n_state;
        let forward_passes = 2_u32;
        let mut fcf = FutureCostFunction::new(n_stages, n_state, forward_passes, 10, 0);
        let mut exchange = exchange_with_states(n_state, vec![vec![10.0], vec![20.0]]);

        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let risk_measures = vec![RiskMeasure::Expectation; n_stages];

        let solution = solution_1_0(100.0, -5.0);
        let solver = MockSolver::always_ok(solution);
        let comm = LocalBackend;
        let mut workspaces = single_workspace(solver, n_state);
        let basis_store = empty_basis_store(exchange.local_count(), n_stages);

        let mut csb = CutSyncBuffers::new(n_state, 64, 1);
        let result = run_backward_pass(
            &mut workspaces,
            &basis_store,
            &StageContext {
                templates: &templates,
                base_rows: &base_rows,
                noise_scale: &[],
                n_hydros: 0,
                n_load_buses: 0,
                load_balance_row_starts: &[],
                load_bus_indices: &[],
                block_counts_per_stage: &[],
                ncs_max_gen: &[],
                discount_factors: &[],
                cumulative_discount_factors: &[],
            },
            &mut fcf,
            &mut empty_cut_batches(templates.len()),
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &[],
            },
            &mut BackwardPassSpec {
                records: &[],
                iteration: 0,
                local_work: exchange.local_count(),
                fwd_offset: 0,
                risk_measures: &risk_measures,
                exchange: &mut exchange,
                cut_activity_tolerance: 0.0,
                cut_sync_bufs: &mut csb,
                probabilities_buf: &mut Vec::new(),
                successor_active_slots_buf: &mut Vec::new(),
                visited_archive: None,
            },
            &comm,
        )
        .unwrap();

        // 3-stage system: cuts at stages 0 and 1; 2 trial points each.
        // Total cuts = 2 stages × 2 trial points = 4.
        assert_eq!(result.cuts_generated, 4);
        assert_eq!(fcf.active_cuts(0).count(), 2);
        assert_eq!(fcf.active_cuts(1).count(), 2);
        assert_eq!(fcf.active_cuts(2).count(), 0);
    }

    #[test]
    fn forward_pass_index_matches_global_scenario_index() {
        // Acceptance criterion: when a cut is generated for global trial point
        // m=5, then `fcf.add_cut(stage, iteration, 5, ...)` is called with
        // forward_pass_index = m = 5.
        //
        // Setup: iteration=2, forward_passes=6 (6 scenarios on 1 rank), 1 opening.
        // ExchangeBuffers: local_count=6, num_ranks=1, total_scenarios=6.
        // state_at(5/6, 5%6) = state_at(0, 5) — valid.
        //
        // Slot formula: slot = warm_start(0) + 2*6 + 5 = 17.
        // The key invariant: forward_pass_index = m = 5.
        let n_stages = 2_usize;
        let stochastic = make_stochastic_context(n_stages, 1);
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template_1_0(); n_stages];
        let base_rows = vec![1_usize; n_stages];

        let n_state = indexer.n_state;
        let forward_passes = 6_u32; // 6 scenarios on a single rank
        let mut fcf = FutureCostFunction::new(n_stages, n_state, forward_passes, 20, 0);

        // 6 trial points (m = 0..5). ExchangeBuffers: local_count=6, num_ranks=1.
        let mut exchange = exchange_with_states(
            n_state,
            vec![
                vec![1.0],
                vec![2.0],
                vec![3.0],
                vec![4.0],
                vec![5.0],
                vec![6.0],
            ],
        );

        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let risk_measures = vec![RiskMeasure::Expectation; n_stages];

        let solution = solution_1_0(50.0, 0.0);
        let solver = MockSolver::always_ok(solution);
        let comm = StubComm;
        let mut workspaces = single_workspace(solver, n_state);
        let basis_store = empty_basis_store(exchange.local_count(), n_stages);

        let mut csb = CutSyncBuffers::new(n_state, 64, 1);
        let _ = run_backward_pass(
            &mut workspaces,
            &basis_store,
            &StageContext {
                templates: &templates,
                base_rows: &base_rows,
                noise_scale: &[],
                n_hydros: 0,
                n_load_buses: 0,
                load_balance_row_starts: &[],
                load_bus_indices: &[],
                block_counts_per_stage: &[],
                ncs_max_gen: &[],
                discount_factors: &[],
                cumulative_discount_factors: &[],
            },
            &mut fcf,
            &mut empty_cut_batches(templates.len()),
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &[],
            },
            &mut BackwardPassSpec {
                records: &[],
                iteration: 2,
                local_work: exchange.local_count(),
                fwd_offset: 0,
                risk_measures: &risk_measures,
                exchange: &mut exchange,
                cut_activity_tolerance: 0.0,
                cut_sync_bufs: &mut csb,
                probabilities_buf: &mut Vec::new(),
                successor_active_slots_buf: &mut Vec::new(),
                visited_archive: None,
            },
            &comm,
        )
        .unwrap();

        // m=5: slot = warm_start(0) + 2*6 + 5 = 17
        // The critical check: forward_pass_index in metadata equals global m=5.
        let meta = &fcf.pools[0].metadata[17];
        assert_eq!(meta.iteration_generated, 2, "iteration_generated must be 2");
        assert_eq!(
            meta.forward_pass_index, 5,
            "forward_pass_index must be 5 (= global m)"
        );
    }

    // ── Unit tests: warm-start basis caching (backward pass) ──────────────────

    /// Warm-start from a pre-populated forward basis: when `BasisStore` has
    /// `Some(Basis)` at `(scenario=0, stage=1)` before the first backward call,
    /// the first opening at the successor stage must call `solve_with_basis`
    /// rather than `solve`.
    ///
    /// AC: Given a 2-stage system, 1 trial point, 1 opening, with
    /// `basis_store.get(0, 1) = Some(Basis::new(...))` pre-populated,
    /// then `solver.warm_start_calls == 1` after the backward pass.
    #[test]
    fn warm_start_uses_prepopulated_forward_basis() {
        let n_stages = 2_usize;
        let n_openings = 1_usize;
        let stochastic = make_stochastic_context(n_stages, n_openings);
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template_1_0(); n_stages];
        let base_rows = vec![1_usize; n_stages];

        let n_state = indexer.n_state;
        let forward_passes = 1_u32;
        let mut fcf = FutureCostFunction::new(n_stages, n_state, forward_passes, 10, 0);
        let mut exchange = exchange_with_states(n_state, vec![vec![10.0]]);

        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let risk_measures = vec![RiskMeasure::Expectation; n_stages];

        let solution = solution_1_0(100.0, -5.0);
        let solver = MockSolver::always_ok(solution);
        let comm = StubComm;

        // Pre-populate the basis store at (scenario=0, stage=1).
        // This simulates a forward pass having already solved stage 1 and cached its basis.
        let pre_basis = Basis::new(templates[1].num_cols, templates[1].num_rows);
        let mut workspaces = single_workspace(solver, n_state);
        let basis_store = basis_store_with_one(exchange.local_count(), n_stages, 0, 1, pre_basis);

        let mut csb = CutSyncBuffers::new(n_state, 64, 1);
        let _ = run_backward_pass(
            &mut workspaces,
            &basis_store,
            &StageContext {
                templates: &templates,
                base_rows: &base_rows,
                noise_scale: &[],
                n_hydros: 0,
                n_load_buses: 0,
                load_balance_row_starts: &[],
                load_bus_indices: &[],
                block_counts_per_stage: &[],
                ncs_max_gen: &[],
                discount_factors: &[],
                cumulative_discount_factors: &[],
            },
            &mut fcf,
            &mut empty_cut_batches(templates.len()),
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &[],
            },
            &mut BackwardPassSpec {
                records: &[],
                iteration: 0,
                local_work: exchange.local_count(),
                fwd_offset: 0,
                risk_measures: &risk_measures,
                exchange: &mut exchange,
                cut_activity_tolerance: 0.0,
                cut_sync_bufs: &mut csb,
                probabilities_buf: &mut Vec::new(),
                successor_active_slots_buf: &mut Vec::new(),
                visited_archive: None,
            },
            &comm,
        )
        .unwrap();

        let warm_start_calls = workspaces[0].solver.warm_start_calls;
        assert_eq!(
            warm_start_calls, 1,
            "first opening at successor stage must call solve_with_basis \
             when basis_store.get(0, 1) is pre-populated (warm_start_calls == 1, got {warm_start_calls})"
        );
    }

    /// Multi-opening P3b behavior: given 3 openings at the same successor stage,
    /// the first opening cold-starts (store slot is None via `solve()`), and
    /// openings 1 and 2 use `HiGHS` internal hot-start via `solve()` instead of
    /// `solve_with_basis(working_basis)`.
    ///
    /// AC: Given a 2-stage system, 1 trial point, 3 openings, empty basis cache,
    /// then `solver.warm_start_calls == 0` after the backward pass (P3b: no
    /// and 3 warm-start; opening 1 cold-starts).
    #[test]
    fn multi_opening_subsequent_openings_use_internal_hotstart() {
        let n_stages = 2_usize;
        let n_openings = 3_usize;
        let stochastic = make_stochastic_context(n_stages, n_openings);
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template_1_0(); n_stages];
        let base_rows = vec![1_usize; n_stages];

        let n_state = indexer.n_state;
        let forward_passes = 1_u32;
        let mut fcf = FutureCostFunction::new(n_stages, n_state, forward_passes, 10, 0);
        let mut exchange = exchange_with_states(n_state, vec![vec![10.0]]);

        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let risk_measures = vec![RiskMeasure::Expectation; n_stages];

        let solution = solution_1_0(100.0, -5.0);
        let solver = MockSolver::always_ok(solution);
        let comm = StubComm;

        // Start with an empty store — opening 1 must cold-start.
        let mut workspaces = single_workspace(solver, n_state);
        let basis_store = empty_basis_store(exchange.local_count(), n_stages);

        let mut csb = CutSyncBuffers::new(n_state, 64, 1);
        let _ = run_backward_pass(
            &mut workspaces,
            &basis_store,
            &StageContext {
                templates: &templates,
                base_rows: &base_rows,
                noise_scale: &[],
                n_hydros: 0,
                n_load_buses: 0,
                load_balance_row_starts: &[],
                load_bus_indices: &[],
                block_counts_per_stage: &[],
                ncs_max_gen: &[],
                discount_factors: &[],
                cumulative_discount_factors: &[],
            },
            &mut fcf,
            &mut empty_cut_batches(templates.len()),
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &[],
            },
            &mut BackwardPassSpec {
                records: &[],
                iteration: 0,
                local_work: exchange.local_count(),
                fwd_offset: 0,
                risk_measures: &risk_measures,
                exchange: &mut exchange,
                cut_activity_tolerance: 0.0,
                cut_sync_bufs: &mut csb,
                probabilities_buf: &mut Vec::new(),
                successor_active_slots_buf: &mut Vec::new(),
                visited_archive: None,
            },
            &comm,
        )
        .unwrap();

        // P3b optimization: opening 0 cold-starts (no basis in store),
        // openings 1 and 2 use solve() (HiGHS internal hot-start) instead of
        // solve_with_basis. No explicit warm-start calls for subsequent openings.
        let warm_start_calls = workspaces[0].solver.warm_start_calls;
        assert_eq!(
            warm_start_calls, 0,
            "P3b: no solve_with_basis calls expected when BasisStore is empty \
             (warm_start_calls == 0, got {warm_start_calls})"
        );
    }

    /// Error propagation: when a backward solve returns `SolverError::Infeasible`,
    /// the error must propagate as `SddpError::Infeasible`.
    ///
    /// In the new per-scenario design, the backward pass uses a local `working_basis`
    /// variable (not written back to `BasisStore`), so there is no shared cache slot
    /// to check after the error. The test verifies that the error is correctly
    /// propagated regardless of what was in the basis store at entry.
    ///
    /// AC: Given a 2-stage system, 1 opening, `MockSolver` returns infeasible on
    /// call 0, then `run_backward_pass` returns `Err(SddpError::Infeasible { .. })`.
    #[test]
    fn backward_solver_error_propagates() {
        let n_stages = 2_usize;
        let n_openings = 1_usize;
        let stochastic = make_stochastic_context(n_stages, n_openings);
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template_1_0(); n_stages];
        let base_rows = vec![1_usize; n_stages];

        let n_state = indexer.n_state;
        let forward_passes = 1_u32;
        let mut fcf = FutureCostFunction::new(n_stages, n_state, forward_passes, 10, 0);
        let mut exchange = exchange_with_states(n_state, vec![vec![10.0]]);

        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let risk_measures = vec![RiskMeasure::Expectation; n_stages];

        let solution = solution_1_0(0.0, 0.0);
        // The first backward solve (call 0) returns infeasible.
        let solver = MockSolver::infeasible_on(solution, 0);
        let comm = StubComm;

        // Pre-populate the store — error should propagate regardless.
        let pre_basis = Basis::new(templates[1].num_cols, templates[1].num_rows);
        let mut workspaces = single_workspace(solver, n_state);
        let basis_store = basis_store_with_one(exchange.local_count(), n_stages, 0, 1, pre_basis);

        let mut csb = CutSyncBuffers::new(n_state, 64, 1);
        let result = run_backward_pass(
            &mut workspaces,
            &basis_store,
            &StageContext {
                templates: &templates,
                base_rows: &base_rows,
                noise_scale: &[],
                n_hydros: 0,
                n_load_buses: 0,
                load_balance_row_starts: &[],
                load_bus_indices: &[],
                block_counts_per_stage: &[],
                ncs_max_gen: &[],
                discount_factors: &[],
                cumulative_discount_factors: &[],
            },
            &mut fcf,
            &mut empty_cut_batches(templates.len()),
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &[],
            },
            &mut BackwardPassSpec {
                records: &[],
                iteration: 0,
                local_work: exchange.local_count(),
                fwd_offset: 0,
                risk_measures: &risk_measures,
                exchange: &mut exchange,
                cut_activity_tolerance: 0.0,
                cut_sync_bufs: &mut csb,
                probabilities_buf: &mut Vec::new(),
                successor_active_slots_buf: &mut Vec::new(),
                visited_archive: None,
            },
            &comm,
        );

        assert!(
            matches!(result, Err(crate::SddpError::Infeasible { .. })),
            "expected SddpError::Infeasible, got: {result:?}",
        );
        // The BasisStore is not mutated by the backward pass — the working_basis
        // is a local variable dropped on error. The store slot at (0, 1) remains
        // as it was; this just verifies the store is untouched by an error path.
        assert!(
            basis_store.get(0, 1).is_some(),
            "BasisStore must not be mutated by the backward pass error path"
        );
    }

    // ── New test: parallel cut determinism ────────────────────────────────────

    /// AC: When `run_backward_pass` runs with 1 workspace vs 4 workspaces given
    /// the same input data, the FCF pools contain identical cuts (same intercept,
    /// coefficient vectors, and slot assignments for each trial point).
    #[test]
    #[allow(
        clippy::too_many_lines,
        clippy::cast_possible_truncation,
        clippy::cast_precision_loss
    )]
    fn test_backward_pass_parallel_cut_determinism() {
        use crate::lp_builder::PatchBuffer;

        let n_stages = 3_usize;
        let n_openings = 2_usize;
        let n_trial_points = 8_usize;

        let stochastic = make_stochastic_context(n_stages, n_openings);
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template_1_0(); n_stages];
        let base_rows = vec![1_usize; n_stages];

        let n_state = indexer.n_state;
        #[allow(clippy::cast_possible_truncation)]
        let forward_passes = n_trial_points as u32;

        // Build 8 distinct trial-point states.
        let states: Vec<Vec<f64>> = (0..n_trial_points).map(|i| vec![i as f64 + 1.0]).collect();
        let mut exchange = exchange_with_states(n_state, states);

        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let risk_measures = vec![RiskMeasure::Expectation; n_stages];
        let solution = solution_1_0(100.0, -5.0);
        let comm = StubComm;

        // --- Run with 1 workspace ---
        let mut fcf_1 = FutureCostFunction::new(n_stages, n_state, forward_passes, 20, 0);
        let solver_1 = MockSolver::always_ok(solution.clone());
        let mut workspaces_1 = vec![SolverWorkspace {
            solver: solver_1,
            patch_buf: PatchBuffer::new(1, 0, 0, 0),
            current_state: Vec::with_capacity(n_state),
            scratch: crate::workspace::ScratchBuffers {
                noise_buf: Vec::new(),
                inflow_m3s_buf: Vec::new(),
                lag_matrix_buf: Vec::new(),
                par_inflow_buf: Vec::new(),
                eta_floor_buf: Vec::new(),
                zero_targets_buf: Vec::new(),
                ncs_col_upper_buf: Vec::new(),
                ncs_col_lower_buf: Vec::new(),
                ncs_col_indices_buf: Vec::new(),
                load_rhs_buf: Vec::new(),
                row_lower_buf: Vec::new(),
                z_inflow_rhs_buf: Vec::new(),
                effective_eta_buf: Vec::new(),
                unscaled_primal: Vec::new(),
                unscaled_dual: Vec::new(),
            },
        }];
        let basis_store_1 = empty_basis_store(exchange.local_count(), n_stages);
        let ctx = StageContext {
            templates: &templates,
            base_rows: &base_rows,
            noise_scale: &[],
            n_hydros: 0,
            n_load_buses: 0,
            load_balance_row_starts: &[],
            load_bus_indices: &[],
            block_counts_per_stage: &[],
            ncs_max_gen: &[],
            discount_factors: &[],
            cumulative_discount_factors: &[],
        };
        let mut csb = CutSyncBuffers::new(n_state, 64, 1);
        let _ = run_backward_pass(
            &mut workspaces_1,
            &basis_store_1,
            &ctx,
            &mut fcf_1,
            &mut empty_cut_batches(n_stages),
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &[],
            },
            &mut BackwardPassSpec {
                records: &[],
                iteration: 0,
                local_work: exchange.local_count(),
                fwd_offset: 0,
                risk_measures: &risk_measures,
                exchange: &mut exchange,
                cut_activity_tolerance: 0.0,
                cut_sync_bufs: &mut csb,
                probabilities_buf: &mut Vec::new(),
                successor_active_slots_buf: &mut Vec::new(),
                visited_archive: None,
            },
            &comm,
        )
        .unwrap();

        // --- Run with 4 workspaces ---
        let mut fcf_4 = FutureCostFunction::new(n_stages, n_state, forward_passes, 20, 0);
        let mut workspaces_4: Vec<SolverWorkspace<MockSolver>> = (0..4)
            .map(|_| SolverWorkspace {
                solver: MockSolver::always_ok(solution.clone()),
                patch_buf: PatchBuffer::new(1, 0, 0, 0),
                current_state: Vec::with_capacity(n_state),
                scratch: crate::workspace::ScratchBuffers {
                    noise_buf: Vec::new(),
                    inflow_m3s_buf: Vec::new(),
                    lag_matrix_buf: Vec::new(),
                    par_inflow_buf: Vec::new(),
                    eta_floor_buf: Vec::new(),
                    zero_targets_buf: Vec::new(),
                    ncs_col_upper_buf: Vec::new(),
                    ncs_col_lower_buf: Vec::new(),
                    ncs_col_indices_buf: Vec::new(),
                    load_rhs_buf: Vec::new(),
                    row_lower_buf: Vec::new(),
                    z_inflow_rhs_buf: Vec::new(),
                    effective_eta_buf: Vec::new(),
                    unscaled_primal: Vec::new(),
                    unscaled_dual: Vec::new(),
                },
            })
            .collect();
        let basis_store_4 = empty_basis_store(exchange.local_count(), n_stages);
        let mut csb = CutSyncBuffers::new(n_state, 64, 1);
        let _ = run_backward_pass(
            &mut workspaces_4,
            &basis_store_4,
            &ctx,
            &mut fcf_4,
            &mut empty_cut_batches(n_stages),
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &[],
            },
            &mut BackwardPassSpec {
                records: &[],
                iteration: 0,
                local_work: exchange.local_count(),
                fwd_offset: 0,
                risk_measures: &risk_measures,
                exchange: &mut exchange,
                cut_activity_tolerance: 0.0,
                cut_sync_bufs: &mut csb,
                probabilities_buf: &mut Vec::new(),
                successor_active_slots_buf: &mut Vec::new(),
                visited_archive: None,
            },
            &comm,
        )
        .unwrap();

        // --- Verify identical FCF contents for all non-last stages ---
        for t in 0..(n_stages - 1) {
            let cuts_1: Vec<_> = fcf_1.active_cuts(t).collect();
            let cuts_4: Vec<_> = fcf_4.active_cuts(t).collect();

            assert_eq!(
                cuts_1.len(),
                cuts_4.len(),
                "stage {t}: cut count differs (1 workspace: {}, 4 workspaces: {})",
                cuts_1.len(),
                cuts_4.len()
            );

            for (idx, ((slot_1, intercept_1, coeff_1), (slot_4, intercept_4, coeff_4))) in
                cuts_1.iter().zip(cuts_4.iter()).enumerate()
            {
                assert_eq!(
                    slot_1, slot_4,
                    "stage {t} cut {idx}: slot mismatch ({slot_1} vs {slot_4})"
                );
                assert!(
                    (intercept_1 - intercept_4).abs() < 1e-12,
                    "stage {t} cut {idx}: intercept mismatch ({intercept_1} vs {intercept_4})"
                );
                assert_eq!(
                    coeff_1.len(),
                    coeff_4.len(),
                    "stage {t} cut {idx}: coefficient vector length mismatch"
                );
                for (j, (c1, c4)) in coeff_1.iter().zip(coeff_4.iter()).enumerate() {
                    assert!(
                        (c1 - c4).abs() < 1e-12,
                        "stage {t} cut {idx} coeff[{j}]: {c1} vs {c4}"
                    );
                }
            }
        }

        // Last stage must have no cuts in both.
        assert_eq!(fcf_1.active_cuts(n_stages - 1).count(), 0);
        assert_eq!(fcf_4.active_cuts(n_stages - 1).count(), 0);
    }

    // ── Load noise wiring tests (backward pass) ──────────────────────────────

    /// Build a 2-stage `StochasticContext` with 1 hydro and 1 stochastic load bus.
    ///
    /// The noise vector dimension is `n_hydros + n_load_buses = 2`.
    /// Stage 0 uses `branching_factor` openings; stage 1 is the successor solved
    /// in the backward pass opening loop.
    #[allow(clippy::too_many_lines)]
    fn make_stochastic_context_with_load(
        n_stages: usize,
        branching_factor: usize,
        mean_mw: f64,
        std_mw: f64,
    ) -> cobre_stochastic::StochasticContext {
        use chrono::NaiveDate;
        use cobre_core::entities::hydro::{Hydro, HydroGenerationModel, HydroPenalties};
        use cobre_core::scenario::{CorrelationModel, InflowModel, LoadModel};
        use cobre_core::temporal::{
            Block, BlockMode, NoiseMethod, ScenarioSourceConfig, Stage, StageRiskConfig,
            StageStateConfig,
        };
        use cobre_core::{Bus, DeficitSegment, EntityId, SystemBuilder};
        use cobre_stochastic::context::build_stochastic_context;

        let bus0 = Bus {
            id: EntityId(0),
            name: "B0".to_string(),
            deficit_segments: vec![DeficitSegment {
                depth_mw: None,
                cost_per_mwh: 1000.0,
            }],
            excess_cost: 0.0,
        };
        let bus1 = Bus {
            id: EntityId(1),
            name: "B1".to_string(),
            deficit_segments: vec![DeficitSegment {
                depth_mw: None,
                cost_per_mwh: 1000.0,
            }],
            excess_cost: 0.0,
        };
        let hydro = Hydro {
            id: EntityId(10),
            name: "H10".to_string(),
            bus_id: EntityId(0),
            downstream_id: None,
            entry_stage_id: None,
            exit_stage_id: None,
            min_storage_hm3: 0.0,
            max_storage_hm3: 100.0,
            min_outflow_m3s: 0.0,
            max_outflow_m3s: None,
            generation_model: HydroGenerationModel::ConstantProductivity {
                productivity_mw_per_m3s: 1.0,
            },
            min_turbined_m3s: 0.0,
            max_turbined_m3s: 100.0,
            min_generation_mw: 0.0,
            max_generation_mw: 100.0,
            tailrace: None,
            hydraulic_losses: None,
            efficiency: None,
            evaporation_coefficients_mm: None,
            evaporation_reference_volumes_hm3: None,
            diversion: None,
            filling: None,
            penalties: HydroPenalties {
                spillage_cost: 0.0,
                diversion_cost: 0.0,
                fpha_turbined_cost: 0.0,
                storage_violation_below_cost: 0.0,
                filling_target_violation_cost: 0.0,
                turbined_violation_below_cost: 0.0,
                outflow_violation_below_cost: 0.0,
                outflow_violation_above_cost: 0.0,
                generation_violation_below_cost: 0.0,
                evaporation_violation_cost: 0.0,
                water_withdrawal_violation_cost: 0.0,
                water_withdrawal_violation_pos_cost: 0.0,
                water_withdrawal_violation_neg_cost: 0.0,
                evaporation_violation_pos_cost: 0.0,
                evaporation_violation_neg_cost: 0.0,
                inflow_nonnegativity_cost: 1000.0,
            },
        };

        #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
        let make_stage = |idx: usize| Stage {
            index: idx,
            id: idx as i32,
            start_date: NaiveDate::from_ymd_opt(2024, 1, 1).unwrap(),
            end_date: NaiveDate::from_ymd_opt(2024, 2, 1).unwrap(),
            season_id: Some(0),
            blocks: vec![Block {
                index: 0,
                name: "S".to_string(),
                duration_hours: 744.0,
            }],
            block_mode: BlockMode::Parallel,
            state_config: StageStateConfig {
                storage: true,
                inflow_lags: false,
            },
            risk_config: StageRiskConfig::Expectation,
            scenario_config: ScenarioSourceConfig {
                branching_factor,
                noise_method: NoiseMethod::Saa,
            },
        };

        let stages: Vec<Stage> = (0..n_stages).map(make_stage).collect();

        #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
        let inflow_models: Vec<InflowModel> = (0..n_stages)
            .map(|idx| InflowModel {
                hydro_id: EntityId(10),
                stage_id: idx as i32,
                mean_m3s: 100.0,
                std_m3s: 30.0,
                ar_coefficients: vec![],
                residual_std_ratio: 1.0,
            })
            .collect();

        #[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
        let load_models: Vec<LoadModel> = (0..n_stages)
            .map(|idx| LoadModel {
                bus_id: EntityId(1),
                stage_id: idx as i32,
                mean_mw,
                std_mw,
            })
            .collect();

        let correlation = CorrelationModel {
            method: "cholesky".to_string(),
            profiles: std::collections::BTreeMap::new(),
            schedule: vec![],
        };

        let system = SystemBuilder::new()
            .buses(vec![bus0, bus1])
            .hydros(vec![hydro])
            .stages(stages)
            .inflow_models(inflow_models)
            .load_models(load_models)
            .correlation(correlation)
            .build()
            .unwrap();

        build_stochastic_context(&system, 42, &[], &[], None).unwrap()
    }

    /// AC: Given a backward pass with 1 stochastic load bus and opening noise
    /// that includes a load component eta, the load balance row RHS in the patch
    /// buffer is set to `max(0, mean + std * eta) * block_factor` before the solve.
    ///
    /// We verify this indirectly: after the backward pass runs, `ws.scratch.load_rhs_buf`
    /// must be non-empty and must contain a positive value (with mean=300, std=30
    /// any reasonable eta produces a positive realization).
    #[test]
    #[allow(clippy::too_many_lines)]
    fn backward_pass_load_patches_applied() {
        // 2-stage system: backward pass solves at successor=1 for each opening.
        // n_hydros=1, n_load_buses=1, 1 block per stage.
        let n_stages = 2_usize;
        let n_openings = 2_usize;
        // mean_mw=300 guarantees a positive realization for any reasonable eta draw.
        let stochastic = make_stochastic_context_with_load(n_stages, n_openings, 300.0, 30.0);
        let indexer = StageIndexer::new(1, 0); // N=1, L=0, n_state=1

        // PatchBuffer: n_hydros=1, max_par_order=0, n_load_buses=1, max_blocks=1.
        let patch_buf = crate::lp_builder::PatchBuffer::new(1, 0, 1, 1);

        // Template: 2 rows (row 0 = state-fixing, row 1 = water-balance).
        // base_rows=[1] → inflow RHS row starts at index 1.
        // noise_scale=[1.0, 1.0] (one per (stage, hydro)).
        let template = StageTemplate {
            num_cols: 3,
            num_rows: 2,
            num_nz: 1,
            col_starts: vec![0_i32, 0, 1, 1],
            row_indices: vec![0_i32],
            values: vec![1.0],
            col_lower: vec![0.0, 0.0, 0.0],
            col_upper: vec![f64::INFINITY, f64::INFINITY, f64::INFINITY],
            objective: vec![0.0, 0.0, 1.0],
            row_lower: vec![50.0, 100.0],
            row_upper: vec![50.0, 100.0],
            n_state: 1,
            n_transfer: 0,
            n_dual_relevant: 1,
            n_hydro: 1,
            max_par_order: 0,
            col_scale: Vec::new(),
            row_scale: Vec::new(),
        };
        let templates = vec![template; n_stages];
        let base_rows = vec![1_usize; n_stages];
        let noise_scale = vec![1.0_f64; n_stages]; // one per (stage, hydro)

        let n_state = indexer.n_state;
        let forward_passes = 1_u32;
        let mut fcf = FutureCostFunction::new(n_stages, n_state, forward_passes, 10, 0);
        let mut exchange = exchange_with_states(n_state, vec![vec![10.0]]);

        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let risk_measures = vec![RiskMeasure::Expectation; n_stages];

        // MockSolver returns a fixed solution (1 state var, 1 dual entry).
        let solution = solution_1_0(100.0, -2.0);

        let ws = SolverWorkspace {
            solver: MockSolver::always_ok(solution),
            patch_buf,
            current_state: Vec::with_capacity(n_state),
            scratch: crate::workspace::ScratchBuffers {
                noise_buf: Vec::new(),
                inflow_m3s_buf: Vec::new(),
                lag_matrix_buf: Vec::new(),
                par_inflow_buf: Vec::new(),
                eta_floor_buf: Vec::new(),
                zero_targets_buf: Vec::new(),
                ncs_col_upper_buf: Vec::new(),
                ncs_col_lower_buf: Vec::new(),
                ncs_col_indices_buf: Vec::new(),
                load_rhs_buf: Vec::with_capacity(1),
                row_lower_buf: Vec::new(),
                z_inflow_rhs_buf: Vec::new(),
                effective_eta_buf: Vec::new(),
                unscaled_primal: Vec::new(),
                unscaled_dual: Vec::new(),
            },
        };
        let mut workspaces = vec![ws];

        let comm = StubComm;
        let basis_store = empty_basis_store(exchange.local_count(), n_stages);

        // load_balance_row_starts[successor=1]=10; load_bus_indices=[0]; 1 block/stage.
        let load_balance_row_starts = vec![10_usize; n_stages];
        let load_bus_indices = vec![0_usize];
        let block_counts_per_stage = vec![1_usize; n_stages];

        let mut csb = CutSyncBuffers::new(n_state, 64, 1);
        let _ = run_backward_pass(
            &mut workspaces,
            &basis_store,
            &StageContext {
                templates: &templates,
                base_rows: &base_rows,
                noise_scale: &noise_scale,
                n_hydros: 1,
                n_load_buses: 1,
                load_balance_row_starts: &load_balance_row_starts,
                load_bus_indices: &load_bus_indices,
                block_counts_per_stage: &block_counts_per_stage,
                ncs_max_gen: &[],
                discount_factors: &[],
                cumulative_discount_factors: &[],
            },
            &mut fcf,
            &mut empty_cut_batches(templates.len()),
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &[],
            },
            &mut BackwardPassSpec {
                records: &[],
                iteration: 0,
                local_work: exchange.local_count(),
                fwd_offset: 0,
                risk_measures: &risk_measures,
                exchange: &mut exchange,
                cut_activity_tolerance: 0.0,
                cut_sync_bufs: &mut csb,
                probabilities_buf: &mut Vec::new(),
                successor_active_slots_buf: &mut Vec::new(),
                visited_archive: None,
            },
            &comm,
        )
        .unwrap();

        // After the backward pass, load_rhs_buf must have been populated with a
        // positive value for the last opening solved (mean=300, std=30 → positive).
        assert_eq!(
            workspaces[0].scratch.load_rhs_buf.len(),
            1,
            "load_rhs_buf must have 1 entry (1 load bus × 1 block)"
        );
        assert!(
            workspaces[0].scratch.load_rhs_buf[0] > 0.0,
            "load realization must be positive with mean=300, std=30: got {}",
            workspaces[0].scratch.load_rhs_buf[0]
        );
    }

    /// AC: Given a backward pass with 0 stochastic load buses, `patch_count`
    /// equals `N*(2+L)` (no load patches) and `load_rhs_buf` stays empty.
    ///
    /// N=1, L=0 → `N*(2+L) = 2`.
    #[test]
    #[allow(clippy::too_many_lines)]
    fn backward_pass_no_load_buses_unchanged() {
        let n_stages = 2_usize;
        let n_openings = 2_usize;
        let stochastic = make_stochastic_context(n_stages, n_openings);
        let indexer = StageIndexer::new(1, 0); // N=1, L=0

        // PatchBuffer with no load buses: n_load_buses=0, max_blocks=1.
        let patch_buf = crate::lp_builder::PatchBuffer::new(1, 0, 0, 0);

        let template = StageTemplate {
            num_cols: 3,
            num_rows: 2,
            num_nz: 1,
            col_starts: vec![0_i32, 0, 1, 1],
            row_indices: vec![0_i32],
            values: vec![1.0],
            col_lower: vec![0.0, 0.0, 0.0],
            col_upper: vec![f64::INFINITY, f64::INFINITY, f64::INFINITY],
            objective: vec![0.0, 0.0, 1.0],
            row_lower: vec![50.0, 100.0],
            row_upper: vec![50.0, 100.0],
            n_state: 1,
            n_transfer: 0,
            n_dual_relevant: 1,
            n_hydro: 1,
            max_par_order: 0,
            col_scale: Vec::new(),
            row_scale: Vec::new(),
        };
        let templates = vec![template; n_stages];
        let base_rows = vec![1_usize; n_stages];
        let noise_scale = vec![1.0_f64; n_stages];

        let n_state = indexer.n_state;
        let forward_passes = 1_u32;
        let mut fcf = FutureCostFunction::new(n_stages, n_state, forward_passes, 10, 0);
        let mut exchange = exchange_with_states(n_state, vec![vec![10.0]]);

        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let risk_measures = vec![RiskMeasure::Expectation; n_stages];

        let solution = solution_1_0(100.0, -2.0);
        let ws = SolverWorkspace {
            solver: MockSolver::always_ok(solution),
            patch_buf,
            current_state: Vec::with_capacity(n_state),
            scratch: crate::workspace::ScratchBuffers {
                noise_buf: Vec::new(),
                inflow_m3s_buf: Vec::new(),
                lag_matrix_buf: Vec::new(),
                par_inflow_buf: Vec::new(),
                eta_floor_buf: Vec::new(),
                zero_targets_buf: Vec::new(),
                ncs_col_upper_buf: Vec::new(),
                ncs_col_lower_buf: Vec::new(),
                ncs_col_indices_buf: Vec::new(),
                load_rhs_buf: Vec::new(),
                row_lower_buf: Vec::new(),
                z_inflow_rhs_buf: Vec::new(),
                effective_eta_buf: Vec::new(),
                unscaled_primal: Vec::new(),
                unscaled_dual: Vec::new(),
            },
        };
        let mut workspaces = vec![ws];
        let comm = StubComm;
        let basis_store = empty_basis_store(exchange.local_count(), n_stages);

        let mut csb = CutSyncBuffers::new(n_state, 64, 1);
        let _ = run_backward_pass(
            &mut workspaces,
            &basis_store,
            &StageContext {
                templates: &templates,
                base_rows: &base_rows,
                noise_scale: &noise_scale,
                n_hydros: 1,
                n_load_buses: 0,
                load_balance_row_starts: &[],
                load_bus_indices: &[],
                block_counts_per_stage: &[1_usize; 2],
                ncs_max_gen: &[],
                discount_factors: &[],
                cumulative_discount_factors: &[],
            },
            &mut fcf,
            &mut empty_cut_batches(templates.len()),
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &[],
            },
            &mut BackwardPassSpec {
                records: &[],
                iteration: 0,
                local_work: exchange.local_count(),
                fwd_offset: 0,
                risk_measures: &risk_measures,
                exchange: &mut exchange,
                cut_activity_tolerance: 0.0,
                cut_sync_bufs: &mut csb,
                probabilities_buf: &mut Vec::new(),
                successor_active_slots_buf: &mut Vec::new(),
                visited_archive: None,
            },
            &comm,
        )
        .unwrap();

        // With n_load_buses=0, forward_patch_count = N*(2+L) + z_inflow = 1*(2+0)+1 = 3.
        assert_eq!(
            workspaces[0].patch_buf.forward_patch_count(),
            3,
            "forward_patch_count must be N*(2+L)+N=3 when n_load_buses=0, got {}",
            workspaces[0].patch_buf.forward_patch_count()
        );
        // load_rhs_buf must remain empty.
        assert!(
            workspaces[0].scratch.load_rhs_buf.is_empty(),
            "load_rhs_buf must be empty when n_load_buses=0"
        );
    }

    /// AC: Given a backward pass with stochastic load, when Benders cut
    /// coefficients are extracted, the cut coefficient array has length `n_state`
    /// unchanged — load adds no state variables.
    ///
    /// Setup: N=1 hydro, L=0 PAR lags → `n_state=1`. After the backward pass with
    /// 1 load bus, each generated cut must have exactly 1 coefficient.
    #[test]
    #[allow(clippy::too_many_lines)]
    fn backward_pass_cut_coefficients_unaffected() {
        let n_stages = 2_usize;
        let n_openings = 2_usize;
        let stochastic = make_stochastic_context_with_load(n_stages, n_openings, 200.0, 20.0);
        let indexer = StageIndexer::new(1, 0); // N=1, L=0, n_state=1

        let patch_buf = crate::lp_builder::PatchBuffer::new(1, 0, 1, 1);

        let template = StageTemplate {
            num_cols: 3,
            num_rows: 2,
            num_nz: 1,
            col_starts: vec![0_i32, 0, 1, 1],
            row_indices: vec![0_i32],
            values: vec![1.0],
            col_lower: vec![0.0, 0.0, 0.0],
            col_upper: vec![f64::INFINITY, f64::INFINITY, f64::INFINITY],
            objective: vec![0.0, 0.0, 1.0],
            row_lower: vec![50.0, 100.0],
            row_upper: vec![50.0, 100.0],
            n_state: 1,
            n_transfer: 0,
            n_dual_relevant: 1,
            n_hydro: 1,
            max_par_order: 0,
            col_scale: Vec::new(),
            row_scale: Vec::new(),
        };
        let templates = vec![template; n_stages];
        let base_rows = vec![1_usize; n_stages];
        let noise_scale = vec![1.0_f64; n_stages];

        let n_state = indexer.n_state; // 1
        let forward_passes = 1_u32;
        let mut fcf = FutureCostFunction::new(n_stages, n_state, forward_passes, 10, 0);
        let mut exchange = exchange_with_states(n_state, vec![vec![10.0]]);

        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let risk_measures = vec![RiskMeasure::Expectation; n_stages];

        let solution = solution_1_0(80.0, -3.0);
        let ws = SolverWorkspace {
            solver: MockSolver::always_ok(solution),
            patch_buf,
            current_state: Vec::with_capacity(n_state),
            scratch: crate::workspace::ScratchBuffers {
                noise_buf: Vec::new(),
                inflow_m3s_buf: Vec::new(),
                lag_matrix_buf: Vec::new(),
                par_inflow_buf: Vec::new(),
                eta_floor_buf: Vec::new(),
                zero_targets_buf: Vec::new(),
                ncs_col_upper_buf: Vec::new(),
                ncs_col_lower_buf: Vec::new(),
                ncs_col_indices_buf: Vec::new(),
                load_rhs_buf: Vec::with_capacity(1),
                row_lower_buf: Vec::new(),
                z_inflow_rhs_buf: Vec::new(),
                effective_eta_buf: Vec::new(),
                unscaled_primal: Vec::new(),
                unscaled_dual: Vec::new(),
            },
        };
        let mut workspaces = vec![ws];
        let comm = StubComm;
        let basis_store = empty_basis_store(exchange.local_count(), n_stages);

        let load_balance_row_starts = vec![10_usize; n_stages];
        let load_bus_indices = vec![0_usize];
        let block_counts_per_stage = vec![1_usize; n_stages];

        let mut csb = CutSyncBuffers::new(n_state, 64, 1);
        let result = run_backward_pass(
            &mut workspaces,
            &basis_store,
            &StageContext {
                templates: &templates,
                base_rows: &base_rows,
                noise_scale: &noise_scale,
                n_hydros: 1,
                n_load_buses: 1,
                load_balance_row_starts: &load_balance_row_starts,
                load_bus_indices: &load_bus_indices,
                block_counts_per_stage: &block_counts_per_stage,
                ncs_max_gen: &[],
                discount_factors: &[],
                cumulative_discount_factors: &[],
            },
            &mut fcf,
            &mut empty_cut_batches(templates.len()),
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &[],
            },
            &mut BackwardPassSpec {
                records: &[],
                iteration: 0,
                local_work: exchange.local_count(),
                fwd_offset: 0,
                risk_measures: &risk_measures,
                exchange: &mut exchange,
                cut_activity_tolerance: 0.0,
                cut_sync_bufs: &mut csb,
                probabilities_buf: &mut Vec::new(),
                successor_active_slots_buf: &mut Vec::new(),
                visited_archive: None,
            },
            &comm,
        )
        .unwrap();

        // Exactly 1 cut generated (1 trial point × 1 stage with a successor).
        assert_eq!(result.cuts_generated, 1);

        // The cut must have exactly n_state=1 coefficient.
        let cuts: Vec<_> = fcf.active_cuts(0).collect();
        assert_eq!(cuts.len(), 1);
        let (_, _intercept, coefficients) = &cuts[0];
        assert_eq!(
            coefficients.len(),
            n_state,
            "cut coefficients length must be n_state={n_state}, got {} — \
             load buses must not add state variables",
            coefficients.len()
        );
    }

    /// BUG-1 structural invariant: per-stage cut sync inside the backward loop.
    ///
    /// Verifies that after `run_backward_pass`, the cut synchronization has been
    /// performed per-stage (not as a separate post-sweep loop). The structural
    /// evidence is:
    ///
    /// 1. `BackwardResult.cut_sync_time_ms` is populated (timing was captured).
    /// 2. The FCF has the expected number of cuts per stage — same as single-rank
    ///    without sync, because single-rank sync is a no-op that does not change
    ///    results but exercises the code path.
    /// 3. Using `LocalBackend` (the production single-rank communicator) instead
    ///    of `StubComm` exercises the full `sync_cuts` → allgatherv → deserialize
    ///    path, confirming no panics or data corruption.
    ///
    /// True multi-rank correctness testing requires actual MPI and is out of
    /// scope for CI. This test validates the structural invariant (sync is
    /// called per-stage inside the loop) and exercises the full code path.
    #[test]
    fn per_stage_cut_sync_invariant_after_bug1_fix() {
        use cobre_comm::LocalBackend;

        let n_stages = 4_usize;
        let n_openings = 2_usize;
        let stochastic = make_stochastic_context(n_stages, n_openings);
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template_1_0(); n_stages];
        let base_rows = vec![1_usize; n_stages];

        let n_state = indexer.n_state;
        let forward_passes = 3_u32;
        let mut fcf = FutureCostFunction::new(n_stages, n_state, forward_passes, 20, 0);
        let mut exchange = exchange_with_states(n_state, vec![vec![10.0], vec![20.0], vec![30.0]]);

        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let risk_measures = vec![RiskMeasure::Expectation; n_stages];

        let solution = solution_1_0(100.0, -5.0);
        let solver = MockSolver::always_ok(solution);
        let comm = LocalBackend;
        let mut workspaces = single_workspace(solver, n_state);
        let basis_store = empty_basis_store(exchange.local_count(), n_stages);

        let mut csb = CutSyncBuffers::new(n_state, forward_passes as usize, 1);
        let result = run_backward_pass(
            &mut workspaces,
            &basis_store,
            &StageContext {
                templates: &templates,
                base_rows: &base_rows,
                noise_scale: &[],
                n_hydros: 0,
                n_load_buses: 0,
                load_balance_row_starts: &[],
                load_bus_indices: &[],
                block_counts_per_stage: &[],
                ncs_max_gen: &[],
                discount_factors: &[],
                cumulative_discount_factors: &[],
            },
            &mut fcf,
            &mut empty_cut_batches(templates.len()),
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &[],
            },
            &mut BackwardPassSpec {
                records: &[],
                iteration: 1,
                local_work: exchange.local_count(),
                fwd_offset: 0,
                risk_measures: &risk_measures,
                exchange: &mut exchange,
                cut_activity_tolerance: 0.0,
                cut_sync_bufs: &mut csb,
                probabilities_buf: &mut Vec::new(),
                successor_active_slots_buf: &mut Vec::new(),
                visited_archive: None,
            },
            &comm,
        )
        .unwrap();

        // 4-stage system: cuts at stages 0, 1, 2; 3 trial points each.
        // Total cuts = 3 stages × 3 trial points = 9.
        assert_eq!(result.cuts_generated, 9);

        // Each non-terminal stage has 3 cuts (one per trial point).
        assert_eq!(fcf.active_cuts(0).count(), 3, "stage 0 must have 3 cuts");
        assert_eq!(fcf.active_cuts(1).count(), 3, "stage 1 must have 3 cuts");
        assert_eq!(fcf.active_cuts(2).count(), 3, "stage 2 must have 3 cuts");
        assert_eq!(
            fcf.active_cuts(3).count(),
            0,
            "terminal stage must have 0 cuts"
        );

        // Verify cut_sync_time_ms was captured (structural evidence that
        // sync_cuts was called inside the backward loop).
        // For single-rank LocalBackend, sync is a no-op, so time should be
        // very small but the field must be populated (not default/garbage).
        // We just verify it's a valid non-negative value.
        assert!(
            result.cut_sync_time_ms < 10_000,
            "cut_sync_time_ms should be reasonable, got {}",
            result.cut_sync_time_ms
        );
    }
}