//! 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]                       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).
//!
//! ## 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::{Basis, RowBatch, SolverError, SolverInterface};
use rayon::iter::{IndexedParallelIterator, IntoParallelRefMutIterator, ParallelIterator};

use crate::{
    FutureCostFunction, SddpError, TrajectoryRecord,
    context::{StageContext, TrainingContext},
    forward::{build_cut_row_batch, partition},
    noise::{transform_inflow_noise, transform_load_noise},
    risk_measure::BackwardOutcome,
    risk_measure::RiskMeasure,
    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-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],
}

/// 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,
}

/// Per-trial-point mutable accumulators reused across openings.
struct TrialAccumulators {
    /// Collected per-opening outcomes (intercept + coefficients + objective).
    outcomes: Vec<BackwardOutcome>,
    /// Binding cut slot increments keyed by slot index.
    slot_increments: std::collections::HashMap<usize, u64>,
}

/// Apply noise and state patches for one opening of the backward pass.
///
/// Loads the stage `s` template, appends cut rows, transforms inflow and load
/// noise into the scratch buffers, fills the patch buffer, and calls
/// `set_row_bounds`. After this call the solver is ready for `solve` /
/// `solve_with_basis`.
fn apply_opening_noise_and_patch<S: SolverInterface + Send>(
    ws: &mut SolverWorkspace<S>,
    ctx: &StageContext<'_>,
    training_ctx: &TrainingContext<'_>,
    cut_batch: &RowBatch,
    raw_noise: &[f64],
    x_hat: &[f64],
    s: usize,
) {
    // n_blks is zero when there are no load buses (avoids out-of-bounds access
    // on block_counts_per_stage which may be empty in deterministic cases).
    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,
    );
    ws.solver.load_model(&ctx.templates[s]);
    ws.solver.add_rows(cut_batch);
    ws.patch_buf.fill_forward_patches(
        training_ctx.indexer,
        x_hat,
        &ws.scratch.noise_buf,
        ctx.base_rows[s],
    );
    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,
        );
    }
    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],
    );
}

/// 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 mut working_basis: Option<Basis> = None;
    let s = succ.successor;

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

        let warm = if omega == 0 {
            succ.basis_store.get(m, s)
        } else {
            working_basis.as_ref()
        };
        let view = (if let Some(rb) = warm {
            ws.solver.solve_with_basis(rb)
        } 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;
        let coefficients: Vec<f64> = view.dual[..indexer.n_state].to_vec();
        let intercept = objective
            - 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] > 0.0 {
                    *accum.slot_increments.entry(slot).or_insert(0) += 1;
                }
            }
        }
        if let Some(rb) = working_basis.as_mut() {
            ws.solver.get_basis(rb);
        } else {
            let mut rb = Basis::new(ctx.templates[s].num_cols, ctx.templates[s].num_rows);
            ws.solver.get_basis(&mut rb);
            working_basis = Some(rb);
        }
        accum.outcomes.push(BackwardOutcome {
            intercept,
            coefficients,
            objective_value: objective,
        });
    }

    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()
        .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);
            let mut staged: Vec<StagedCut> = Vec::with_capacity(end_m - start_m);
            let mut accum = TrialAccumulators {
                outcomes: Vec::with_capacity(n_openings),
                slot_increments: std::collections::HashMap::new(),
            };

            for m in start_m..end_m {
                accum.outcomes.clear();
                accum.slot_increments.clear();
                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`
pub fn run_backward_pass<S: SolverInterface + Send, C: Communicator>(
    workspaces: &mut [SolverWorkspace<S>],
    basis_store: &BasisStore,
    ctx: &StageContext<'_>,
    fcf: &mut FutureCostFunction,
    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 tree_view = stochastic.tree_view();

    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() {
            spec.exchange.exchange(spec.records, t, num_stages, comm)?;
        }

        let successor = t + 1;
        let n_openings = tree_view.n_openings(successor);
        #[allow(clippy::cast_precision_loss)]
        let probabilities: Vec<f64> = vec![1.0_f64 / n_openings as f64; n_openings];

        let cut_batch = build_cut_row_batch(fcf, successor, indexer);
        let num_cuts_at_successor = cut_batch.num_rows;
        let template_num_rows = ctx.templates[successor].num_rows;
        let successor_active_slots: Vec<usize> = fcf
            .active_cuts(successor)
            .map(|(slot, _, _)| slot)
            .collect();

        let succ_spec = SuccessorSpec {
            t,
            successor,
            my_rank,
            probabilities: &probabilities,
            cut_batch: &cut_batch,
            num_cuts_at_successor,
            template_num_rows,
            successor_active_slots: &successor_active_slots,
            basis_store,
        };
        let worker_staged = process_stage_backward(workspaces, ctx, training_ctx, spec, &succ_spec);

        let mut all_staged: Vec<StagedCut> = Vec::with_capacity(spec.local_work);
        for worker_result in worker_staged {
            all_staged.extend(worker_result?);
        }
        all_staged.sort_by_key(|c| c.trial_point_idx);

        for cut in all_staged {
            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;
            }
        }
    }

    #[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,
    })
}

#[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},
        workspace::{BasisStore, SolverWorkspace},
    };

    /// 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> {
            unreachable!("StubComm allgatherv not used in backward pass tests")
        }

        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,
        }
    }

    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(),
                load_rhs_buf: Vec::new(),
                row_lower_buf: 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,
            },
        };

        #[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,
        };
        assert_eq!(r.cuts_generated, 6);
        assert_eq!(r.elapsed_ms, 42);
    }

    #[test]
    fn backward_result_clone_and_debug() {
        let r = BackwardResult {
            cuts_generated: 3,
            elapsed_ms: 100,
            lp_solves: 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 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: &[],
            },
            &mut fcf,
            &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,
            },
            &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 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: &[],
            },
            &mut fcf,
            &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,
            },
            &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 _ = 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: &[],
            },
            &mut fcf,
            &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,
            },
            &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 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: &[],
            },
            &mut fcf,
            &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,
            },
            &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 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: &[],
            },
            &mut fcf,
            &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,
            },
            &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 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: &[],
            },
            &mut fcf,
            &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,
            },
            &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 _ = 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: &[],
            },
            &mut fcf,
            &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,
            },
            &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 _ = 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: &[],
            },
            &mut fcf,
            &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,
            },
            &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 _ = 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: &[],
            },
            &mut fcf,
            &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,
            },
            &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 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: &[],
            },
            &mut fcf,
            &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,
            },
            &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 _ = 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: &[],
            },
            &mut fcf,
            &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,
            },
            &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 _ = 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: &[],
            },
            &mut fcf,
            &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,
            },
            &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 warm-start: given 3 openings at the same successor stage,
    /// the first opening cold-starts (store slot is None), and openings 2 and 3
    /// warm-start from the first opening's basis via `working_basis`.
    ///
    /// AC: Given a 2-stage system, 1 trial point, 3 openings, empty basis cache,
    /// then `solver.warm_start_calls == 2` after the backward pass (openings 2
    /// and 3 warm-start; opening 1 cold-starts).
    #[test]
    fn multi_opening_warm_start_subsequent_openings_use_basis() {
        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 _ = 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: &[],
            },
            &mut fcf,
            &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,
            },
            &comm,
        )
        .unwrap();

        // Opening 1: cold-start (store slot was None for scenario 0, stage 1).
        // Opening 2: warm-start (working_basis populated by opening 1).
        // Opening 3: warm-start (working_basis updated by opening 2).
        let warm_start_calls = workspaces[0].solver.warm_start_calls;
        assert_eq!(
            warm_start_calls, 2,
            "openings 2 and 3 must call solve_with_basis \
             (warm_start_calls == 2, 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 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: &[],
            },
            &mut fcf,
            &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,
            },
            &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(),
                load_rhs_buf: Vec::new(),
                row_lower_buf: 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: &[],
        };
        let _ = run_backward_pass(
            &mut workspaces_1,
            &basis_store_1,
            &ctx,
            &mut fcf_1,
            &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,
            },
            &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(),
                    load_rhs_buf: Vec::new(),
                    row_lower_buf: Vec::new(),
                },
            })
            .collect();
        let basis_store_4 = empty_basis_store(exchange.local_count(), n_stages);
        let _ = run_backward_pass(
            &mut workspaces_4,
            &basis_store_4,
            &ctx,
            &mut fcf_4,
            &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,
            },
            &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,
            },
        };

        #[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]
    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,
        };
        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(),
                load_rhs_buf: Vec::with_capacity(1),
                row_lower_buf: 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 _ = 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,
            },
            &mut fcf,
            &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,
            },
            &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]
    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,
        };
        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(),
                load_rhs_buf: Vec::new(),
                row_lower_buf: Vec::new(),
            },
        };
        let mut workspaces = vec![ws];
        let comm = StubComm;
        let basis_store = empty_basis_store(exchange.local_count(), n_stages);

        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],
            },
            &mut fcf,
            &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,
            },
            &comm,
        )
        .unwrap();

        // With n_load_buses=0, forward_patch_count must equal N*(2+L)=1*(2+0)=2.
        assert_eq!(
            workspaces[0].patch_buf.forward_patch_count(),
            2,
            "forward_patch_count must be N*(2+L)=2 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]
    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,
        };
        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(),
                load_rhs_buf: Vec::with_capacity(1),
                row_lower_buf: 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 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,
            },
            &mut fcf,
            &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,
            },
            &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()
        );
    }
}
