cobre-sddp 0.6.2

//! Training loop orchestrator for the SDDP algorithm.
//!
//! [`train`] wires together the forward pass, forward synchronization, state
//! exchange, backward pass, cut synchronization, lower bound evaluation, and
//! convergence check into a single iterative loop.
//!
//! ## Iteration lifecycle
//!
//! Each iteration follows the corrected ordering from the LB plan fix (F-019):
//!
//! 1. Forward pass — scenario simulation, local UB statistics.
//! 2. Forward sync — global synchronization for UB statistics.
//! 3. State exchange — `allgatherv` trial points for the backward pass.
//! 4. Backward pass — Benders cut generation.
//! 5. Cut sync — `allgatherv` new cuts across ranks.
//!    5a. Cut selection — optional periodic pool pruning via `CutSelectionStrategy`.
//!    5b. Budget enforcement — active-cut hard cap (every iteration when set).
//!    5c. Template baking — rebuild per-stage baked LP templates.
//! 6. Lower bound evaluation — rank 0 solves stage-0 openings, broadcasts scalar.
//! 7. Convergence check — stopping rules evaluated.
//! 8. Event emission — `IterationSummary` and per-step events via channel.
//!
//! ## Pre-allocation discipline
//!
//! All workspace buffers (`PatchBuffer`, `TrajectoryRecord` flat vec,
//! `ExchangeBuffers`, `CutSyncBuffers`) are allocated once before the
//! iteration loop and reused across all iterations. No heap allocation
//! occurs on the hot path.

use cobre_comm::Communicator;
use cobre_solver::{SolverInterface, StageTemplate};

use crate::{
    SddpError, TrainingConfig,
    context::{StageContext, TrainingContext},
    cut::fcf::FutureCostFunction,
    solver_stats::SolverStatsEntry,
    training_session::{IterationOutcome, TrainingSession},
    workspace::CapturedBasis,
};

// ---------------------------------------------------------------------------
// TrainingResult
// ---------------------------------------------------------------------------

/// Result of a training run that always carries partial results.
///
/// When training completes normally, `error` is `None` and `result` contains
/// the full training statistics. When training fails mid-iteration, `error`
/// carries the failure cause and `result` contains statistics from all
/// fully completed iterations (the failing iteration is excluded).
#[derive(Debug)]
pub struct TrainingOutcome {
    /// Training result from completed iterations. Always populated, even when
    /// `error` is `Some` -- in that case, `result.iterations` reflects only
    /// the iterations that completed without error.
    pub result: TrainingResult,

    /// If training was interrupted by an error, the cause. `None` when
    /// training completed normally (convergence, iteration limit, or
    /// time limit).
    pub error: Option<SddpError>,
}

/// Summary statistics produced when the training loop terminates.
#[non_exhaustive]
#[derive(Debug, Clone)]
pub struct TrainingResult {
    /// Final lower bound at termination.
    pub final_lb: f64,

    /// Final upper bound mean at termination.
    pub final_ub: f64,

    /// Final upper bound standard deviation at termination.
    pub final_ub_std: f64,

    /// Final convergence gap: `(UB - LB) / max(1.0, |UB|)`.
    pub final_gap: f64,

    /// Total number of iterations completed.
    pub iterations: u64,

    /// Human-readable termination reason (e.g., `"iteration_limit"`, `"graceful_shutdown"`).
    pub reason: String,

    /// Total wall-clock time for the training run, in milliseconds.
    pub total_time_ms: u64,

    /// Per-stage captured basis from the last iteration, indexed 0-based.
    ///
    /// Each entry is `Some(captured)` if the stage was solved at least once
    /// during the final iteration, or `None` if no solve occurred (e.g.,
    /// the last stage in finite-horizon mode has no successor cuts).
    /// Used for policy checkpoint persistence and simulation warm-start
    /// via slot-tracked reconstruction.
    pub basis_cache: Vec<Option<CapturedBasis>>,

    /// Per-iteration, per-phase solver statistics log.
    ///
    /// Each entry is `(iteration, phase_name, stage_index, delta)`.
    /// Phase names: `"forward"`, `"backward"`, `"lower_bound"`.
    /// Stage index is `-1` for forward and lower bound (which span all stages),
    /// and the actual stage index for backward per-stage entries (added in T-004).
    pub solver_stats_log: Vec<SolverStatsEntry>,

    /// Visited states archive containing all forward-pass trial points.
    ///
    /// Always populated during training. The caller decides whether to
    /// persist it to the policy checkpoint based on `exports.states`.
    pub visited_archive: Option<crate::visited_states::VisitedStatesArchive>,

    /// Final-iteration baked templates, one per stage.
    ///
    /// Always `Some` — templates are baked unconditionally before the first
    /// iteration of the training loop and never reverted.
    pub baked_templates: Option<Vec<StageTemplate>>,
}

impl TrainingResult {
    /// Construct a `TrainingResult` with explicit field values.
    ///
    /// Prefer this constructor over struct-literal syntax for compiler-enforced
    /// parity when adding new fields.
    ///
    /// Structurally independent parameters: each of the 11 fields is the output
    /// of a distinct training-loop stage (convergence metrics, timing, basis
    /// cache, solver stats, visited archive, baked templates). Absorbing them
    /// into a context struct would not reduce the argument count at the call
    /// sites, because each call site constructs the full `TrainingResult` as
    /// the terminal value — there is no shared upstream context to forward.
    // RATIONALE: 11 args map 1-to-1 to distinct output fields of the training result.
    // No shared context struct can be forwarded from the call sites; each call site
    // is the terminal aggregation point that collects all training outputs.
    #[must_use]
    #[allow(clippy::too_many_arguments, clippy::similar_names)]
    pub fn new(
        final_lb: f64,
        final_ub: f64,
        final_ub_std: f64,
        final_gap: f64,
        iterations: u64,
        reason: String,
        total_time_ms: u64,
        basis_cache: Vec<Option<CapturedBasis>>,
        solver_stats_log: Vec<SolverStatsEntry>,
        visited_archive: Option<crate::visited_states::VisitedStatesArchive>,
        baked_templates: Option<Vec<StageTemplate>>,
    ) -> Self {
        Self {
            final_lb,
            final_ub,
            final_ub_std,
            final_gap,
            iterations,
            reason,
            total_time_ms,
            basis_cache,
            solver_stats_log,
            visited_archive,
            baked_templates,
        }
    }
}

// ---------------------------------------------------------------------------
// Internal helpers
// ---------------------------------------------------------------------------

/// Convert a buffer length to `i32` for use as an MPI broadcast count.
///
/// Returns `Err(SddpError::Communication(CommError::InvalidBufferSize { .. }))`
/// when `len > i32::MAX`, carrying the operation name, the maximum legal value
/// (`i32::MAX as usize`), and the actual (oversized) length.
fn checked_broadcast_len(len: usize, operation: &'static str) -> Result<i32, SddpError> {
    i32::try_from(len).map_err(|_| {
        SddpError::Communication(cobre_comm::CommError::InvalidBufferSize {
            operation,
            expected: i32::MAX as usize,
            actual: len,
        })
    })
}

/// Build a `basis_cache` from the canonical global scenario 0, broadcasting
/// rank 0's bases to all other ranks so that every rank has an identical
/// warm-start basis for the simulation phase.
///
/// ## Why global scenario 0?
///
/// `basis_store` is indexed by **local** scenario index. Local scenario 0 on
/// rank 0 is always global scenario 0 (rank 0's `my_fwd_offset` is 0). On
/// rank *r > 0*, local scenario 0 maps to a different global scenario. Using
/// the last local scenario (`my_actual_fwd - 1`) therefore produces different
/// bases on different ranks, because each rank follows a distinct noise
/// realisation. Broadcasting rank 0's scenario 0 ensures all ranks start
/// simulation from the identical LP vertex.
///
/// ## Serialization format
///
/// Two broadcast buffers per call (four broadcasts total: length then payload
/// for each):
///
/// **i32 buffer** — for each stage *t* in `0..num_stages`:
/// - `0_i32` sentinel when `basis_store.get(0, t)` is `None`
/// - When `Some(captured)`: `1_i32` sentinel, then `col_len`, `row_len`,
///   `base_row_count`, `cut_slot_count`, `state_len` (all `i32`), then
///   `col_status[..]`, `row_status[..]`, `cut_row_slots[..] as i32`
///
/// **f64 buffer** — for each stage *t* in `0..num_stages`:
/// - When `sentinel == 1`: `state_at_capture[..]` appended sequentially
/// - When `sentinel == 0`: no entries appended
///
/// `u32` slot values are cast to `i32` for transmission; they are always
/// non-negative (LP pool indices) and fit comfortably in `i32`.
///
/// ## Single-rank optimization
///
/// When `comm.size() == 1` the broadcast is skipped; only the extraction
/// from local scenario 0 runs. All metadata is preserved via `Clone`.
///
/// # Errors
///
/// Returns `SddpError::Communication(CommError::InvalidBufferSize { .. })`
/// if either buffer length exceeds `i32::MAX` (the MPI count limit).
/// Buffer length scalars are broadcast as i32 (MPI `CommData` requires Copy,
/// ruling out usize); the usize→i32 conversions are checked via
/// [`checked_broadcast_len`] before any MPI call is issued.
/// Returns `SddpError::Communication` if any `comm.broadcast` call fails.
/// Returns `SddpError::Validation` if any length prefix is inconsistent with
/// the received buffer size, naming the offending stage in the message.
pub(crate) fn broadcast_basis_cache<C: Communicator>(
    basis_store: &crate::workspace::BasisStore,
    num_stages: usize,
    comm: &C,
) -> Result<Vec<Option<CapturedBasis>>, SddpError> {
    // Single-rank fast path: no communication needed — clone the full
    // CapturedBasis including metadata (cut_row_slots, state_at_capture,
    // base_row_count) so that simulation reconstruction has full slot
    // identity on single-rank runs.
    if comm.size() == 1 {
        let cache = (0..num_stages)
            .map(|t| basis_store.get(0, t).cloned())
            .collect();
        return Ok(cache);
    }

    // Multi-rank path: pack rank 0's scenario-0 full CapturedBasis into two
    // flat buffers (one i32 for integer/status fields, one f64 for
    // state_at_capture) and broadcast both.
    //
    // Wire format is owned by CapturedBasis::to_broadcast_payload /
    // CapturedBasis::try_from_broadcast_payload. The None case writes a
    // 0_i32 sentinel directly; the Some case delegates to the method.
    let mut buf: Vec<i32> = Vec::new();
    let mut f64_buf: Vec<f64> = Vec::new();
    if comm.rank() == 0 {
        for t in 0..num_stages {
            match basis_store.get(0, t) {
                None => buf.push(0_i32),
                Some(captured) => captured.to_broadcast_payload(&mut buf, &mut f64_buf),
            }
        }
    }

    // Step 1: broadcast the i32 buffer length so all ranks can allocate.
    let mut len_buf = [checked_broadcast_len(
        buf.len(),
        "broadcast_basis_cache_i32",
    )?];
    comm.broadcast(&mut len_buf, 0).map_err(SddpError::from)?;
    let total_len = usize::try_from(len_buf[0]).map_err(|_| {
        SddpError::Validation(format!(
            "broadcast_basis_cache_i32: received negative length {}",
            len_buf[0]
        ))
    })?;

    // Step 2: resize non-root i32 buffers and broadcast the i32 payload.
    buf.resize(total_len, 0_i32);
    comm.broadcast(&mut buf, 0).map_err(SddpError::from)?;

    // Step 3: broadcast the f64 buffer length so all ranks can allocate.
    let mut f64_len_buf = [checked_broadcast_len(
        f64_buf.len(),
        "broadcast_basis_cache_f64",
    )?];
    comm.broadcast(&mut f64_len_buf, 0)
        .map_err(SddpError::from)?;
    let f64_total_len = usize::try_from(f64_len_buf[0]).map_err(|_| {
        SddpError::Validation(format!(
            "broadcast_basis_cache_f64: received negative length {}",
            f64_len_buf[0]
        ))
    })?;

    // Step 4: resize non-root f64 buffers and broadcast the f64 payload.
    f64_buf.resize(f64_total_len, 0.0_f64);
    comm.broadcast(&mut f64_buf, 0).map_err(SddpError::from)?;

    // Step 5: deserialize back into Vec<Option<CapturedBasis>>.
    let mut cache: Vec<Option<CapturedBasis>> = Vec::with_capacity(num_stages);
    let mut pos = 0_usize;
    let mut f64_pos = 0_usize;
    for stage in 0..num_stages {
        let captured = CapturedBasis::try_from_broadcast_payload(
            stage,
            &buf,
            &mut pos,
            &f64_buf,
            &mut f64_pos,
        )?;
        cache.push(captured);
    }

    Ok(cache)
}

// ---------------------------------------------------------------------------
// train
// ---------------------------------------------------------------------------

/// Execute the SDDP training loop.
///
/// Allocates all workspace buffers, runs the iteration loop until a stopping
/// rule triggers or `config.max_iterations` is reached, and returns a
/// [`TrainingOutcome`] summarising the final convergence statistics.
///
/// ## Error handling
///
/// Returns `Err(SddpError)` immediately on:
/// - LP infeasibility in the forward or backward pass → `SddpError::Infeasible`
/// - Solver failure → `SddpError::Solver`
/// - Communication failure → `SddpError::Communication`
///
/// ## Event channel
///
/// When `config.event_sender` is `Some`, typed [`cobre_core::TrainingEvent`] values are
/// emitted at each lifecycle step boundary. Event send failures (receiver
/// dropped) are silently ignored so they cannot interrupt training.
///
/// ## Cut selection
///
/// When `cut_selection` is `Some(strategy)`, step 5a runs after every cut
/// synchronisation. The strategy's `should_run(iteration)` gate controls the
/// frequency; at eligible iterations every stage's cut pool is scanned and
/// inactive cuts are deactivated. When `cut_selection` is `None`, step 5a is
/// skipped entirely and no [`cobre_core::TrainingEvent::PolicySelectionComplete`] events are
/// emitted.
///
/// # Errors
///
/// Returns `Err(SddpError::Infeasible { .. })` when an LP has no feasible
/// solution. Returns `Err(SddpError::Solver(_))` for other solver failures.
/// Returns `Err(SddpError::Communication(_))` when a collective operation
/// fails.
///
/// # Examples
///
/// ```rust,ignore
/// use cobre_sddp::{train, TrainingConfig, LoopConfig, CutManagementConfig, EventConfig};
/// use cobre_sddp::{StoppingRuleSet, StoppingRule, RiskMeasure, HorizonMode};
///
/// let mut solver = HiggsBackend::new();
/// let config = TrainingConfig {
///     loop_config: LoopConfig { forward_passes: 100, max_iterations: 100, ..LoopConfig::default() },
///     cut_management: CutManagementConfig {
///         risk_measures: vec![RiskMeasure::Expectation; num_stages],
///         ..CutManagementConfig::default()
///     },
///     events: EventConfig::default(),
/// };
/// let mut fcf = FutureCostFunction::new(num_stages - 1, n_state, capacity);
///
/// let result = train(
///     &mut solver, config, &mut fcf, &stage_ctx, &training_ctx, &comm,
///     || HiggsBackend::new(),
/// )?;
///
/// println!("converged in {} iterations, gap={:.4}", result.result.iterations, result.result.final_gap);
/// ```
///
/// # Panics
///
/// In debug builds, panics if `templates.len() != horizon.num_stages()` or if
/// `config.cut_management.risk_measures.len() != horizon.num_stages()` or if
/// `training_ctx.stochastic.opening_tree().n_openings(0) == 0`.
///
/// Always panics if `comm.rank() > i32::MAX`. MPI world sizes are bounded well
/// below this on all real systems.
pub fn train<S: SolverInterface + Send, C: Communicator>(
    solver: &mut S,
    config: TrainingConfig,
    fcf: &mut FutureCostFunction,
    stage_ctx: &StageContext<'_>,
    training_ctx: &TrainingContext<'_>,
    comm: &C,
    solver_factory: impl Fn() -> Result<S, cobre_solver::SolverError>,
) -> Result<TrainingOutcome, SddpError> {
    let mut session = TrainingSession::new(
        solver,
        config,
        fcf,
        stage_ctx,
        training_ctx,
        comm,
        solver_factory,
    )?;
    for iteration in session.iteration_range() {
        match session.run_iteration(iteration) {
            Ok(IterationOutcome::Continue) => {}
            Ok(IterationOutcome::Converged | IterationOutcome::Shutdown) => break,
            Err(e) => return session.finalize_with_error(e),
        }
    }
    session.finalize()
}

#[cfg(test)]
#[allow(
    clippy::unwrap_used,
    clippy::expect_used,
    clippy::panic,
    clippy::float_cmp,
    clippy::cast_precision_loss,
    clippy::cast_possible_truncation,
    clippy::cast_possible_wrap,
    clippy::too_many_lines,
    clippy::doc_markdown,
    clippy::needless_range_loop
)]
mod tests {
    use std::collections::BTreeMap;
    use std::sync::mpsc;

    use chrono::NaiveDate;
    use cobre_comm::{CommData, CommError, Communicator, ReduceOp};
    use cobre_core::{
        Bus, EntityId, SystemBuilder, TrainingEvent, WorkerTimingPhase,
        scenario::{
            CorrelationEntity, CorrelationGroup, CorrelationModel, CorrelationProfile,
            SamplingScheme,
        },
        temporal::{
            Block, BlockMode, NoiseMethod, ScenarioSourceConfig, Stage, StageRiskConfig,
            StageStateConfig,
        },
    };
    use cobre_solver::{
        Basis, LpSolution, RowBatch, SolverError, SolverInterface, SolverStatistics, StageTemplate,
    };
    use cobre_stochastic::{
        ClassSchemes, OpeningTreeInputs, StochasticContext, build_stochastic_context,
    };

    use super::train;
    use crate::{
        StoppingMode, StoppingRule, StoppingRuleSet, TrainingConfig,
        config::{CutManagementConfig, EventConfig, LoopConfig},
        context::{StageContext, TrainingContext},
        cut::fcf::FutureCostFunction,
        error::SddpError,
        horizon_mode::HorizonMode,
        indexer::StageIndexer,
        inflow_method::InflowNonNegativityMethod,
        risk_measure::RiskMeasure,
        solver_stats::SolverStatsDelta,
    };

    /// Minimal LP for N=1 hydro, L=0 PAR order.
    ///
    /// Column layout (N=1, L=0):
    /// - col 0: `storage_out` (no NZ in structural rows)
    /// - col 1: `z_inflow` (no NZ — `z_inflow` row at row 1)
    /// - col 2: `storage_in` (1 NZ: row 0, storage-fixing row)
    /// - col 3: `theta` (no NZ)
    ///
    /// Row layout:
    /// - row 0: storage-fixing (`storage_out` fixed to incoming state)
    /// - row 1: `z_inflow` definition row
    fn minimal_template(n_state: usize) -> StageTemplate {
        let _ = n_state;
        StageTemplate {
            num_cols: 4,
            num_rows: 2,
            num_nz: 1,
            // CSC col_starts: 4 cols + 1 sentinel = 5 entries.
            // col 0 (storage_out): 0 NZ
            // col 1 (z_inflow):    0 NZ
            // col 2 (storage_in):  1 NZ at row 0
            // col 3 (theta):       0 NZ
            col_starts: vec![0_i32, 0, 0, 1, 1],
            row_indices: vec![0_i32],
            values: vec![1.0],
            col_lower: vec![0.0, f64::NEG_INFINITY, 0.0, 0.0],
            col_upper: vec![f64::INFINITY; 4],
            objective: vec![0.0, 0.0, 0.0, 1.0],
            row_lower: vec![0.0, 0.0],
            row_upper: vec![0.0, 0.0],
            n_state: 1,
            n_transfer: 0,
            n_dual_relevant: 1,
            n_hydro: 1,
            max_par_order: 0,
            col_scale: Vec::new(),
            row_scale: Vec::new(),
        }
    }

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

    /// Mock solver that returns fixed objective values in sequence.
    ///
    /// Each call to `solve()` returns the next value from `objectives`,
    /// wrapping around. If `infeasible_on_first` is set, the first call
    /// returns `SolverError::Infeasible`.
    struct MockSolver {
        objectives: Vec<f64>,
        call_count: usize,
        infeasible_on_first: bool,
    }

    impl MockSolver {
        fn with_fixed(objective: f64) -> Self {
            Self {
                objectives: vec![objective],
                call_count: 0,
                infeasible_on_first: false,
            }
        }

        fn infeasible() -> Self {
            Self {
                objectives: vec![0.0],
                call_count: 0,
                infeasible_on_first: true,
            }
        }
    }

    impl SolverInterface for MockSolver {
        fn solver_name_version(&self) -> String {
            "MockSolver 0.0.0".to_string()
        }
        fn load_model(&mut self, _t: &StageTemplate) {}
        fn add_rows(&mut self, _r: &RowBatch) {}
        fn set_row_bounds(&mut self, _i: &[usize], _l: &[f64], _u: &[f64]) {}
        fn set_col_bounds(&mut self, _i: &[usize], _l: &[f64], _u: &[f64]) {}

        fn solve(
            &mut self,
            _basis: Option<&Basis>,
        ) -> Result<cobre_solver::SolutionView<'_>, SolverError> {
            let call = self.call_count;
            self.call_count += 1;
            if self.infeasible_on_first && call == 0 {
                return Err(SolverError::Infeasible);
            }
            let obj = self.objectives[call % self.objectives.len()];
            // Return a view with primal[3] = 0.0 (theta = 0, N=1 L=0 → theta at col 3)
            // so that the forward pass computes stage_cost = objective - primal[theta]
            // = obj - 0 = obj.  The fixed_solution helper provides compatible arrays.
            let sol = fixed_solution(obj);
            // We cannot borrow from a temporary, so we use static empty slices.
            // training.rs mock only needs to satisfy the SolverInterface bound;
            // the actual slice contents are not checked by the training loop.
            let _ = sol;
            Ok(cobre_solver::SolutionView {
                objective: obj,
                primal: &[0.0, 0.0, 0.0, 0.0],
                dual: &[0.0, 0.0],
                reduced_costs: &[0.0, 0.0, 0.0, 0.0],
                iterations: 0,
                solve_time_seconds: 0.0,
            })
        }

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

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

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

    struct StubComm;

    impl Communicator for StubComm {
        fn allgatherv<T: CommData>(
            &self,
            send: &[T],
            recv: &mut [T],
            _counts: &[usize],
            _displs: &[usize],
        ) -> Result<(), CommError> {
            recv[..send.len()].clone_from_slice(send);
            Ok(())
        }

        fn allreduce<T: CommData>(
            &self,
            send: &[T],
            recv: &mut [T],
            _op: ReduceOp,
        ) -> Result<(), CommError> {
            recv.clone_from_slice(send);
            Ok(())
        }

        fn broadcast<T: CommData>(&self, _buf: &mut [T], _root: usize) -> Result<(), CommError> {
            Ok(())
        }

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

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

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

        fn abort(&self, error_code: i32) -> ! {
            std::process::exit(error_code)
        }
    }

    /// Build a minimal `StochasticContext` with `n_stages` stages and a single
    /// hydro entity.
    ///
    /// Used to provide the `stochastic` argument to `train`. Follows the same
    /// [`SystemBuilder`] + `build_stochastic_context` pattern as the forward-pass
    /// integration tests.  The `n_openings` parameter controls the branching
    /// factor of the opening tree.
    #[allow(clippy::cast_possible_wrap)]
    fn make_stochastic_context(n_stages: usize, n_openings: usize) -> StochasticContext {
        use cobre_core::entities::hydro::{Hydro, HydroGenerationModel, HydroPenalties};
        use cobre_core::scenario::InflowModel;

        let bus = Bus {
            id: EntityId(0),
            name: "B0".to_string(),
            deficit_segments: vec![cobre_core::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,
            min_turbined_m3s: 0.0,
            max_turbined_m3s: 100.0,
            specific_productivity_mw_per_m3s_per_m: None,
            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,
                turbined_cost: 0.0,
                storage_violation_below_cost: 0.0,
                filling_target_violation_cost: 0.0,
                turbined_violation_below_cost: 0.0,
                outflow_violation_below_cost: 0.0,
                outflow_violation_above_cost: 0.0,
                generation_violation_below_cost: 0.0,
                evaporation_violation_cost: 0.0,
                water_withdrawal_violation_cost: 0.0,
                water_withdrawal_violation_pos_cost: 0.0,
                water_withdrawal_violation_neg_cost: 0.0,
                evaporation_violation_pos_cost: 0.0,
                evaporation_violation_neg_cost: 0.0,
                inflow_nonnegativity_cost: 1000.0,
            },
        };

        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: n_openings,
                noise_method: NoiseMethod::Saa,
            },
        };

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

        let inflow_models: Vec<InflowModel> = (0..n_stages)
            .map(|i| InflowModel {
                hydro_id: EntityId(1),
                stage_id: i as i32,
                mean_m3s: 100.0,
                std_m3s: 30.0,
                ar_coefficients: vec![],
                residual_std_ratio: 1.0,
                annual: None,
            })
            .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: "spectral".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,
            &[],
            &[],
            OpeningTreeInputs::default(),
            ClassSchemes {
                inflow: Some(SamplingScheme::InSample),
                load: Some(SamplingScheme::InSample),
                ncs: Some(SamplingScheme::InSample),
            },
        )
        .unwrap()
    }

    /// Build `n_stages` minimal [`Stage`] values with sequential `id`s (0..n_stages).
    ///
    /// Used to populate [`TrainingContext::stages`] so that
    /// [`cobre_stochastic::build_forward_sampler`] can read per-stage noise methods.
    fn make_stages(n_stages: usize) -> Vec<Stage> {
        (0..n_stages)
            .map(|i| Stage {
                index: i,
                id: i as i32,
                start_date: chrono::NaiveDate::from_ymd_opt(2024, 1, 1).unwrap(),
                end_date: chrono::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: cobre_core::temporal::StageStateConfig {
                    storage: true,
                    inflow_lags: false,
                },
                risk_config: cobre_core::temporal::StageRiskConfig::Expectation,
                scenario_config: ScenarioSourceConfig {
                    branching_factor: 1,
                    noise_method: NoiseMethod::Saa,
                },
            })
            .collect()
    }

    fn make_fcf(
        n_stages: usize,
        n_state: usize,
        forward_passes: u32,
        max_iter: u64,
    ) -> FutureCostFunction {
        FutureCostFunction::new(
            n_stages,
            n_state,
            forward_passes,
            max_iter,
            &vec![0; n_stages],
        )
    }

    fn iteration_limit_rules(limit: u64) -> StoppingRuleSet {
        StoppingRuleSet {
            rules: vec![StoppingRule::IterationLimit { limit }],
            mode: StoppingMode::Any,
        }
    }

    /// AC: `train_completes_with_iteration_limit`
    ///
    /// Given `max_iterations: 5`, a `StoppingRuleSet` with
    /// `IterationLimit { limit: 5 }` in `Any` mode, a mock solver returning
    /// fixed objectives, and `StubComm` as communicator, when the function
    /// completes, then `result.iterations == 5` and
    /// `result.reason == "iteration_limit"`.
    #[test]
    fn ac_train_completes_with_iteration_limit() {
        let n_stages = 2;
        let indexer = StageIndexer::new(1, 0); // N=1, L=0
        let templates = vec![minimal_template(indexer.n_state); n_stages];
        let base_rows = vec![2usize; n_stages];
        let initial_state = vec![0.0_f64; indexer.n_state];
        let stochastic = make_stochastic_context(n_stages, 1);
        let stages = make_stages(n_stages);
        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let mut fcf = make_fcf(n_stages, indexer.n_state, 1, 10);

        let config = TrainingConfig {
            loop_config: LoopConfig {
                forward_passes: 1,
                max_iterations: 5,
                start_iteration: 0,
                n_fwd_threads: 1,
                max_blocks: 1,
                stopping_rules: iteration_limit_rules(5),
            },
            cut_management: CutManagementConfig {
                cut_selection: None,
                budget: None,
                cut_activity_tolerance: 0.0,
                basis_activity_window: crate::basis_reconstruct::DEFAULT_BASIS_ACTIVITY_WINDOW,
                warm_start_cuts: 0,
                risk_measures: vec![RiskMeasure::Expectation; n_stages],
            },
            events: EventConfig {
                event_sender: None,
                checkpoint_interval: None,
                shutdown_flag: None,
                export_states: false,
            },
        };

        let mut solver = MockSolver::with_fixed(100.0);
        let comm = StubComm;

        let stage_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: &[1usize, 1],
            ncs_max_gen: &[],
            ncs_allow_curtailment: &[],
            discount_factors: &[],
            cumulative_discount_factors: &[],
            stage_lag_transitions: &[],
            noise_group_ids: &[],
            downstream_par_order: 0,
        };
        let result = train(
            &mut solver,
            config,
            &mut fcf,
            &stage_ctx,
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &initial_state,
                inflow_scheme: SamplingScheme::InSample,
                load_scheme: SamplingScheme::InSample,
                ncs_scheme: SamplingScheme::InSample,
                stages: &stages,
                historical_library: None,
                external_inflow_library: None,
                external_load_library: None,
                external_ncs_library: None,
                recent_accum_seed: &[],
                recent_weight_seed: 0.0,
            },
            &comm,
            || Ok(MockSolver::with_fixed(100.0)),
        )
        .unwrap();

        assert!(result.error.is_none(), "expected no error");
        assert_eq!(result.result.iterations, 5, "expected 5 iterations");
        assert_eq!(result.result.reason, "iteration_limit");
    }

    /// AC: `train_returns_partial_on_infeasible`
    ///
    /// Given a mock solver that returns `SolverError::Infeasible` on the first
    /// forward pass solve, when the function is called, then it returns
    /// `Ok(TrainingOutcome)` with `error: Some(SddpError::Infeasible { .. })`
    /// and `result.iterations == 0`.
    #[test]
    fn ac_train_returns_partial_on_infeasible() {
        let n_stages = 2;
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template(indexer.n_state); n_stages];
        let base_rows = vec![2usize; n_stages];
        let initial_state = vec![0.0_f64; indexer.n_state];
        let stochastic = make_stochastic_context(n_stages, 1);
        let stages = make_stages(n_stages);
        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let mut fcf = make_fcf(n_stages, indexer.n_state, 1, 10);

        let config = TrainingConfig {
            loop_config: LoopConfig {
                forward_passes: 1,
                max_iterations: 5,
                start_iteration: 0,
                n_fwd_threads: 1,
                max_blocks: 1,
                stopping_rules: iteration_limit_rules(5),
            },
            cut_management: CutManagementConfig {
                cut_selection: None,
                budget: None,
                cut_activity_tolerance: 0.0,
                basis_activity_window: crate::basis_reconstruct::DEFAULT_BASIS_ACTIVITY_WINDOW,
                warm_start_cuts: 0,
                risk_measures: vec![RiskMeasure::Expectation; n_stages],
            },
            events: EventConfig {
                event_sender: None,
                checkpoint_interval: None,
                shutdown_flag: None,
                export_states: false,
            },
        };

        let mut solver = MockSolver::infeasible();
        let comm = StubComm;

        let stage_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: &[1usize, 1],
            ncs_max_gen: &[],
            ncs_allow_curtailment: &[],
            discount_factors: &[],
            cumulative_discount_factors: &[],
            stage_lag_transitions: &[],
            noise_group_ids: &[],
            downstream_par_order: 0,
        };
        let result = train(
            &mut solver,
            config,
            &mut fcf,
            &stage_ctx,
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &initial_state,
                inflow_scheme: SamplingScheme::InSample,
                load_scheme: SamplingScheme::InSample,
                ncs_scheme: SamplingScheme::InSample,
                stages: &stages,
                historical_library: None,
                external_inflow_library: None,
                external_load_library: None,
                external_ncs_library: None,
                recent_accum_seed: &[],
                recent_weight_seed: 0.0,
            },
            &comm,
            || Ok(MockSolver::infeasible()),
        );

        let outcome = result.unwrap();
        assert!(
            outcome.error.is_some(),
            "expected error in TrainingOutcome, got: {outcome:?}"
        );
        assert!(
            matches!(outcome.error, Some(SddpError::Infeasible { stage: 0, .. })),
            "expected SddpError::Infeasible at stage 0, got: {:?}",
            outcome.error
        );
        assert_eq!(
            outcome.result.iterations, 0,
            "no iterations should have completed"
        );
        assert_eq!(outcome.result.reason, "error");
    }

    /// AC: `train_emits_correct_event_sequence`
    ///
    /// Given `train` with `event_sender: Some(tx)`, runs for 2 iterations
    /// before `IterationLimit(2)` triggers. The receiver must collect exactly:
    ///
    /// - 1 `TrainingStarted`
    /// - 2 × (`WorkerTiming` (Forward), `ForwardPassComplete`, `ForwardSyncComplete`,
    ///   `WorkerTiming` (Backward), `BackwardPassComplete`, `PolicySyncComplete`,
    ///   `PolicyTemplateBakeComplete`, `ConvergenceUpdate`, `IterationSummary`)
    /// - 1 `TrainingFinished`
    ///
    /// = 1 + 18 + 1 = 20 events.
    #[test]
    fn ac_train_emits_correct_event_sequence() {
        let n_stages = 2;
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template(indexer.n_state); n_stages];
        let base_rows = vec![2usize; n_stages];
        let initial_state = vec![0.0_f64; indexer.n_state];
        let stochastic = make_stochastic_context(n_stages, 1);
        let stages = make_stages(n_stages);
        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let mut fcf = make_fcf(n_stages, indexer.n_state, 1, 10);

        let (tx, rx) = mpsc::channel::<TrainingEvent>();

        let config = TrainingConfig {
            loop_config: LoopConfig {
                forward_passes: 1,
                max_iterations: 10,
                start_iteration: 0,
                n_fwd_threads: 1,
                max_blocks: 1,
                stopping_rules: iteration_limit_rules(2),
            },
            cut_management: CutManagementConfig {
                cut_selection: None,
                budget: None,
                cut_activity_tolerance: 0.0,
                basis_activity_window: crate::basis_reconstruct::DEFAULT_BASIS_ACTIVITY_WINDOW,
                warm_start_cuts: 0,
                risk_measures: vec![RiskMeasure::Expectation; n_stages],
            },
            events: EventConfig {
                event_sender: Some(tx),
                checkpoint_interval: None,
                shutdown_flag: None,
                export_states: false,
            },
        };

        let mut solver = MockSolver::with_fixed(100.0);
        let comm = StubComm;

        let stage_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: &[1usize, 1],
            ncs_max_gen: &[],
            ncs_allow_curtailment: &[],
            discount_factors: &[],
            cumulative_discount_factors: &[],
            stage_lag_transitions: &[],
            noise_group_ids: &[],
            downstream_par_order: 0,
        };
        train(
            &mut solver,
            config,
            &mut fcf,
            &stage_ctx,
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &initial_state,
                inflow_scheme: SamplingScheme::InSample,
                load_scheme: SamplingScheme::InSample,
                ncs_scheme: SamplingScheme::InSample,
                stages: &stages,
                historical_library: None,
                external_inflow_library: None,
                external_load_library: None,
                external_ncs_library: None,
                recent_accum_seed: &[],
                recent_weight_seed: 0.0,
            },
            &comm,
            || Ok(MockSolver::with_fixed(100.0)),
        )
        .unwrap();

        // Drain all events.
        drop(fcf); // not needed; just for clarity
        let events: Vec<TrainingEvent> = rx.try_iter().collect();

        // 1 TrainingStarted + 2*(9 per-iteration) + 1 TrainingFinished = 20
        // Per-iteration: WorkerTiming(Forward), ForwardPassComplete,
        //   ForwardSyncComplete, WorkerTiming(Backward), BackwardPassComplete,
        //   PolicySyncComplete, PolicyTemplateBakeComplete, ConvergenceUpdate, IterationSummary
        assert_eq!(
            events.len(),
            20,
            "expected 20 events, got {} ({events:?})",
            events.len()
        );

        assert!(
            matches!(events[0], TrainingEvent::TrainingStarted { .. }),
            "first event must be TrainingStarted"
        );
        assert!(
            matches!(events.last(), Some(TrainingEvent::TrainingFinished { .. })),
            "last event must be TrainingFinished"
        );

        // Check per-iteration event pattern for iteration 1 (events[1..10])
        assert!(matches!(
            events[1],
            TrainingEvent::WorkerTiming {
                phase: WorkerTimingPhase::Forward,
                ..
            }
        ));
        assert!(matches!(
            events[2],
            TrainingEvent::ForwardPassComplete { .. }
        ));
        assert!(matches!(
            events[3],
            TrainingEvent::ForwardSyncComplete { .. }
        ));
        assert!(matches!(
            events[4],
            TrainingEvent::WorkerTiming {
                phase: WorkerTimingPhase::Backward,
                ..
            }
        ));
        assert!(matches!(
            events[5],
            TrainingEvent::BackwardPassComplete { .. }
        ));
        assert!(matches!(
            events[6],
            TrainingEvent::PolicySyncComplete { .. }
        ));
        assert!(matches!(
            events[7],
            TrainingEvent::PolicyTemplateBakeComplete { .. }
        ));
        assert!(matches!(events[8], TrainingEvent::ConvergenceUpdate { .. }));
        assert!(matches!(events[9], TrainingEvent::IterationSummary { .. }));

        // Iteration 2 (events[10..19]) follows the same pattern.
        assert!(matches!(
            events[10],
            TrainingEvent::WorkerTiming {
                phase: WorkerTimingPhase::Forward,
                ..
            }
        ));
        assert!(matches!(
            events[11],
            TrainingEvent::ForwardPassComplete { .. }
        ));
        assert!(matches!(
            events[12],
            TrainingEvent::ForwardSyncComplete { .. }
        ));
        assert!(matches!(
            events[13],
            TrainingEvent::WorkerTiming {
                phase: WorkerTimingPhase::Backward,
                ..
            }
        ));
        assert!(matches!(
            events[14],
            TrainingEvent::BackwardPassComplete { .. }
        ));
        assert!(matches!(
            events[15],
            TrainingEvent::PolicySyncComplete { .. }
        ));
        assert!(matches!(
            events[16],
            TrainingEvent::PolicyTemplateBakeComplete { .. }
        ));
        assert!(matches!(
            events[17],
            TrainingEvent::ConvergenceUpdate { .. }
        ));
        assert!(matches!(events[18], TrainingEvent::IterationSummary { .. }));
    }

    /// with 4 workers and 1 iteration, exactly
    /// `2 * n_workers_local = 8` WorkerTiming events are emitted with `rank = 0`
    /// and `worker_id ∈ {0, 1, 2, 3}`; AND the sum of `timings[BWD_SETUP]` across
    /// the 4 backward emissions equals `BackwardPassComplete.setup_time_ms` for
    /// the same iteration (within ±1 ms numerical tolerance).
    #[test]
    fn ac_worker_timing_per_worker_event_count_and_setup_invariant() {
        use cobre_core::WorkerTimingPhase;

        let n_stages = 2;
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template(indexer.n_state); n_stages];
        let base_rows = vec![2usize; n_stages];
        let initial_state = vec![0.0_f64; indexer.n_state];
        let stochastic = make_stochastic_context(n_stages, 1);
        let stages = make_stages(n_stages);
        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let mut fcf = make_fcf(n_stages, indexer.n_state, 1, 10);

        let (tx, rx) = mpsc::channel::<TrainingEvent>();

        let config = TrainingConfig {
            loop_config: LoopConfig {
                forward_passes: 1,
                max_iterations: 10,
                start_iteration: 0,
                n_fwd_threads: 4,
                max_blocks: 1,
                stopping_rules: iteration_limit_rules(1),
            },
            cut_management: CutManagementConfig {
                cut_selection: None,
                budget: None,
                cut_activity_tolerance: 0.0,
                basis_activity_window: crate::basis_reconstruct::DEFAULT_BASIS_ACTIVITY_WINDOW,
                warm_start_cuts: 0,
                risk_measures: vec![RiskMeasure::Expectation; n_stages],
            },
            events: EventConfig {
                event_sender: Some(tx),
                checkpoint_interval: None,
                shutdown_flag: None,
                export_states: false,
            },
        };

        let mut solver = MockSolver::with_fixed(100.0);
        let comm = StubComm;

        let stage_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: &[1usize, 1],
            ncs_max_gen: &[],
            ncs_allow_curtailment: &[],
            discount_factors: &[],
            cumulative_discount_factors: &[],
            stage_lag_transitions: &[],
            noise_group_ids: &[],
            downstream_par_order: 0,
        };
        train(
            &mut solver,
            config,
            &mut fcf,
            &stage_ctx,
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &initial_state,
                inflow_scheme: SamplingScheme::InSample,
                load_scheme: SamplingScheme::InSample,
                ncs_scheme: SamplingScheme::InSample,
                stages: &stages,
                historical_library: None,
                external_inflow_library: None,
                external_load_library: None,
                external_ncs_library: None,
                recent_accum_seed: &[],
                recent_weight_seed: 0.0,
            },
            &comm,
            || Ok(MockSolver::with_fixed(100.0)),
        )
        .unwrap();

        let events: Vec<TrainingEvent> = rx.try_iter().collect();

        // Collect all WorkerTiming events and the BackwardPassComplete event.
        let worker_events: Vec<&TrainingEvent> = events
            .iter()
            .filter(|e| matches!(e, TrainingEvent::WorkerTiming { .. }))
            .collect();
        // C2: exactly 8 WorkerTiming events (4 workers × 2 phases × 1 iteration).
        assert_eq!(
            worker_events.len(),
            8,
            "expected 8 WorkerTiming events (4 workers × 2 phases × 1 iter), got {}",
            worker_events.len()
        );

        // C2: rank=0, worker_id ∈ {0,1,2,3}, iteration=1 for every event.
        let mut fwd_workers = std::collections::BTreeSet::new();
        let mut bwd_workers = std::collections::BTreeSet::new();
        let mut bwd_setup_sum_ms = 0.0_f64;
        for ev in &worker_events {
            let TrainingEvent::WorkerTiming {
                rank,
                worker_id,
                iteration,
                phase,
                timings,
            } = ev
            else {
                unreachable!()
            };
            assert_eq!(*rank, 0, "expected rank=0 in single-rank stub");
            assert!(
                (0..4).contains(worker_id),
                "worker_id {worker_id} out of [0,4)"
            );
            assert_eq!(*iteration, 1, "expected iteration=1 (max_iterations=1)");
            match phase {
                WorkerTimingPhase::Forward => {
                    assert!(
                        fwd_workers.insert(*worker_id),
                        "worker_id {worker_id} duplicated in Forward emissions"
                    );
                    // Forward-only fields can be non-zero; backward-only fields must be 0.
                    assert_eq!(timings.bwd_setup_ms, 0.0);
                }
                WorkerTimingPhase::Backward => {
                    assert!(
                        bwd_workers.insert(*worker_id),
                        "worker_id {worker_id} duplicated in Backward emissions"
                    );
                    bwd_setup_sum_ms += timings.bwd_setup_ms;
                    // Backward-only fields can be non-zero; forward-only fields must be 0.
                    assert_eq!(timings.fwd_setup_ms, 0.0);
                }
            }
        }
        assert_eq!(fwd_workers.len(), 4, "expected 4 distinct forward workers");
        assert_eq!(bwd_workers.len(), 4, "expected 4 distinct backward workers");

        // C3: setup-sum invariant — sum of per-worker BWD_SETUP equals
        // BackwardPassComplete.setup_time_ms within ±1 ms tolerance.
        // (BackwardPassComplete.setup_time_ms is u64; per-worker timings are f64.)
        let bwd_setup_total_ms_u64 = events
            .iter()
            .find_map(|e| match e {
                TrainingEvent::BackwardPassComplete { setup_time_ms, .. } => Some(*setup_time_ms),
                _ => None,
            })
            .expect("BackwardPassComplete event must exist");
        #[allow(clippy::cast_precision_loss)]
        let bwd_setup_total_ms = bwd_setup_total_ms_u64 as f64;
        assert!(
            (bwd_setup_sum_ms - bwd_setup_total_ms).abs() < 1.0,
            "sum of per-worker BWD_SETUP ({bwd_setup_sum_ms} ms) must match \
             BackwardPassComplete.setup_time_ms ({bwd_setup_total_ms} ms) within ±1 ms"
        );
    }

    /// AC: `train_result_fields_populated`
    ///
    /// Verify that all `TrainingResult` fields are non-default after a
    /// successful 5-iteration run.
    #[test]
    fn ac_train_result_fields_populated() {
        let n_stages = 2;
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template(indexer.n_state); n_stages];
        let base_rows = vec![2usize; n_stages];
        let initial_state = vec![0.0_f64; indexer.n_state];
        let stochastic = make_stochastic_context(n_stages, 1);
        let stages = make_stages(n_stages);
        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let mut fcf = make_fcf(n_stages, indexer.n_state, 1, 10);

        let config = TrainingConfig {
            loop_config: LoopConfig {
                forward_passes: 1,
                max_iterations: 5,
                start_iteration: 0,
                n_fwd_threads: 1,
                max_blocks: 1,
                stopping_rules: iteration_limit_rules(5),
            },
            cut_management: CutManagementConfig {
                cut_selection: None,
                budget: None,
                cut_activity_tolerance: 0.0,
                basis_activity_window: crate::basis_reconstruct::DEFAULT_BASIS_ACTIVITY_WINDOW,
                warm_start_cuts: 0,
                risk_measures: vec![RiskMeasure::Expectation; n_stages],
            },
            events: EventConfig {
                event_sender: None,
                checkpoint_interval: None,
                shutdown_flag: None,
                export_states: false,
            },
        };

        let mut solver = MockSolver::with_fixed(100.0);
        let comm = StubComm;

        let stage_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: &[1usize, 1],
            ncs_max_gen: &[],
            ncs_allow_curtailment: &[],
            discount_factors: &[],
            cumulative_discount_factors: &[],
            stage_lag_transitions: &[],
            noise_group_ids: &[],
            downstream_par_order: 0,
        };
        let result = train(
            &mut solver,
            config,
            &mut fcf,
            &stage_ctx,
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &initial_state,
                inflow_scheme: SamplingScheme::InSample,
                load_scheme: SamplingScheme::InSample,
                ncs_scheme: SamplingScheme::InSample,
                stages: &stages,
                historical_library: None,
                external_inflow_library: None,
                external_load_library: None,
                external_ncs_library: None,
                recent_accum_seed: &[],
                recent_weight_seed: 0.0,
            },
            &comm,
            || Ok(MockSolver::with_fixed(100.0)),
        )
        .unwrap();

        assert!(result.error.is_none(), "expected no error");
        assert_eq!(result.result.iterations, 5);
        assert!(!result.result.reason.is_empty(), "reason must not be empty");
    }

    /// AC: `train_with_no_event_sender`
    ///
    /// Verify that `event_sender: None` does not panic and training completes
    /// normally.
    #[test]
    fn ac_train_with_no_event_sender() {
        let n_stages = 2;
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template(indexer.n_state); n_stages];
        let base_rows = vec![2usize; n_stages];
        let initial_state = vec![0.0_f64; indexer.n_state];
        let stochastic = make_stochastic_context(n_stages, 1);
        let stages = make_stages(n_stages);
        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let mut fcf = make_fcf(n_stages, indexer.n_state, 1, 10);

        let config = TrainingConfig {
            loop_config: LoopConfig {
                forward_passes: 1,
                max_iterations: 2,
                start_iteration: 0,
                n_fwd_threads: 1,
                max_blocks: 1,
                stopping_rules: iteration_limit_rules(2),
            },
            cut_management: CutManagementConfig {
                cut_selection: None,
                budget: None,
                cut_activity_tolerance: 0.0,
                basis_activity_window: crate::basis_reconstruct::DEFAULT_BASIS_ACTIVITY_WINDOW,
                warm_start_cuts: 0,
                risk_measures: vec![RiskMeasure::Expectation; n_stages],
            },
            events: EventConfig {
                event_sender: None,
                checkpoint_interval: None,
                shutdown_flag: None,
                export_states: false,
            },
        };

        let mut solver = MockSolver::with_fixed(100.0);
        let comm = StubComm;

        let stage_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: &[1usize, 1],
            ncs_max_gen: &[],
            ncs_allow_curtailment: &[],
            discount_factors: &[],
            cumulative_discount_factors: &[],
            stage_lag_transitions: &[],
            noise_group_ids: &[],
            downstream_par_order: 0,
        };
        let result = train(
            &mut solver,
            config,
            &mut fcf,
            &stage_ctx,
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &initial_state,
                inflow_scheme: SamplingScheme::InSample,
                load_scheme: SamplingScheme::InSample,
                ncs_scheme: SamplingScheme::InSample,
                stages: &stages,
                historical_library: None,
                external_inflow_library: None,
                external_load_library: None,
                external_ncs_library: None,
                recent_accum_seed: &[],
                recent_weight_seed: 0.0,
            },
            &comm,
            || Ok(MockSolver::with_fixed(100.0)),
        );

        assert!(result.is_ok(), "train with no event_sender must not panic");
    }

    /// AC: train result `total_time_ms` is greater than 0
    ///
    /// After a successful run, `result.total_time_ms` must be >= 0 and the
    /// `TrainingResult` must be constructible without panicking.
    #[test]
    fn ac_total_time_ms_is_non_negative() {
        let n_stages = 2;
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template(indexer.n_state); n_stages];
        let base_rows = vec![2usize; n_stages];
        let initial_state = vec![0.0_f64; indexer.n_state];
        let stochastic = make_stochastic_context(n_stages, 1);
        let stages = make_stages(n_stages);
        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let mut fcf = make_fcf(n_stages, indexer.n_state, 1, 10);

        let config = TrainingConfig {
            loop_config: LoopConfig {
                forward_passes: 1,
                max_iterations: 1,
                start_iteration: 0,
                n_fwd_threads: 1,
                max_blocks: 1,
                stopping_rules: iteration_limit_rules(1),
            },
            cut_management: CutManagementConfig {
                cut_selection: None,
                budget: None,
                cut_activity_tolerance: 0.0,
                basis_activity_window: crate::basis_reconstruct::DEFAULT_BASIS_ACTIVITY_WINDOW,
                warm_start_cuts: 0,
                risk_measures: vec![RiskMeasure::Expectation; n_stages],
            },
            events: EventConfig {
                event_sender: None,
                checkpoint_interval: None,
                shutdown_flag: None,
                export_states: false,
            },
        };

        let mut solver = MockSolver::with_fixed(100.0);
        let comm = StubComm;

        let stage_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: &[1usize, 1],
            ncs_max_gen: &[],
            ncs_allow_curtailment: &[],
            discount_factors: &[],
            cumulative_discount_factors: &[],
            stage_lag_transitions: &[],
            noise_group_ids: &[],
            downstream_par_order: 0,
        };
        let result = train(
            &mut solver,
            config,
            &mut fcf,
            &stage_ctx,
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &initial_state,
                inflow_scheme: SamplingScheme::InSample,
                load_scheme: SamplingScheme::InSample,
                ncs_scheme: SamplingScheme::InSample,
                stages: &stages,
                historical_library: None,
                external_inflow_library: None,
                external_load_library: None,
                external_ncs_library: None,
                recent_accum_seed: &[],
                recent_weight_seed: 0.0,
            },
            &comm,
            || Ok(MockSolver::with_fixed(100.0)),
        )
        .unwrap();

        assert!(result.error.is_none(), "expected no error");
        assert!(
            result.result.total_time_ms > 0,
            "total_time_ms must be > 0, got {}",
            result.result.total_time_ms,
        );
    }

    /// `cut_selection_none_skips_step`
    ///
    /// Given `train` with `cut_selection: None` running for 5 iterations, then
    /// no `PolicySelectionComplete` event is emitted.
    #[test]
    fn cut_selection_none_skips_step() {
        let n_stages = 2;
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template(indexer.n_state); n_stages];
        let base_rows = vec![2usize; n_stages];
        let initial_state = vec![0.0_f64; indexer.n_state];
        let stochastic = make_stochastic_context(n_stages, 1);
        let stages = make_stages(n_stages);
        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let mut fcf = make_fcf(n_stages, indexer.n_state, 1, 10);

        let (tx, rx) = mpsc::channel::<TrainingEvent>();

        let config = TrainingConfig {
            loop_config: LoopConfig {
                forward_passes: 1,
                max_iterations: 10,
                start_iteration: 0,
                n_fwd_threads: 1,
                max_blocks: 1,
                stopping_rules: iteration_limit_rules(5),
            },
            cut_management: CutManagementConfig {
                cut_selection: None,
                budget: None,
                cut_activity_tolerance: 0.0,
                basis_activity_window: crate::basis_reconstruct::DEFAULT_BASIS_ACTIVITY_WINDOW,
                warm_start_cuts: 0,
                risk_measures: vec![RiskMeasure::Expectation; n_stages],
            },
            events: EventConfig {
                event_sender: Some(tx),
                checkpoint_interval: None,
                shutdown_flag: None,
                export_states: false,
            },
        };

        let mut solver = MockSolver::with_fixed(100.0);
        let comm = StubComm;

        let stage_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: &[1usize, 1],
            ncs_max_gen: &[],
            ncs_allow_curtailment: &[],
            discount_factors: &[],
            cumulative_discount_factors: &[],
            stage_lag_transitions: &[],
            noise_group_ids: &[],
            downstream_par_order: 0,
        };
        train(
            &mut solver,
            config,
            &mut fcf,
            &stage_ctx,
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &initial_state,
                inflow_scheme: SamplingScheme::InSample,
                load_scheme: SamplingScheme::InSample,
                ncs_scheme: SamplingScheme::InSample,
                stages: &stages,
                historical_library: None,
                external_inflow_library: None,
                external_load_library: None,
                external_ncs_library: None,
                recent_accum_seed: &[],
                recent_weight_seed: 0.0,
            },
            &comm,
            || Ok(MockSolver::with_fixed(100.0)),
        )
        .unwrap();

        let events: Vec<TrainingEvent> = rx.try_iter().collect();
        let cut_sel_count = events
            .iter()
            .filter(|e| matches!(e, TrainingEvent::PolicySelectionComplete { .. }))
            .count();

        assert_eq!(
            cut_sel_count, 0,
            "expected no PolicySelectionComplete events with cut_selection: None"
        );
    }

    /// `cut_selection_level1_runs_at_frequency`
    ///
    /// Given `train` with `cut_selection: Some(Level1 { threshold: 0,
    /// check_frequency: 3 })` running for 5 iterations, then
    /// `PolicySelectionComplete` is emitted exactly once (at iteration 3).
    #[test]
    fn cut_selection_level1_runs_at_frequency() {
        use crate::cut_selection::CutSelectionStrategy;

        let n_stages = 2;
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template(indexer.n_state); n_stages];
        let base_rows = vec![2usize; n_stages];
        let initial_state = vec![0.0_f64; indexer.n_state];
        let stochastic = make_stochastic_context(n_stages, 1);
        let stages = make_stages(n_stages);
        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let mut fcf = make_fcf(n_stages, indexer.n_state, 1, 10);

        let (tx, rx) = mpsc::channel::<TrainingEvent>();

        let config = TrainingConfig {
            loop_config: LoopConfig {
                forward_passes: 1,
                max_iterations: 10,
                start_iteration: 0,
                n_fwd_threads: 1,
                max_blocks: 1,
                stopping_rules: iteration_limit_rules(5),
            },
            cut_management: CutManagementConfig {
                cut_selection: Some(CutSelectionStrategy::Level1 {
                    threshold: 0,
                    check_frequency: 3,
                }),
                budget: None,
                cut_activity_tolerance: 0.0,
                basis_activity_window: crate::basis_reconstruct::DEFAULT_BASIS_ACTIVITY_WINDOW,
                warm_start_cuts: 0,
                risk_measures: vec![RiskMeasure::Expectation; n_stages],
            },
            events: EventConfig {
                event_sender: Some(tx),
                checkpoint_interval: None,
                shutdown_flag: None,
                export_states: false,
            },
        };

        let mut solver = MockSolver::with_fixed(100.0);
        let comm = StubComm;

        let stage_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: &[1usize, 1],
            ncs_max_gen: &[],
            ncs_allow_curtailment: &[],
            discount_factors: &[],
            cumulative_discount_factors: &[],
            stage_lag_transitions: &[],
            noise_group_ids: &[],
            downstream_par_order: 0,
        };
        train(
            &mut solver,
            config,
            &mut fcf,
            &stage_ctx,
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &initial_state,
                inflow_scheme: SamplingScheme::InSample,
                load_scheme: SamplingScheme::InSample,
                ncs_scheme: SamplingScheme::InSample,
                stages: &stages,
                historical_library: None,
                external_inflow_library: None,
                external_load_library: None,
                external_ncs_library: None,
                recent_accum_seed: &[],
                recent_weight_seed: 0.0,
            },
            &comm,
            || Ok(MockSolver::with_fixed(100.0)),
        )
        .unwrap();

        let events: Vec<TrainingEvent> = rx.try_iter().collect();
        let sel_events: Vec<&TrainingEvent> = events
            .iter()
            .filter(|e| matches!(e, TrainingEvent::PolicySelectionComplete { .. }))
            .collect();

        assert_eq!(
            sel_events.len(),
            1,
            "expected exactly 1 PolicySelectionComplete event for check_frequency=3 over 5 iterations"
        );

        // Verify the event was emitted at iteration 3.
        let TrainingEvent::PolicySelectionComplete { iteration, .. } = sel_events[0] else {
            panic!("wrong variant");
        };
        assert_eq!(
            *iteration, 3,
            "PolicySelectionComplete must fire at iteration 3"
        );
    }

    /// `cut_selection_deactivates_inactive_cuts`
    ///
    /// Stage 0 is exempt from cut selection because its cuts have no
    /// backward-pass activity tracking. With a 2-stage system, only stage 0
    /// has cuts, so `rows_deactivated` must be 0.
    #[test]
    fn cut_selection_stage0_exempt_preserves_cuts() {
        use crate::cut_selection::CutSelectionStrategy;

        let n_stages = 2;
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template(indexer.n_state); n_stages];
        let base_rows = vec![2usize; n_stages];
        let initial_state = vec![0.0_f64; indexer.n_state];
        let stochastic = make_stochastic_context(n_stages, 1);
        let stages = make_stages(n_stages);
        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let mut fcf = make_fcf(n_stages, indexer.n_state, 1, 10);

        let (tx, rx) = mpsc::channel::<TrainingEvent>();

        let config = TrainingConfig {
            loop_config: LoopConfig {
                forward_passes: 1,
                max_iterations: 10,
                start_iteration: 0,
                n_fwd_threads: 1,
                max_blocks: 1,
                stopping_rules: iteration_limit_rules(2),
            },
            cut_management: CutManagementConfig {
                cut_selection: Some(CutSelectionStrategy::Level1 {
                    threshold: 0,
                    check_frequency: 2,
                }),
                budget: None,
                cut_activity_tolerance: 0.0,
                basis_activity_window: crate::basis_reconstruct::DEFAULT_BASIS_ACTIVITY_WINDOW,
                warm_start_cuts: 0,
                risk_measures: vec![RiskMeasure::Expectation; n_stages],
            },
            events: EventConfig {
                event_sender: Some(tx),
                checkpoint_interval: None,
                shutdown_flag: None,
                export_states: false,
            },
        };

        let mut solver = MockSolver::with_fixed(100.0);
        let comm = StubComm;

        let stage_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: &[1usize, 1],
            ncs_max_gen: &[],
            ncs_allow_curtailment: &[],
            discount_factors: &[],
            cumulative_discount_factors: &[],
            stage_lag_transitions: &[],
            noise_group_ids: &[],
            downstream_par_order: 0,
        };
        train(
            &mut solver,
            config,
            &mut fcf,
            &stage_ctx,
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &initial_state,
                inflow_scheme: SamplingScheme::InSample,
                load_scheme: SamplingScheme::InSample,
                ncs_scheme: SamplingScheme::InSample,
                stages: &stages,
                historical_library: None,
                external_inflow_library: None,
                external_load_library: None,
                external_ncs_library: None,
                recent_accum_seed: &[],
                recent_weight_seed: 0.0,
            },
            &comm,
            || Ok(MockSolver::with_fixed(100.0)),
        )
        .unwrap();

        let events: Vec<TrainingEvent> = rx.try_iter().collect();
        let sel_events: Vec<&TrainingEvent> = events
            .iter()
            .filter(|e| matches!(e, TrainingEvent::PolicySelectionComplete { .. }))
            .collect();

        assert_eq!(
            sel_events.len(),
            1,
            "expected exactly 1 PolicySelectionComplete event at iteration 2"
        );

        let TrainingEvent::PolicySelectionComplete {
            iteration,
            rows_deactivated,
            per_stage,
            ..
        } = sel_events[0]
        else {
            panic!("wrong variant");
        };

        assert_eq!(*iteration, 2, "selection must fire at iteration 2");
        assert_eq!(
            *rows_deactivated, 0,
            "stage 0 is exempt from cut selection, so no cuts should be deactivated"
        );
        // Verify per-stage records are populated and stage 0 is exempt.
        assert!(
            !per_stage.is_empty(),
            "per_stage must contain at least the stage 0 record"
        );
        assert_eq!(per_stage[0].stage, 0, "first record must be stage 0");
        assert_eq!(
            per_stage[0].rows_deactivated, 0,
            "stage 0 must have zero deactivations"
        );
    }

    /// `existing_train_tests_pass_with_none`
    ///
    /// Verify backward compatibility: calling `train` with `cut_selection:
    /// None` produces the same result as before. This is
    /// implicitly verified by the existing `ac_train_completes_with_iteration_limit`
    /// test. This test is an explicit additional check with an explicit `None`.
    #[test]
    fn existing_train_tests_pass_with_none() {
        let n_stages = 2;
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template(indexer.n_state); n_stages];
        let base_rows = vec![2usize; n_stages];
        let initial_state = vec![0.0_f64; indexer.n_state];
        let stochastic = make_stochastic_context(n_stages, 1);
        let stages = make_stages(n_stages);
        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let mut fcf = make_fcf(n_stages, indexer.n_state, 1, 10);

        let config = TrainingConfig {
            loop_config: LoopConfig {
                forward_passes: 1,
                max_iterations: 3,
                start_iteration: 0,
                n_fwd_threads: 1,
                max_blocks: 1,
                stopping_rules: iteration_limit_rules(3),
            },
            cut_management: CutManagementConfig {
                cut_selection: None,
                budget: None,
                cut_activity_tolerance: 0.0,
                basis_activity_window: crate::basis_reconstruct::DEFAULT_BASIS_ACTIVITY_WINDOW,
                warm_start_cuts: 0,
                risk_measures: vec![RiskMeasure::Expectation; n_stages],
            },
            events: EventConfig {
                event_sender: None,
                checkpoint_interval: None,
                shutdown_flag: None,
                export_states: false,
            },
        };

        let mut solver = MockSolver::with_fixed(100.0);
        let comm = StubComm;

        let stage_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: &[1usize, 1],
            ncs_max_gen: &[],
            ncs_allow_curtailment: &[],
            discount_factors: &[],
            cumulative_discount_factors: &[],
            stage_lag_transitions: &[],
            noise_group_ids: &[],
            downstream_par_order: 0,
        };
        let result = train(
            &mut solver,
            config,
            &mut fcf,
            &stage_ctx,
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &initial_state,
                inflow_scheme: SamplingScheme::InSample,
                load_scheme: SamplingScheme::InSample,
                ncs_scheme: SamplingScheme::InSample,
                stages: &stages,
                historical_library: None,
                external_inflow_library: None,
                external_load_library: None,
                external_ncs_library: None,
                recent_accum_seed: &[],
                recent_weight_seed: 0.0,
            },
            &comm,
            || Ok(MockSolver::with_fixed(100.0)),
        )
        .unwrap();

        assert!(result.error.is_none(), "expected no error");
        assert_eq!(result.result.iterations, 3);
        assert_eq!(result.result.reason, "iteration_limit");
    }

    /// AC: `train_partial_result_on_mid_iteration_failure`
    ///
    /// Given a mock solver that fails on the 3rd solve call (which occurs
    /// during iteration 1 since each iteration performs multiple solves),
    /// when `train` is called, it returns `Ok(TrainingOutcome)` with
    /// `error: Some(...)`, `result.iterations == 0` (no completed iterations),
    /// `result.reason == "error"`, and `solver_stats_log` containing entries
    /// from the phases that completed before the error.
    #[test]
    fn ac_train_partial_result_on_mid_iteration_failure() {
        let n_stages = 2;
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template(indexer.n_state); n_stages];
        let base_rows = vec![2usize; n_stages];
        let initial_state = vec![0.0_f64; indexer.n_state];
        let stochastic = make_stochastic_context(n_stages, 1);
        let stages = make_stages(n_stages);
        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let mut fcf = make_fcf(n_stages, indexer.n_state, 1, 10);

        let (tx, rx) = mpsc::channel::<TrainingEvent>();

        let config = TrainingConfig {
            loop_config: LoopConfig {
                forward_passes: 1,
                max_iterations: 5,
                start_iteration: 0,
                n_fwd_threads: 1,
                max_blocks: 1,
                stopping_rules: iteration_limit_rules(5),
            },
            cut_management: CutManagementConfig {
                cut_selection: None,
                budget: None,
                cut_activity_tolerance: 0.0,
                basis_activity_window: crate::basis_reconstruct::DEFAULT_BASIS_ACTIVITY_WINDOW,
                warm_start_cuts: 0,
                risk_measures: vec![RiskMeasure::Expectation; n_stages],
            },
            events: EventConfig {
                event_sender: Some(tx),
                checkpoint_interval: None,
                shutdown_flag: None,
                export_states: false,
            },
        };

        // Mock solver that fails on the Nth call. With 2 stages and 1 forward
        // pass, the forward pass solves 2 LPs (stage 0, stage 1). A failure
        // on the 1st call (index 0) means failure in the forward pass of
        // iteration 1.
        let mut solver = MockSolver::infeasible();
        let comm = StubComm;

        let stage_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: &[1usize, 1],
            ncs_max_gen: &[],
            ncs_allow_curtailment: &[],
            discount_factors: &[],
            cumulative_discount_factors: &[],
            stage_lag_transitions: &[],
            noise_group_ids: &[],
            downstream_par_order: 0,
        };
        let outcome = train(
            &mut solver,
            config,
            &mut fcf,
            &stage_ctx,
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &initial_state,
                inflow_scheme: SamplingScheme::InSample,
                load_scheme: SamplingScheme::InSample,
                ncs_scheme: SamplingScheme::InSample,
                stages: &stages,
                historical_library: None,
                external_inflow_library: None,
                external_load_library: None,
                external_ncs_library: None,
                recent_accum_seed: &[],
                recent_weight_seed: 0.0,
            },
            &comm,
            || Ok(MockSolver::infeasible()),
        )
        .unwrap();

        // Verify partial result semantics.
        assert!(outcome.error.is_some(), "expected error in TrainingOutcome");
        assert_eq!(
            outcome.result.iterations, 0,
            "no iterations should have completed (failure in iteration 1)"
        );
        assert_eq!(outcome.result.reason, "error");
        assert!(
            outcome.result.total_time_ms > 0,
            "total_time_ms must be > 0"
        );

        // Verify TrainingFinished event was emitted with reason "error".
        let events: Vec<TrainingEvent> = rx.try_iter().collect();
        let finished = events
            .iter()
            .find(|e| matches!(e, TrainingEvent::TrainingFinished { .. }));
        assert!(
            finished.is_some(),
            "TrainingFinished event must be emitted even on error"
        );
        if let Some(TrainingEvent::TrainingFinished { reason, .. }) = finished {
            assert_eq!(reason, "error", "TrainingFinished reason must be 'error'");
        }
    }

    /// When `start_iteration = 3` and `max_iterations = 5`, the training loop
    /// executes exactly 2 iterations (4 and 5) and reports `iterations = 5`.
    #[test]
    fn start_iteration_resumes_from_offset() {
        let n_stages = 2;
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template(indexer.n_state); n_stages];
        let base_rows = vec![2usize; n_stages];
        let initial_state = vec![0.0_f64; indexer.n_state];
        let stochastic = make_stochastic_context(n_stages, 1);
        let stages = make_stages(n_stages);
        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let mut fcf = make_fcf(n_stages, indexer.n_state, 1, 10);

        let config = TrainingConfig {
            loop_config: LoopConfig {
                forward_passes: 1,
                max_iterations: 5,
                start_iteration: 3,
                n_fwd_threads: 1,
                max_blocks: 1,
                stopping_rules: iteration_limit_rules(5),
            },
            cut_management: CutManagementConfig {
                cut_selection: None,
                budget: None,
                cut_activity_tolerance: 0.0,
                basis_activity_window: crate::basis_reconstruct::DEFAULT_BASIS_ACTIVITY_WINDOW,
                warm_start_cuts: 0,
                risk_measures: vec![RiskMeasure::Expectation; n_stages],
            },
            events: EventConfig {
                event_sender: None,
                checkpoint_interval: None,
                shutdown_flag: None,
                export_states: false,
            },
        };

        let mut solver = MockSolver::with_fixed(100.0);
        let comm = StubComm;

        let stage_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: &[1usize, 1],
            ncs_max_gen: &[],
            ncs_allow_curtailment: &[],
            discount_factors: &[],
            cumulative_discount_factors: &[],
            stage_lag_transitions: &[],
            noise_group_ids: &[],
            downstream_par_order: 0,
        };
        let outcome = train(
            &mut solver,
            config,
            &mut fcf,
            &stage_ctx,
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &initial_state,
                inflow_scheme: SamplingScheme::InSample,
                load_scheme: SamplingScheme::InSample,
                ncs_scheme: SamplingScheme::InSample,
                stages: &stages,
                historical_library: None,
                external_inflow_library: None,
                external_load_library: None,
                external_ncs_library: None,
                recent_accum_seed: &[],
                recent_weight_seed: 0.0,
            },
            &comm,
            || Ok(MockSolver::with_fixed(100.0)),
        )
        .unwrap();

        assert_eq!(
            outcome.result.iterations, 5,
            "iterations must report the absolute iteration number (5), not the delta (2)"
        );
        assert_eq!(outcome.result.reason, "iteration_limit");
    }

    /// When `start_iteration >= max_iterations`, the training loop executes zero
    /// iterations and returns immediately with `iterations = start_iteration`.
    #[test]
    fn start_iteration_at_or_beyond_max_runs_zero_iterations() {
        let n_stages = 2;
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template(indexer.n_state); n_stages];
        let base_rows = vec![2usize; n_stages];
        let initial_state = vec![0.0_f64; indexer.n_state];
        let stochastic = make_stochastic_context(n_stages, 1);
        let stages = make_stages(n_stages);
        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let mut fcf = make_fcf(n_stages, indexer.n_state, 1, 10);

        let config = TrainingConfig {
            loop_config: LoopConfig {
                forward_passes: 1,
                max_iterations: 5,
                start_iteration: 5,
                n_fwd_threads: 1,
                max_blocks: 1,
                stopping_rules: iteration_limit_rules(5),
            },
            cut_management: CutManagementConfig {
                cut_selection: None,
                budget: None,
                cut_activity_tolerance: 0.0,
                basis_activity_window: crate::basis_reconstruct::DEFAULT_BASIS_ACTIVITY_WINDOW,
                warm_start_cuts: 0,
                risk_measures: vec![RiskMeasure::Expectation; n_stages],
            },
            events: EventConfig {
                event_sender: None,
                checkpoint_interval: None,
                shutdown_flag: None,
                export_states: false,
            },
        };

        let mut solver = MockSolver::with_fixed(100.0);
        let comm = StubComm;

        let stage_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: &[1usize, 1],
            ncs_max_gen: &[],
            ncs_allow_curtailment: &[],
            discount_factors: &[],
            cumulative_discount_factors: &[],
            stage_lag_transitions: &[],
            noise_group_ids: &[],
            downstream_par_order: 0,
        };
        let outcome = train(
            &mut solver,
            config,
            &mut fcf,
            &stage_ctx,
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &initial_state,
                inflow_scheme: SamplingScheme::InSample,
                load_scheme: SamplingScheme::InSample,
                ncs_scheme: SamplingScheme::InSample,
                stages: &stages,
                historical_library: None,
                external_inflow_library: None,
                external_load_library: None,
                external_ncs_library: None,
                recent_accum_seed: &[],
                recent_weight_seed: 0.0,
            },
            &comm,
            || Ok(MockSolver::with_fixed(100.0)),
        )
        .unwrap();

        assert_eq!(
            outcome.result.iterations, 5,
            "iterations must equal start_iteration when no loop iterations execute"
        );
        assert_eq!(
            outcome.result.reason, "iteration_limit",
            "reason should be iteration_limit when loop range is empty"
        );
    }

    // ── broadcast_basis_cache unit tests ─────────────────────────────────────

    /// AC: `broadcast_basis_cache` returns scenario 0's bases, not scenario N-1's.
    ///
    /// Constructs a `BasisStore` with 2 scenarios and 3 stages. Scenario 0 is
    /// populated with known distinct values; scenario 1 is populated with
    /// different values. The helper must return scenario 0's bases on
    /// `LocalBackend` (single-rank, no broadcast).
    #[test]
    fn ac_broadcast_basis_cache_uses_scenario_0_not_last() {
        use super::broadcast_basis_cache;
        use crate::workspace::BasisStore;

        let num_scenarios = 4; // simulates total_forward_passes=4, num_ranks=1
        let num_stages = 3;
        let mut store = BasisStore::new(num_scenarios, num_stages);

        // Populate scenario 0 with col_status=[10,20], row_status=[30].
        for t in 0..num_stages {
            // test shim: zero metadata is acceptable for tests exercising broadcast path
            *store.get_mut(0, t) = Some(crate::workspace::CapturedBasis {
                basis: Basis {
                    col_status: vec![10_i32 + t as i32, 20_i32 + t as i32],
                    row_status: vec![30_i32 + t as i32],
                },
                base_row_count: 0,
                cut_row_slots: Vec::new(),
                state_at_capture: Vec::new(),
            });
        }

        // Populate scenario 3 (last) with completely different values.
        for t in 0..num_stages {
            // test shim: zero metadata is acceptable for tests exercising broadcast path
            *store.get_mut(3, t) = Some(crate::workspace::CapturedBasis {
                basis: Basis {
                    col_status: vec![99_i32, 88_i32],
                    row_status: vec![77_i32],
                },
                base_row_count: 0,
                cut_row_slots: Vec::new(),
                state_at_capture: Vec::new(),
            });
        }

        let comm = StubComm; // single-rank, no broadcast
        let cache = broadcast_basis_cache(&store, num_stages, &comm).unwrap();

        assert_eq!(cache.len(), num_stages);
        for (t, entry) in cache.iter().enumerate() {
            let captured = entry
                .as_ref()
                .expect("stage {t} must have a captured basis");
            assert_eq!(
                captured.basis.col_status,
                vec![10_i32 + t as i32, 20_i32 + t as i32],
                "stage {t} col_status must come from scenario 0, not scenario 3"
            );
            assert_eq!(
                captured.basis.row_status,
                vec![30_i32 + t as i32],
                "stage {t} row_status must come from scenario 0, not scenario 3"
            );
        }
    }

    /// AC: `broadcast_basis_cache` handles `None` slots correctly.
    ///
    /// When scenario 0 has no basis stored for some stages (e.g. training
    /// stopped before any forward pass completed), those stages must be `None`
    /// in the returned cache.
    #[test]
    fn ac_broadcast_basis_cache_none_slots_preserved() {
        use super::broadcast_basis_cache;
        use crate::workspace::BasisStore;

        let num_stages = 2;
        // Scenario 0 is left unpopulated (all None).
        let store = BasisStore::new(1, num_stages);

        let comm = StubComm;
        let cache = broadcast_basis_cache(&store, num_stages, &comm).unwrap();

        assert_eq!(cache.len(), num_stages);
        for t in 0..num_stages {
            assert!(
                cache[t].is_none(),
                "stage {t} must be None when basis store has no entry for scenario 0"
            );
        }
    }

    /// `broadcast_basis_cache` with `comm.size() == 1` must clone
    /// the full `CapturedBasis` including metadata (`cut_row_slots`,
    /// `state_at_capture`, `base_row_count`), not just the bare basis body.
    #[test]
    fn broadcast_basis_cache_single_rank_preserves_metadata() {
        use super::broadcast_basis_cache;
        use crate::workspace::{BasisStore, CapturedBasis};

        let num_stages = 2;
        let mut store = BasisStore::new(1, num_stages);

        // Populate stage 0 with non-empty metadata.
        *store.get_mut(0, 0) = Some(CapturedBasis {
            basis: Basis {
                col_status: vec![1_i32, 2_i32],
                row_status: vec![3_i32, 4_i32, 5_i32],
            },
            base_row_count: 2,
            cut_row_slots: vec![10_u32, 11_u32, 12_u32],
            state_at_capture: vec![1.5_f64, 2.5_f64],
        });
        // Stage 1 left None.

        let comm = StubComm; // size == 1
        let cache = broadcast_basis_cache(&store, num_stages, &comm).unwrap();

        assert_eq!(cache.len(), num_stages);
        let cb = cache[0].as_ref().expect("stage 0 must have captured basis");
        assert_eq!(
            cb.cut_row_slots.len(),
            3,
            "single-rank path must preserve cut_row_slots"
        );
        assert_eq!(cb.base_row_count, 2, "base_row_count must be preserved");
        assert_eq!(
            cb.state_at_capture,
            vec![1.5_f64, 2.5_f64],
            "state_at_capture must be preserved"
        );
        assert!(cache[1].is_none(), "stage 1 must remain None");
    }

    /// A discriminated payload stored by `MultiRankMockComm`.
    ///
    /// `broadcast_basis_cache` issues exactly four `broadcast` calls:
    /// two `i32` calls (length then payload) and two `f64` calls (length then
    /// payload). We store each call as a typed variant so that rank 1 can
    /// deserialize without any `unsafe` code.
    #[derive(Clone)]
    enum MockPayload {
        Ints(Vec<i32>),
        Floats(Vec<f64>),
    }

    /// A multi-rank mock communicator that simulates 2-rank broadcasts.
    ///
    /// On rank 0, each `broadcast<T>` call records the outgoing buffer as a
    /// `MockPayload` variant. On rank 1, each `broadcast<T>` call pops the
    /// next recorded entry and copies it into the caller's mutable slice.
    ///
    /// Only `T = i32` and `T = f64` are supported (matching the wire types
    /// used by `broadcast_basis_cache`). `allgatherv` and `allreduce` are not
    /// called by that function and are left as `unreachable!()`.
    ///
    /// `Mutex` is used instead of `RefCell` to satisfy the `Sync` bound
    /// required by `Communicator`.
    struct MultiRankMockComm {
        rank: usize,
        queue: std::sync::Mutex<std::collections::VecDeque<MockPayload>>,
    }

    impl MultiRankMockComm {
        fn new_root() -> Self {
            Self {
                rank: 0,
                queue: std::sync::Mutex::new(std::collections::VecDeque::new()),
            }
        }

        /// Build a rank-1 peer by snapshotting the root's recorded queue.
        ///
        /// Must be called **after** `broadcast_basis_cache` has returned on
        /// rank 0 so the snapshot contains all four recorded payloads.
        fn new_peer(root: &MultiRankMockComm) -> Self {
            Self {
                rank: 1,
                queue: std::sync::Mutex::new(root.queue.lock().unwrap().clone()),
            }
        }

        /// Build a rank-1 peer from an explicit replay queue.
        ///
        /// Used by corruption tests that tamper with the recorded payloads
        /// before replaying them to rank 1.
        fn new_peer_from_queue(queue: std::collections::VecDeque<MockPayload>) -> Self {
            Self {
                rank: 1,
                queue: std::sync::Mutex::new(queue),
            }
        }

        /// Extract a snapshot of the recorded queue (for corruption tests).
        fn snapshot(&self) -> std::collections::VecDeque<MockPayload> {
            self.queue.lock().unwrap().clone()
        }
    }

    impl Communicator for MultiRankMockComm {
        fn allgatherv<T: CommData>(
            &self,
            _send: &[T],
            _recv: &mut [T],
            _counts: &[usize],
            _displs: &[usize],
        ) -> Result<(), CommError> {
            unreachable!("broadcast_basis_cache does not call allgatherv")
        }

        fn allreduce<T: CommData>(
            &self,
            _send: &[T],
            _recv: &mut [T],
            _op: ReduceOp,
        ) -> Result<(), CommError> {
            unreachable!("broadcast_basis_cache does not call allreduce")
        }

        fn broadcast<T: CommData>(&self, buf: &mut [T], root: usize) -> Result<(), CommError> {
            // We need MockRecord to dispatch safely. We cannot add that bound
            // to the Communicator trait, so we use a private helper that
            // downcasts via a local trait object. This is the only clean safe
            // approach without unsafe in a forbid-unsafe crate.
            //
            // The actual dispatch is done by calling into a monomorphic helper
            // through a function pointer selected at the call site where T is
            // known. We implement this by defining a local function that takes
            // &dyn Any and casts it:
            self.broadcast_typed(buf, root)
        }

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

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

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

        fn abort(&self, code: i32) -> ! {
            std::process::exit(code)
        }
    }

    impl MultiRankMockComm {
        // Result<(), CommError> is required to match the Communicator::broadcast
        // return type this delegates to, even though the body always returns Ok(()).
        #[allow(clippy::unnecessary_wraps)]
        fn broadcast_typed<T: CommData>(
            &self,
            buf: &mut [T],
            _root: usize,
        ) -> Result<(), CommError> {
            use std::any::Any;
            // T: CommData implies T: 'static + Copy, so Box<T>: Any.
            // We identify the concrete type by boxing a probe value (Default
            // if the slice is empty) and downcasting. No raw pointer casts.
            let probe: Box<dyn Any> = Box::new(T::default());

            if probe.downcast_ref::<i32>().is_some() {
                // T == i32
                if self.rank == 0 {
                    let ints: Vec<i32> = buf
                        .iter()
                        .map(|v| {
                            *Box::<dyn Any>::from(Box::new(*v))
                                .downcast::<i32>()
                                .expect("T proved i32 above")
                        })
                        .collect();
                    self.queue
                        .lock()
                        .unwrap()
                        .push_back(MockPayload::Ints(ints));
                } else {
                    let payload = self
                        .queue
                        .lock()
                        .unwrap()
                        .pop_front()
                        .expect("MultiRankMockComm: no payload to replay for i32 broadcast");
                    let MockPayload::Ints(src) = payload else {
                        panic!("MultiRankMockComm: expected Ints payload for i32 broadcast");
                    };
                    assert_eq!(src.len(), buf.len(), "i32 replay length mismatch");
                    for (dst, v) in buf.iter_mut().zip(src.iter()) {
                        let boxed: Box<dyn Any> = Box::new(*v);
                        *dst = *boxed.downcast::<T>().expect("T proved i32 above");
                    }
                }
            } else if probe.downcast_ref::<f64>().is_some() {
                // T == f64
                if self.rank == 0 {
                    let floats: Vec<f64> = buf
                        .iter()
                        .map(|v| {
                            *Box::<dyn Any>::from(Box::new(*v))
                                .downcast::<f64>()
                                .expect("T proved f64 above")
                        })
                        .collect();
                    self.queue
                        .lock()
                        .unwrap()
                        .push_back(MockPayload::Floats(floats));
                } else {
                    let payload = self
                        .queue
                        .lock()
                        .unwrap()
                        .pop_front()
                        .expect("MultiRankMockComm: no payload to replay for f64 broadcast");
                    let MockPayload::Floats(src) = payload else {
                        panic!("MultiRankMockComm: expected Floats payload for f64 broadcast");
                    };
                    assert_eq!(src.len(), buf.len(), "f64 replay length mismatch");
                    for (dst, v) in buf.iter_mut().zip(src.iter()) {
                        let boxed: Box<dyn Any> = Box::new(*v);
                        *dst = *boxed.downcast::<T>().expect("T proved f64 above");
                    }
                }
            } else {
                panic!("MultiRankMockComm: unsupported broadcast type (expected i32 or f64)");
            }
            Ok(())
        }
    }

    /// `broadcast_basis_cache` with `size=2` must transmit the
    /// full `CapturedBasis` metadata (`cut_row_slots`, `state_at_capture`,
    /// `base_row_count`) to rank 1. The `MultiRankMockComm` pair records
    /// rank-0's four broadcasts and replays them exactly to rank 1.
    #[test]
    fn broadcast_basis_cache_multi_rank_round_trips_full_metadata() {
        use super::broadcast_basis_cache;
        use crate::workspace::{BasisStore, CapturedBasis};

        let mut store = BasisStore::new(1, 2);
        *store.get_mut(0, 0) = Some(CapturedBasis {
            basis: Basis {
                col_status: vec![1_i32, 2_i32, 3_i32],
                row_status: vec![10_i32, 20_i32],
            },
            base_row_count: 4,
            cut_row_slots: vec![10_u32, 11_u32, 12_u32],
            state_at_capture: vec![1.5_f64, 2.5_f64],
        });
        // Stage 1 left None.

        // Step 1: run broadcast_basis_cache from rank-0's perspective.
        // This populates the mock's recorded buffers.
        let root_comm = MultiRankMockComm::new_root();
        let _cache_rank0 = broadcast_basis_cache(&store, 2, &root_comm).unwrap();

        // Step 2: build a rank-1 peer that replays rank-0's recorded buffers.
        let peer_comm = MultiRankMockComm::new_peer(&root_comm);
        // Rank 1's basis_store is empty — all data must come from the broadcast.
        let empty_store = BasisStore::new(1, 2);
        let cache = broadcast_basis_cache(&empty_store, 2, &peer_comm).unwrap();

        assert_eq!(cache.len(), 2);
        let cb0 = cache[0]
            .as_ref()
            .expect("stage 0 must deserialise into CapturedBasis on rank 1");
        assert_eq!(
            cb0.basis.col_status,
            vec![1_i32, 2_i32, 3_i32],
            "col_status must round-trip"
        );
        assert_eq!(
            cb0.basis.row_status,
            vec![10_i32, 20_i32],
            "row_status must round-trip"
        );
        assert_eq!(
            cb0.cut_row_slots,
            vec![10_u32, 11_u32, 12_u32],
            "cut_row_slots must round-trip on non-root rank"
        );
        assert_eq!(
            cb0.state_at_capture,
            vec![1.5_f64, 2.5_f64],
            "state_at_capture must round-trip on non-root rank"
        );
        assert_eq!(
            cb0.base_row_count, 4,
            "base_row_count must round-trip on non-root rank"
        );
        assert!(cache[1].is_none(), "stage 1 had no basis → None");
    }

    /// When rank 0's `CapturedBasis` has an empty `cut_row_slots`
    /// (legitimate for stages that never produced cuts), the round-trip must
    /// produce `cut_row_slots.is_empty()` on rank 1 without error.
    #[test]
    fn broadcast_basis_cache_empty_cut_slots_round_trips_ok() {
        use super::broadcast_basis_cache;
        use crate::workspace::{BasisStore, CapturedBasis};

        let mut store = BasisStore::new(1, 1);
        *store.get_mut(0, 0) = Some(CapturedBasis {
            basis: Basis {
                col_status: vec![5_i32, 6_i32],
                row_status: vec![7_i32],
            },
            base_row_count: 1,
            cut_row_slots: vec![], // deliberately empty
            state_at_capture: vec![3.75_f64],
        });

        let root_comm = MultiRankMockComm::new_root();
        let _ = broadcast_basis_cache(&store, 1, &root_comm).unwrap();

        let peer_comm = MultiRankMockComm::new_peer(&root_comm);
        let empty_store = BasisStore::new(1, 1);
        let cache = broadcast_basis_cache(&empty_store, 1, &peer_comm).unwrap();

        assert_eq!(cache.len(), 1);
        let cb = cache[0]
            .as_ref()
            .expect("stage 0 must be Some after broadcast");
        assert!(
            cb.cut_row_slots.is_empty(),
            "empty cut_row_slots must round-trip without error or panic"
        );
        assert_eq!(
            cb.state_at_capture,
            vec![3.75_f64],
            "state_at_capture must still round-trip when cut_row_slots is empty"
        );
        assert_eq!(cb.base_row_count, 1, "base_row_count must round-trip");
    }

    /// When the i32 broadcast buffer is truncated mid-
    /// `cut_row_slots`, `broadcast_basis_cache` must return
    /// `SddpError::Validation` with a message containing "cut_row_slots" and
    /// the stage index.
    ///
    /// The queue produced by a successful rank-0 run contains four payloads:
    ///   [0] Ints(i32-len-scalar)   — one i32 (the total i32 count)
    ///   [1] Ints(i32-payload)      — all the integer data
    ///   [2] Ints(f64-len-scalar)   — one i32 (the total f64 count)
    ///   [3] Floats(f64-payload)    — all the f64 state data
    ///
    /// To simulate a truncated i32 payload we replace entry [1] with a
    /// shorter `Ints` vector (missing the last cut slot) and patch entry [0]
    /// to reflect the new count, so that rank 1 allocates the right buffer
    /// size but then fails the `cut_row_slots` bounds check.
    #[test]
    fn broadcast_basis_cache_truncated_cut_slots_returns_validation() {
        use super::broadcast_basis_cache;
        use crate::workspace::{BasisStore, CapturedBasis};

        // Build a store with cut_row_slots = [10, 11, 12].
        let mut store = BasisStore::new(1, 1);
        *store.get_mut(0, 0) = Some(CapturedBasis {
            basis: Basis {
                col_status: vec![1_i32],
                row_status: vec![2_i32],
            },
            base_row_count: 1,
            cut_row_slots: vec![10_u32, 11_u32, 12_u32],
            state_at_capture: vec![0.0_f64],
        });

        // Record rank-0 payloads.
        let root_comm = MultiRankMockComm::new_root();
        let _ = broadcast_basis_cache(&store, 1, &root_comm).unwrap();
        let mut snapshot = root_comm.snapshot();

        // snapshot[1] is the Ints(payload) entry. Remove the last i32 value
        // (the last cut slot) to simulate a truncated buffer. Also patch
        // snapshot[0] (the length scalar) to match the reduced count.
        let truncated_len = {
            let entry = snapshot.get_mut(1).expect("i32 payload entry must exist");
            let MockPayload::Ints(ref mut ints) = *entry else {
                panic!("entry [1] must be Ints");
            };
            ints.pop(); // remove one cut slot
            ints.len() as i32
        };
        // Patch the length scalar (entry [0]).
        let len_entry = snapshot.get_mut(0).expect("i32 length entry must exist");
        let MockPayload::Ints(ref mut len_vec) = *len_entry else {
            panic!("entry [0] must be Ints");
        };
        assert_eq!(len_vec.len(), 1, "length entry must hold a single scalar");
        len_vec[0] = truncated_len;

        let peer_comm = MultiRankMockComm::new_peer_from_queue(snapshot);
        let empty_store = BasisStore::new(1, 1);
        let result = broadcast_basis_cache(&empty_store, 1, &peer_comm);

        match result {
            Err(SddpError::Validation(msg)) => {
                assert!(
                    msg.contains("cut_row_slots"),
                    "error message must mention 'cut_row_slots', got: {msg}"
                );
                assert!(
                    msg.contains('0'),
                    "error message must contain stage index 0, got: {msg}"
                );
            }
            other => panic!("expected SddpError::Validation, got: {other:?}"),
        }
    }

    /// When the f64 broadcast buffer is truncated mid-
    /// `state_at_capture`, `broadcast_basis_cache` must return
    /// `SddpError::Validation` with a message containing "state_at_capture"
    /// and the stage index.
    ///
    /// We truncate entry [3] (Floats payload) and patch entry [2] (f64-len
    /// scalar) to match, so rank 1 allocates a shorter f64 buffer and the
    /// `state_at_capture` bounds check fires.
    #[test]
    fn broadcast_basis_cache_truncated_state_returns_validation() {
        use super::broadcast_basis_cache;
        use crate::workspace::{BasisStore, CapturedBasis};

        let mut store = BasisStore::new(1, 1);
        *store.get_mut(0, 0) = Some(CapturedBasis {
            basis: Basis {
                col_status: vec![1_i32],
                row_status: vec![2_i32],
            },
            base_row_count: 1,
            cut_row_slots: vec![],
            state_at_capture: vec![1.0_f64, 2.0_f64, 3.0_f64],
        });

        let root_comm = MultiRankMockComm::new_root();
        let _ = broadcast_basis_cache(&store, 1, &root_comm).unwrap();
        let mut snapshot = root_comm.snapshot();

        // snapshot[3] is the Floats(f64-payload) entry with 3 values.
        // Keep only 1 value so that rank 1's buffer is too short for the
        // state_len=3 embedded in the i32 payload.
        let truncated_f64_len = {
            let entry = snapshot.get_mut(3).expect("f64 payload entry must exist");
            let MockPayload::Floats(ref mut floats) = *entry else {
                panic!("entry [3] must be Floats");
            };
            floats.truncate(1); // keep only 1 f64
            floats.len() as i32
        };
        // Patch entry [2] (f64 length scalar).
        let f64_len_entry = snapshot.get_mut(2).expect("f64 length entry must exist");
        let MockPayload::Ints(ref mut f64_len_vec) = *f64_len_entry else {
            panic!("entry [2] must be Ints (f64 length is broadcast as i32)");
        };
        assert_eq!(
            f64_len_vec.len(),
            1,
            "f64 length entry must hold a single scalar"
        );
        f64_len_vec[0] = truncated_f64_len;

        let peer_comm = MultiRankMockComm::new_peer_from_queue(snapshot);
        let empty_store = BasisStore::new(1, 1);
        let result = broadcast_basis_cache(&empty_store, 1, &peer_comm);

        match result {
            Err(SddpError::Validation(msg)) => {
                assert!(
                    msg.contains("state_at_capture"),
                    "error message must mention 'state_at_capture', got: {msg}"
                );
                assert!(
                    msg.contains('0'),
                    "error message must contain stage index 0, got: {msg}"
                );
            }
            other => panic!("expected SddpError::Validation, got: {other:?}"),
        }
    }

    /// Regression test for the i32-truncation guard in `broadcast_basis_cache`.
    ///
    /// A buffer length that exceeds `i32::MAX` must be rejected by
    /// `checked_broadcast_len` before any MPI broadcast call is issued.
    /// The returned error must be
    /// `SddpError::Communication(CommError::InvalidBufferSize { .. })`
    /// with `actual == (i32::MAX as usize) + 1` and the operation string
    /// `"broadcast_basis_cache_i32"`.
    #[test]
    fn broadcast_basis_cache_rejects_oversized_i32_payload() {
        use super::checked_broadcast_len;

        let oversized: usize = (i32::MAX as usize) + 1;
        let result = checked_broadcast_len(oversized, "broadcast_basis_cache_i32");

        match result {
            Err(SddpError::Communication(CommError::InvalidBufferSize {
                operation,
                expected,
                actual,
            })) => {
                assert_eq!(actual, oversized, "actual must equal the oversized length");
                assert_eq!(
                    expected,
                    i32::MAX as usize,
                    "expected must equal i32::MAX as usize"
                );
                assert_eq!(
                    operation, "broadcast_basis_cache_i32",
                    "operation string must be 'broadcast_basis_cache_i32'"
                );
            }
            other => panic!(
                "expected SddpError::Communication(CommError::InvalidBufferSize {{ .. }}), got: {other:?}"
            ),
        }
    }

    /// AC: `template_bake_event_emitted`
    ///
    /// Verify that `PolicyTemplateBakeComplete` is emitted exactly once per iteration
    /// with the correct `stages_processed` count. Also verifies that
    /// `total_rows_baked > 0` on iteration 2 (because the backward pass on
    /// iteration 1 generates cuts before step 4c runs on that same iteration).
    #[test]
    fn template_bake_event_emitted() {
        let n_stages = 2;
        let indexer = StageIndexer::new(1, 0);
        let templates = vec![minimal_template(indexer.n_state); n_stages];
        let base_rows = vec![2usize; n_stages];
        let initial_state = vec![0.0_f64; indexer.n_state];
        let stochastic = make_stochastic_context(n_stages, 1);
        let stages = make_stages(n_stages);
        let horizon = HorizonMode::Finite {
            num_stages: n_stages,
        };
        let mut fcf = make_fcf(n_stages, indexer.n_state, 1, 10);

        let (tx, rx) = mpsc::channel::<TrainingEvent>();

        let config = TrainingConfig {
            loop_config: LoopConfig {
                forward_passes: 1,
                max_iterations: 10,
                start_iteration: 0,
                n_fwd_threads: 1,
                max_blocks: 1,
                stopping_rules: iteration_limit_rules(2),
            },
            cut_management: CutManagementConfig {
                cut_selection: None,
                budget: None,
                cut_activity_tolerance: 0.0,
                basis_activity_window: crate::basis_reconstruct::DEFAULT_BASIS_ACTIVITY_WINDOW,
                warm_start_cuts: 0,
                risk_measures: vec![RiskMeasure::Expectation; n_stages],
            },
            events: EventConfig {
                event_sender: Some(tx),
                checkpoint_interval: None,
                shutdown_flag: None,
                export_states: false,
            },
        };

        let mut solver = MockSolver::with_fixed(100.0);
        let comm = StubComm;

        let stage_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: &[1usize, 1],
            ncs_max_gen: &[],
            ncs_allow_curtailment: &[],
            discount_factors: &[],
            cumulative_discount_factors: &[],
            stage_lag_transitions: &[],
            noise_group_ids: &[],
            downstream_par_order: 0,
        };

        train(
            &mut solver,
            config,
            &mut fcf,
            &stage_ctx,
            &TrainingContext {
                horizon: &horizon,
                indexer: &indexer,
                inflow_method: &InflowNonNegativityMethod::None,
                stochastic: &stochastic,
                initial_state: &initial_state,
                inflow_scheme: SamplingScheme::InSample,
                load_scheme: SamplingScheme::InSample,
                ncs_scheme: SamplingScheme::InSample,
                stages: &stages,
                historical_library: None,
                external_inflow_library: None,
                external_load_library: None,
                external_ncs_library: None,
                recent_accum_seed: &[],
                recent_weight_seed: 0.0,
            },
            &comm,
            || Ok(MockSolver::with_fixed(100.0)),
        )
        .unwrap();

        let events: Vec<TrainingEvent> = rx.try_iter().collect();

        // Collect all PolicyTemplateBakeComplete events.
        let bake_events: Vec<&TrainingEvent> = events
            .iter()
            .filter(|e| matches!(e, TrainingEvent::PolicyTemplateBakeComplete { .. }))
            .collect();

        // Exactly one per iteration (2 iterations).
        assert_eq!(
            bake_events.len(),
            2,
            "expected exactly 2 PolicyTemplateBakeComplete events, got {}",
            bake_events.len()
        );

        // Each event must report stages_processed == n_stages.
        for event in &bake_events {
            let TrainingEvent::PolicyTemplateBakeComplete {
                stages_processed, ..
            } = event
            else {
                panic!("wrong variant")
            };
            assert_eq!(
                *stages_processed, n_stages as u32,
                "stages_processed must equal num_stages"
            );
        }

        // On iteration 2, the backward pass from iteration 1 will have added
        // cuts, so total_rows_baked must be > 0.
        let second_bake = bake_events[1];
        let TrainingEvent::PolicyTemplateBakeComplete {
            total_rows_baked, ..
        } = second_bake
        else {
            panic!("wrong variant")
        };
        assert!(
            *total_rows_baked > 0,
            "iteration 2 bake must have baked at least one cut row (backward pass \
             generated cuts on iteration 1)"
        );
    }

    /// `TrainingResult::new` assigns every field correctly.
    ///
    /// Calls the canonical constructor with 11 explicit, distinct values and
    /// asserts each field on the returned struct. The test fails at compile time
    /// if any field is renamed, reordered, or removed without a corresponding
    /// update to the constructor signature.
    #[test]
    fn ac_training_result_new_assigns_all_fields() {
        use crate::workspace::CapturedBasis;

        let basis_cache = vec![Some(CapturedBasis {
            basis: Basis {
                col_status: vec![1_i32],
                row_status: vec![2_i32],
            },
            base_row_count: 3,
            cut_row_slots: vec![4_u32],
            state_at_capture: vec![5.0_f64],
        })];
        // SolverStatsEntry is a 7-tuple:
        // (iteration: u64, phase: &'static str, stage: i32, opening: i32,
        //  rank: i32, worker_id: i32, delta: SolverStatsDelta)
        let solver_stats_log = vec![(
            7_u64,
            "forward",
            -1_i32,
            -1_i32,
            0_i32,
            -1_i32,
            SolverStatsDelta::default(),
        )];

        let result = super::TrainingResult::new(
            1.5_f64,                       // final_lb
            2.5_f64,                       // final_ub
            0.25_f64,                      // final_ub_std
            0.1_f64,                       // final_gap
            42_u64,                        // iterations
            "iteration_limit".to_string(), // reason
            9_999_u64,                     // total_time_ms
            basis_cache,
            solver_stats_log,
            None, // visited_archive
            None, // baked_templates
        );

        assert_eq!(result.final_lb, 1.5_f64, "final_lb");
        assert_eq!(result.final_ub, 2.5_f64, "final_ub");
        assert_eq!(result.final_ub_std, 0.25_f64, "final_ub_std");
        assert_eq!(result.final_gap, 0.1_f64, "final_gap");
        assert_eq!(result.iterations, 42_u64, "iterations");
        assert_eq!(result.reason, "iteration_limit", "reason");
        assert_eq!(result.total_time_ms, 9_999_u64, "total_time_ms");
        assert_eq!(result.basis_cache.len(), 1, "basis_cache length");
        let captured = result.basis_cache[0].as_ref().expect("basis_cache[0]");
        assert_eq!(captured.base_row_count, 3, "basis_cache[0].base_row_count");
        assert_eq!(result.solver_stats_log.len(), 1, "solver_stats_log length");
        assert_eq!(
            result.solver_stats_log[0].0, 7_u64,
            "solver_stats_log[0].0 (iteration)"
        );
        assert!(result.visited_archive.is_none(), "visited_archive");
        assert!(result.baked_templates.is_none(), "baked_templates");
    }
}