cobre-sddp 0.8.2

use std::collections::BTreeMap;

use chrono::NaiveDate;
use cobre_comm::{CommData, Communicator, ReduceOp};
use cobre_core::entities::hydro::{Hydro, HydroGenerationModel, HydroPenalties};
use cobre_core::scenario::{
    CorrelationEntity, CorrelationGroup, CorrelationModel, CorrelationProfile, InflowModel,
    LoadModel, SamplingScheme,
};
use cobre_core::temporal::{
    Block, BlockMode, NoiseMethod, ScenarioSourceConfig, Stage, StageRiskConfig, StageStateConfig,
};
use cobre_core::{Bus, DeficitSegment, EntityId, SystemBuilder};
use cobre_solver::{
    Basis, LpSolution, ProfiledSolver, RowBatch, SolverError, SolverInterface, SolverStatistics,
    StageTemplate,
};
use cobre_stochastic::StochasticContext;
use cobre_stochastic::context::{ClassSchemes, OpeningTreeInputs, build_stochastic_context};

use cobre_comm::LocalBackend;

use super::{
    ForwardPassBatch, ForwardResult, SyncResult, build_delta_cut_row_batch_into, run_forward_pass,
    sync_forward,
};
use crate::cut::row::build_cut_row_batch_into;
use crate::solve::partition;
use crate::{
    StoppingMode, StoppingRule, StoppingRuleSet, TrainingConfig,
    config::{CutManagementConfig, EventConfig, LoopConfig},
    context::{StageContext, TrainingContext},
    cut::FutureCostFunction,
    horizon_mode::HorizonMode,
    indexer::StageIndexer,
    inflow_method::InflowNonNegativityMethod,
    risk_measure::RiskMeasure,
    trajectory::TrajectoryRecord,
    workspace::{BackwardAccumulators, BasisStore, SolverWorkspace},
};

// ── Mock solver ──────────────────────────────────────────────────────────

/// Mock solver that returns a configurable fixed `LpSolution` on every `solve()`.
///
/// Optionally returns `SolverError::Infeasible` at a specific
/// `(scenario, stage)` pair (counted across calls in the scenario-outer,
/// stage-inner traversal order). `infeasible_at` counts global solve
/// calls starting from 0.
///
/// `warm_start_calls` is incremented each time `solve(Some(&basis))`
/// is called, enabling warm-start invocation tests.
struct MockSolver {
    solution: LpSolution,
    /// If `Some(n)`, the n-th solve call (0-indexed, counting both cold-start
    /// and warm-start calls) returns infeasible.
    infeasible_at: Option<usize>,
    call_count: usize,
    /// Number of times `solve(Some(&basis))` has been called.
    warm_start_calls: usize,
    /// Internal buffers that `SolutionView` borrows from.
    buf_primal: Vec<f64>,
    buf_dual: Vec<f64>,
    buf_reduced_costs: Vec<f64>,
}

impl MockSolver {
    /// Create a solver that always returns `solution`.
    fn always_ok(solution: LpSolution) -> Self {
        let buf_primal = solution.primal.clone();
        let buf_dual = solution.dual.clone();
        let buf_reduced_costs = solution.reduced_costs.clone();
        Self {
            solution,
            infeasible_at: None,
            call_count: 0,
            warm_start_calls: 0,
            buf_primal,
            buf_dual,
            buf_reduced_costs,
        }
    }

    /// Create a solver that returns infeasible on the `n`-th solve call.
    fn infeasible_on(solution: LpSolution, n: usize) -> Self {
        let buf_primal = solution.primal.clone();
        let buf_dual = solution.dual.clone();
        let buf_reduced_costs = solution.reduced_costs.clone();
        Self {
            solution,
            infeasible_at: Some(n),
            call_count: 0,
            warm_start_calls: 0,
            buf_primal,
            buf_dual,
            buf_reduced_costs,
        }
    }

    /// Shared solve logic used by both cold-start and warm-start paths.
    ///
    /// Increments `call_count` and returns `Infeasible` when `call_count`
    /// matches `infeasible_at`, otherwise returns the stored solution.
    fn do_solve(&mut self) -> Result<cobre_solver::SolutionView<'_>, SolverError> {
        let call = self.call_count;
        self.call_count += 1;
        if self.infeasible_at == Some(call) {
            return Err(SolverError::Infeasible);
        }
        // Fill internal buffers from the stored solution.
        self.buf_primal.clone_from(&self.solution.primal);
        self.buf_dual.clone_from(&self.solution.dual);
        self.buf_reduced_costs
            .clone_from(&self.solution.reduced_costs);
        Ok(cobre_solver::SolutionView {
            objective: self.solution.objective,
            primal: &self.buf_primal,
            dual: &self.buf_dual,
            reduced_costs: &self.buf_reduced_costs,
            iterations: self.solution.iterations,
            solve_time_seconds: self.solution.solve_time_seconds,
        })
    }
}

impl SolverInterface for MockSolver {
    type Profile = cobre_solver::ActiveProfile;

    fn apply_profile(&mut self, _profile: &cobre_solver::ActiveProfile) {}

    fn solver_name_version(&self) -> String {
        "MockSolver 0.0.0".to_string()
    }
    fn load_model(&mut self, _template: &StageTemplate) {}

    fn add_rows(&mut self, _cuts: &RowBatch) {}

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

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

    fn solve(
        &mut self,
        basis: Option<&Basis>,
    ) -> Result<cobre_solver::SolutionView<'_>, SolverError> {
        if basis.is_some() {
            self.warm_start_calls += 1;
        }
        self.do_solve()
    }

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

    fn statistics(&self) -> SolverStatistics {
        SolverStatistics {
            solve_count: self.call_count as u64,
            ..SolverStatistics::default()
        }
    }

    fn statistics_into(&self, out: &mut SolverStatistics) {
        out.copy_from(&self.statistics());
    }

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

// ── Helpers ──────────────────────────────────────────────────────────────

/// Minimal valid stage template for N=1 hydro, L=0 PAR order.
///
/// Column layout: [storage (0), `storage_in` (1), theta (2)]
/// Row layout: [`storage_fixing` (0)]
fn minimal_template_1_0() -> StageTemplate {
    // N=1, L=0:
    //   storage      = 0..1
    //   storage_in   = 1..2
    //   theta        = 2
    //   n_state      = 1
    //   n_transfer   = 0
    //   n_dual_relevant = 1
    //
    // Column layout for N=1, L=0:
    //   col 0: storage_out, col 1: z_inflow, col 2: storage_in, col 3: theta
    //
    // LP: min theta  s.t. storage_in = ? (patched)  x >= 0
    //
    // CSC matrix has 1 non-zero: storage_in coefficient in storage_fixing row.
    // Simplified to a structurally valid but otherwise no-op LP for testing.
    StageTemplate {
        num_cols: 4,
        num_rows: 1,
        num_nz: 1,
        col_starts: vec![0_i32, 0, 0, 1, 1], // col 2 (storage_in) has NZ at row 0
        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], // minimise theta (at col 3)
        row_lower: vec![0.0],
        row_upper: vec![0.0],
        n_state: 1,
        n_transfer: 0,
        n_dual_relevant: 1,
        n_hydro: 1,
        max_par_order: 0,
        col_scale: Vec::new(),
        row_scale: Vec::new(),
    }
}

/// Build a fixed `LpSolution` with `num_cols` columns.
///
/// `objective` is passed directly. `primal[theta_col]` is set to
/// `theta_val`; all other primal entries are zero.
fn fixed_solution(num_cols: usize, objective: f64, theta_col: usize, theta_val: f64) -> LpSolution {
    let mut primal = vec![0.0_f64; num_cols];
    primal[theta_col] = theta_val;
    let num_rows = 1; // single structural row for minimal template
    LpSolution {
        objective,
        primal,
        dual: vec![0.0; num_rows],
        reduced_costs: vec![0.0; num_cols],
        iterations: 0,
        solve_time_seconds: 0.0,
    }
}

/// Allocate `n` empty `TrajectoryRecord`s.
fn empty_records(n: usize) -> Vec<TrajectoryRecord> {
    (0..n)
        .map(|_| TrajectoryRecord {
            primal: Vec::new(),
            dual: Vec::new(),
            stage_cost: 0.0,
            state: Vec::new(),
        })
        .collect()
}

/// Build a minimal `StochasticContext` for a single-hydro, 3-stage system.
///
/// Used by integration tests that call `run_forward_pass`. The `MockSolver`
/// ignores the noise values produced by `sample_forward`, so the exact
/// stochastic parameterisation does not affect correctness; it only needs
/// to be structurally valid for the sampling API.
#[allow(clippy::too_many_lines)]
fn make_stochastic_context_1_hydro_3_stages() -> StochasticContext {
    let bus = Bus {
        id: EntityId(0),
        name: "B0".to_string(),
        deficit_segments: vec![DeficitSegment {
            depth_mw: None,
            cost_per_mwh: 1000.0,
        }],
        excess_cost: 0.0,
    };
    let hydro = Hydro {
        id: EntityId(1),
        name: "H1".to_string(),
        bus_id: EntityId(0),
        downstream_id: None,
        entry_stage_id: None,
        exit_stage_id: None,
        min_storage_hm3: 0.0,
        max_storage_hm3: 100.0,
        min_outflow_m3s: 0.0,
        max_outflow_m3s: None,
        generation_model: HydroGenerationModel::ConstantProductivity,
        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, id: i32| Stage {
        index: idx,
        id,
        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: 3,
            noise_method: NoiseMethod::Saa,
        },
    };
    let stages = vec![make_stage(0, 0), make_stage(1, 1), make_stage(2, 2)];
    let inflow = |stage_id: i32| InflowModel {
        hydro_id: EntityId(1),
        stage_id,
        mean_m3s: 100.0,
        std_m3s: 30.0,
        ar_coefficients: vec![],
        residual_std_ratio: 1.0,
        annual: None,
    };
    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(vec![inflow(0), inflow(1), inflow(2)])
        .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()
}

// ── Unit tests: ForwardResult ────────────────────────────────────────────

#[test]
fn forward_result_field_access() {
    let r = ForwardResult {
        scenario_costs: vec![60.0, 70.0, 80.0, 90.0],
        elapsed_ms: 123,
        lp_solves: 0,
        setup_time_ms: 0,
        load_imbalance_ms: 0,
        scheduling_overhead_ms: 0,
        stage_stats: Vec::new(),
    };
    assert_eq!(r.scenario_costs.len(), 4);
    assert_eq!(r.scenario_costs[0], 60.0);
    assert_eq!(r.elapsed_ms, 123);
}

#[test]
fn forward_result_clone_and_debug() {
    let r = ForwardResult {
        scenario_costs: vec![1.0, 2.0],
        elapsed_ms: 5,
        lp_solves: 0,
        setup_time_ms: 0,
        load_imbalance_ms: 0,
        scheduling_overhead_ms: 0,
        stage_stats: Vec::new(),
    };
    let c = r.clone();
    assert_eq!(c.scenario_costs.len(), r.scenario_costs.len());
    assert_eq!(c.scenario_costs[0].to_bits(), r.scenario_costs[0].to_bits());
    let s = format!("{r:?}");
    assert!(s.contains("ForwardResult"));
}

// ── Unit tests: forward overhead decomposition ───────────────────────────

/// Verify the three-component decomposition: 4 workers with solve times
/// 500/600/550/580 ms and setup times 50/60/45/55 ms yields setup=210,
/// imbalance=50, scheduling varies by `parallel_wall_ms`.
#[test]
fn forward_overhead_decomposition_four_workers() {
    use cobre_solver::SolverStatistics;

    use crate::solver_stats::SolverStatsDelta;

    fn make_stats(
        solve_s: f64,
        load_model_s: f64,
        set_bounds_s: f64,
        basis_set_s: f64,
    ) -> SolverStatistics {
        SolverStatistics {
            total_solve_time_seconds: solve_s,
            total_load_model_time_seconds: load_model_s,
            total_set_bounds_time_seconds: set_bounds_s,
            total_basis_set_time_seconds: basis_set_s,
            ..SolverStatistics::default()
        }
    }

    // Solve times (s): 0.5, 0.6, 0.55, 0.58; setup times (s): 0.05, 0.06, 0.045, 0.055
    let befores = [
        make_stats(0.0, 0.0, 0.0, 0.0),
        make_stats(0.0, 0.0, 0.0, 0.0),
        make_stats(0.0, 0.0, 0.0, 0.0),
        make_stats(0.0, 0.0, 0.0, 0.0),
    ];
    let afters = [
        make_stats(0.500, 0.050, 0.0, 0.0),
        make_stats(0.600, 0.060, 0.0, 0.0),
        make_stats(0.550, 0.045, 0.0, 0.0),
        make_stats(0.580, 0.055, 0.0, 0.0),
    ];

    let deltas: Vec<SolverStatsDelta> = befores
        .iter()
        .zip(&afters)
        .map(|(b, a)| SolverStatsDelta::from_snapshots(b, a))
        .collect();

    // setup_time_ms: sum of all workers' setup phases.
    let setup_ms: f64 = deltas
        .iter()
        .map(|d| d.load_model_time_ms + d.set_bounds_time_ms + d.basis_set_time_ms)
        .sum();

    // Per-worker totals: solve + setup.
    let worker_totals: Vec<f64> = deltas
        .iter()
        .map(|d| {
            d.solve_time_ms + d.load_model_time_ms + d.set_bounds_time_ms + d.basis_set_time_ms
        })
        .collect();

    let n_workers_f = 4.0_f64;
    let max_ms = worker_totals.iter().copied().fold(0.0_f64, f64::max);
    let avg_ms = worker_totals.iter().sum::<f64>() / n_workers_f;
    let imbalance_ms = (max_ms - avg_ms).max(0.0);
    let parallel_wall_ms = 700_u64; // 40 ms above the slowest worker
    #[allow(clippy::cast_precision_loss)]
    let scheduling_ms = (parallel_wall_ms as f64 - max_ms).max(0.0);

    // setup_time_ms = 50 + 60 + 45 + 55 = 210
    assert!(
        (setup_ms - 210.0).abs() < 0.001,
        "setup_time_ms should be 210, got {setup_ms}"
    );
    // max_worker total = 660 (worker 1: 600+60)
    assert!(
        (max_ms - 660.0).abs() < 0.001,
        "max_worker_ms should be 660, got {max_ms}"
    );
    // avg = (550+660+595+635)/4 = 610
    assert!(
        (avg_ms - 610.0).abs() < 0.001,
        "avg_worker_ms should be 610, got {avg_ms}"
    );
    // imbalance = 660 - 610 = 50
    assert!(
        (imbalance_ms - 50.0).abs() < 0.001,
        "load_imbalance_ms should be 50, got {imbalance_ms}"
    );
    // scheduling = 700 - 660 = 40
    assert!(
        (scheduling_ms - 40.0).abs() < 0.001,
        "scheduling_overhead_ms should be 40, got {scheduling_ms}"
    );
}

/// Edge case: single worker — load imbalance must be exactly zero.
#[test]
fn forward_overhead_decomposition_single_worker_zero_imbalance() {
    use cobre_solver::SolverStatistics;

    use crate::solver_stats::SolverStatsDelta;

    let before = SolverStatistics::default();
    let after = SolverStatistics {
        total_solve_time_seconds: 1.0,
        total_load_model_time_seconds: 0.1,
        ..SolverStatistics::default()
    };

    let deltas = [SolverStatsDelta::from_snapshots(&before, &after)];
    let worker_totals: Vec<f64> = deltas
        .iter()
        .map(|d| {
            d.solve_time_ms + d.load_model_time_ms + d.set_bounds_time_ms + d.basis_set_time_ms
        })
        .collect();

    let n_workers_f = 1.0_f64;
    let max_ms = worker_totals.iter().copied().fold(0.0_f64, f64::max);
    let avg_ms = worker_totals.iter().sum::<f64>() / n_workers_f;
    let imbalance_ms = (max_ms - avg_ms).max(0.0);

    assert_eq!(
        imbalance_ms, 0.0,
        "load_imbalance_ms must be 0.0 for a single worker"
    );
}

/// Edge case: scheduling overhead is clamped to zero when parallel wall
/// time is less than the max worker total (clock-skew scenario).
#[test]
fn forward_overhead_scheduling_clamped_to_zero_on_clock_skew() {
    use cobre_solver::SolverStatistics;

    use crate::solver_stats::SolverStatsDelta;

    let before = SolverStatistics::default();
    let after = SolverStatistics {
        total_solve_time_seconds: 1.0, // 1000 ms
        ..SolverStatistics::default()
    };

    let deltas = [SolverStatsDelta::from_snapshots(&before, &after)];
    let max_ms = deltas
        .iter()
        .map(|d| d.solve_time_ms)
        .fold(0.0_f64, f64::max);

    // Wall time (800 ms) < max worker total (1000 ms) — clock skew.
    let parallel_wall_ms = 800_u64;
    #[allow(clippy::cast_precision_loss)]
    let scheduling_ms = (parallel_wall_ms as f64 - max_ms).max(0.0);

    assert_eq!(
        scheduling_ms, 0.0,
        "scheduling_overhead_ms must clamp to 0.0 on clock skew"
    );
}

/// Build a single-workspace from a solver sized for the given `indexer`.
fn single_workspace(solver: MockSolver, indexer: &StageIndexer) -> SolverWorkspace<MockSolver> {
    SolverWorkspace {
        rank: 0,
        worker_id: 0,
        solver: ProfiledSolver::new(solver),
        patch_buf: crate::lp_builder::PatchBuffer::new(
            indexer.hydro_count,
            indexer.max_par_order,
            0,
            0,
            0,
            0,
        ),
        current_state: Vec::with_capacity(indexer.n_state),
        scratch: crate::workspace::ScratchBuffers {
            noise_buf: Vec::with_capacity(indexer.hydro_count),
            inflow_m3s_buf: Vec::with_capacity(indexer.hydro_count),
            lag_matrix_buf: Vec::with_capacity(indexer.max_par_order * indexer.hydro_count),
            par_inflow_buf: Vec::with_capacity(indexer.hydro_count),
            eta_floor_buf: Vec::with_capacity(indexer.hydro_count),
            zero_targets_buf: vec![0.0_f64; indexer.hydro_count],
            ncs_col_upper_buf: Vec::new(),
            ncs_col_lower_buf: Vec::new(),
            ncs_col_indices_buf: Vec::new(),
            load_rhs_buf: Vec::new(),
            row_lower_buf: Vec::new(),
            z_inflow_rhs_buf: Vec::new(),
            effective_eta_buf: Vec::new(),
            unscaled_primal: Vec::new(),
            unscaled_dual: Vec::new(),
            lag_accumulator: vec![],
            lag_weight_accum: 0.0,
            downstream_accumulator: Vec::new(),
            downstream_weight_accum: 0.0,
            downstream_completed_lags: Vec::new(),
            downstream_n_completed: 0,
            recon_slot_lookup: Vec::new(),
            trajectory_costs_buf: Vec::new(),
            raw_noise_buf: Vec::new(),
            perm_scratch: Vec::new(),
            anticipated_state_buf: Vec::new(),
        },
        scratch_basis: Basis::new(0, 0),
        backward_accum: BackwardAccumulators::default(),
        worker_timing_buf: cobre_core::WorkerPhaseTimings::default(),
    }
}

/// Build 3 minimal [`Stage`] values matching `make_stochastic_context_1_hydro_3_stages`.
///
/// Provides the `stages` slice required by [`TrainingContext`] so that
/// [`cobre_stochastic::build_forward_sampler`] can read per-stage noise methods.
fn make_stages_3() -> Vec<Stage> {
    let make_stage = |idx: usize, id: i32| Stage {
        index: idx,
        id,
        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: 3,
            noise_method: NoiseMethod::Saa,
        },
    };
    vec![make_stage(0, 0), make_stage(1, 1), make_stage(2, 2)]
}

// ── Acceptance criteria integration tests ───────────────────────────────

/// AC: 2 scenarios, 3 stages, fixed `LpSolution(objective=100, theta=30)`.
/// Expected: `scenario_count=2`, all 6 records with `stage_cost=70_000`.
#[test]
#[allow(clippy::too_many_lines)]
fn ac_two_scenarios_three_stages_fixed_solution() {
    // StageIndexer: N=1, L=0 → n_state=1, theta=3, num_cols=4
    let indexer = {
        let mut ix = StageIndexer::new(1, 0);
        ix.finalize_for_test();
        ix
    };
    let solution = fixed_solution(4, 100.0, indexer.theta, 30.0);
    let solver = MockSolver::always_ok(solution);
    let fcf = FutureCostFunction::new(3, indexer.n_state, 2, 100, &[0; 3]);
    let config = TrainingConfig {
        loop_config: LoopConfig {
            forward_passes: 2,
            max_iterations: 100,
            start_iteration: 0,
            n_fwd_threads: 1,
            max_blocks: 1,
            stopping_rules: StoppingRuleSet {
                rules: vec![StoppingRule::IterationLimit { limit: 100 }],
                mode: StoppingMode::Any,
            },
        },
        cut_management: CutManagementConfig {
            cut_selection: None,
            budget: None,
            cut_activity_tolerance: 0.0,
            warm_start_cuts: 0,
            risk_measures: vec![RiskMeasure::Expectation],
        },
        events: EventConfig {
            event_sender: None,
            checkpoint_interval: None,
            shutdown_flag: None,
            export_states: false,
        },
    };

    let horizon = HorizonMode::Finite { num_stages: 3 };

    let templates = vec![
        minimal_template_1_0(),
        minimal_template_1_0(),
        minimal_template_1_0(),
    ];
    let base_rows = vec![2usize, 2, 2];
    let initial_state = vec![0.0_f64; indexer.n_state];
    let mut records = empty_records(2 * 3);
    let stochastic = make_stochastic_context_1_hydro_3_stages();
    let stages = make_stages_3();
    let mut ws = single_workspace(solver, &indexer);
    let mut basis_store =
        BasisStore::new(config.loop_config.forward_passes as usize, templates.len());

    let ctx = StageContext {
        templates: &templates,
        base_rows: &base_rows,
        noise_scale: &[],
        n_hydros: 0,
        n_load_buses: 0,
        load_balance_row_starts: &[],
        load_bus_indices: &[],
        block_counts_per_stage: &[1usize, 1, 1],
        ncs_max_gen: &[],
        ncs_allow_curtailment: &[],
        discount_factors: &[],
        cumulative_discount_factors: &[],
        stage_lag_transitions: &[],
        noise_group_ids: &[],
        downstream_par_order: 0,
    };
    let result = run_forward_pass(
        std::slice::from_mut(&mut ws),
        &mut basis_store,
        &ctx,
        &templates,
        &fcf,
        &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,
            dcs: None,
            noise_key_diag: None,
        },
        &ForwardPassBatch {
            local_forward_passes: config.loop_config.forward_passes as usize,
            total_forward_passes: config.loop_config.forward_passes as usize,
            iteration: 0,
            fwd_offset: 0,
            event_sender: None,
        },
        &mut records,
    )
    .unwrap();

    // AC: scenario_costs has exactly 2 entries (one per forward pass).
    assert_eq!(result.scenario_costs.len(), 2);
    // AC: all 6 records have stage_cost = (100 - 30) * COST_SCALE_FACTOR = 70_000_000.
    for (i, record) in records.iter().enumerate() {
        assert_eq!(
            record.stage_cost, 70_000_000.0,
            "record[{i}].stage_cost should be 70_000_000.0 ((objective - theta) * COST_SCALE_FACTOR)"
        );
    }
    // AC: each scenario cost = 70_000_000 * 3 stages = 210_000_000.
    assert_eq!(result.scenario_costs[0], 210_000_000.0);
    assert_eq!(result.scenario_costs[1], 210_000_000.0);
}

/// AC: mock solver returns `Infeasible` at stage 1, scenario 0.
///
/// Call 0 = scenario 0 stage 0 (succeeds). Call 1 = scenario 0 stage 1
/// (infeasible). The function must return `SddpError::Infeasible { stage: 1,
/// scenario: 0 }`.
#[test]
#[allow(clippy::too_many_lines)]
fn ac_infeasible_at_stage_1_scenario_0_returns_infeasible_error() {
    let indexer = {
        let mut ix = StageIndexer::new(1, 0);
        ix.finalize_for_test();
        ix
    };
    let solution = fixed_solution(4, 100.0, indexer.theta, 30.0);
    // Stage-first loop: with 2 scenarios and 3 stages, the solve order is
    // (s0,t0), (s1,t0), (s0,t1), (s1,t1), ... — the 3rd call (index 2)
    // is stage 1 of scenario 0.
    let solver = MockSolver::infeasible_on(solution, 2);
    let fcf = FutureCostFunction::new(3, indexer.n_state, 2, 100, &[0; 3]);
    let config = TrainingConfig {
        loop_config: LoopConfig {
            forward_passes: 2,
            max_iterations: 100,
            start_iteration: 0,
            n_fwd_threads: 1,
            max_blocks: 1,
            stopping_rules: StoppingRuleSet {
                rules: vec![StoppingRule::IterationLimit { limit: 100 }],
                mode: StoppingMode::Any,
            },
        },
        cut_management: CutManagementConfig {
            cut_selection: None,
            budget: None,
            cut_activity_tolerance: 0.0,
            warm_start_cuts: 0,
            risk_measures: vec![RiskMeasure::Expectation],
        },
        events: EventConfig {
            event_sender: None,
            checkpoint_interval: None,
            shutdown_flag: None,
            export_states: false,
        },
    };

    let horizon = HorizonMode::Finite { num_stages: 3 };

    let templates = vec![
        minimal_template_1_0(),
        minimal_template_1_0(),
        minimal_template_1_0(),
    ];
    let base_rows = vec![2usize, 2, 2];
    let initial_state = vec![0.0_f64; indexer.n_state];
    let mut records = empty_records(2 * 3);
    let stochastic = make_stochastic_context_1_hydro_3_stages();
    let stages = make_stages_3();
    let mut ws = single_workspace(solver, &indexer);
    let mut basis_store =
        BasisStore::new(config.loop_config.forward_passes as usize, templates.len());

    let ctx = StageContext {
        templates: &templates,
        base_rows: &base_rows,
        noise_scale: &[],
        n_hydros: 0,
        n_load_buses: 0,
        load_balance_row_starts: &[],
        load_bus_indices: &[],
        block_counts_per_stage: &[1usize, 1, 1],
        ncs_max_gen: &[],
        ncs_allow_curtailment: &[],
        discount_factors: &[],
        cumulative_discount_factors: &[],
        stage_lag_transitions: &[],
        noise_group_ids: &[],
        downstream_par_order: 0,
    };
    let result = run_forward_pass(
        std::slice::from_mut(&mut ws),
        &mut basis_store,
        &ctx,
        &templates,
        &fcf,
        &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,
            dcs: None,
            noise_key_diag: None,
        },
        &ForwardPassBatch {
            local_forward_passes: config.loop_config.forward_passes as usize,
            total_forward_passes: config.loop_config.forward_passes as usize,
            iteration: 0,
            fwd_offset: 0,
            event_sender: None,
        },
        &mut records,
    );

    // AC: must return SddpError::Infeasible with stage=1, scenario=0.
    match result {
        Err(crate::SddpError::Infeasible {
            stage, scenario, ..
        }) => {
            assert_eq!(stage, 1, "expected stage=1");
            assert_eq!(scenario, 0, "expected scenario=0");
        }
        other => panic!("expected Infeasible, got {other:?}"),
    }
}

/// AC: with `forward_passes=3`, rank=1, size=2, `global_scenario` for m=0 is 3.
#[test]
fn ac_global_scenario_index_rank1_scenario0() {
    // global_scenario = rank * forward_passes + m = 1 * 3 + 0 = 3
    let rank = 1usize;
    let forward_passes = 3usize;
    let m = 0usize;
    let global_scenario = rank * forward_passes + m;
    assert_eq!(global_scenario, 3);
}

/// Behavioral: `cost_sum` and `cost_sum_sq` are correctly accumulated.
///
/// With 2 scenarios and `stage_cost=70_000` at every `(scenario, stage)`:
/// - `total_cost` per scenario = `70_000` \* 3 = `210_000`
/// - `cost_sum` = `210_000` + `210_000` = `420_000`
/// - `cost_sum_sq` = `210_000`^2 + `210_000`^2 = `88_200_000_000`
#[test]
#[allow(clippy::too_many_lines)]
fn cost_statistics_accumulated_correctly() {
    let indexer = {
        let mut ix = StageIndexer::new(1, 0);
        ix.finalize_for_test();
        ix
    };
    let solution = fixed_solution(4, 100.0, indexer.theta, 30.0);
    let solver = MockSolver::always_ok(solution);
    let fcf = FutureCostFunction::new(3, indexer.n_state, 2, 100, &[0; 3]);
    let config = TrainingConfig {
        loop_config: LoopConfig {
            forward_passes: 2,
            max_iterations: 100,
            start_iteration: 0,
            n_fwd_threads: 1,
            max_blocks: 1,
            stopping_rules: StoppingRuleSet {
                rules: vec![StoppingRule::IterationLimit { limit: 100 }],
                mode: StoppingMode::Any,
            },
        },
        cut_management: CutManagementConfig {
            cut_selection: None,
            budget: None,
            cut_activity_tolerance: 0.0,
            warm_start_cuts: 0,
            risk_measures: vec![RiskMeasure::Expectation],
        },
        events: EventConfig {
            event_sender: None,
            checkpoint_interval: None,
            shutdown_flag: None,
            export_states: false,
        },
    };

    let horizon = HorizonMode::Finite { num_stages: 3 };

    let templates = vec![
        minimal_template_1_0(),
        minimal_template_1_0(),
        minimal_template_1_0(),
    ];
    let base_rows = vec![2usize, 2, 2];
    let initial_state = vec![0.0_f64; indexer.n_state];
    let mut records = empty_records(2 * 3);
    let stochastic = make_stochastic_context_1_hydro_3_stages();
    let stages = make_stages_3();
    let mut ws = single_workspace(solver, &indexer);
    let mut basis_store =
        BasisStore::new(config.loop_config.forward_passes as usize, templates.len());

    let ctx = StageContext {
        templates: &templates,
        base_rows: &base_rows,
        noise_scale: &[],
        n_hydros: 0,
        n_load_buses: 0,
        load_balance_row_starts: &[],
        load_bus_indices: &[],
        block_counts_per_stage: &[1usize, 1, 1],
        ncs_max_gen: &[],
        ncs_allow_curtailment: &[],
        discount_factors: &[],
        cumulative_discount_factors: &[],
        stage_lag_transitions: &[],
        noise_group_ids: &[],
        downstream_par_order: 0,
    };
    let result = run_forward_pass(
        std::slice::from_mut(&mut ws),
        &mut basis_store,
        &ctx,
        &templates,
        &fcf,
        &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,
            dcs: None,
            noise_key_diag: None,
        },
        &ForwardPassBatch {
            local_forward_passes: config.loop_config.forward_passes as usize,
            total_forward_passes: config.loop_config.forward_passes as usize,
            iteration: 0,
            fwd_offset: 0,
            event_sender: None,
        },
        &mut records,
    )
    .unwrap();

    // stage_cost per solve = (100 - 30) * COST_SCALE_FACTOR = 70_000_000
    // total_cost per scenario = 70_000_000 * 3 stages = 210_000_000
    assert_eq!(result.scenario_costs.len(), 2);
    assert_eq!(result.scenario_costs[0], 210_000_000.0);
    assert_eq!(result.scenario_costs[1], 210_000_000.0);
    // Derived statistics: sum = 420_000_000, sum_sq = 210_000_000^2 * 2.
    let cost_sum: f64 = result.scenario_costs.iter().sum();
    let cost_sum_sq: f64 = result.scenario_costs.iter().map(|c| c * c).sum();
    assert_eq!(cost_sum, 420_000_000.0);
    assert_eq!(cost_sum_sq, 210_000_000.0_f64.powi(2) * 2.0);
}

// ── Unit tests: SyncResult ───────────────────────────────────────────────

#[test]
fn sync_result_field_access() {
    let r = SyncResult {
        global_ub_mean: 75.0,
        global_ub_std: 12.909,
        ci_95_half_width: 12.651,
        sync_time_ms: 7,
    };
    assert_eq!(r.global_ub_mean, 75.0);
    assert_eq!(r.global_ub_std, 12.909);
    assert_eq!(r.ci_95_half_width, 12.651);
    assert_eq!(r.sync_time_ms, 7);
}

#[test]
fn sync_result_clone_and_debug() {
    let r = SyncResult {
        global_ub_mean: 2.0,
        global_ub_std: 3.0,
        ci_95_half_width: 4.0,
        sync_time_ms: 5,
    };
    let c = r.clone();
    assert_eq!(c.global_ub_mean, r.global_ub_mean);
    assert_eq!(c.global_ub_std, r.global_ub_std);
    let s = format!("{r:?}");
    assert!(s.contains("SyncResult"));
}

// ── Unit tests: UB statistics computation ───────────────────────────────

/// AC: 4 scenarios with costs [60, 70, 80, 90].
///
/// `cost_sum` = 300, `cost_sum_sq` = 60²+70²+80²+90² = 23000, count = 4.
/// mean = 75.0
/// variance = (23000 - 4 * 75^2) / 3 = (23000 - 22500) / 3 = 500/3
/// std = sqrt(500/3) ≈ 12.910
/// `ci_95` = 1.96 * std / sqrt(4)
#[test]
fn ub_statistics_four_scenarios_correct_mean_and_std() {
    let local = ForwardResult {
        scenario_costs: vec![60.0, 70.0, 80.0, 90.0],
        elapsed_ms: 0,
        lp_solves: 0,
        setup_time_ms: 0,
        load_imbalance_ms: 0,
        scheduling_overhead_ms: 0,
        stage_stats: Vec::new(),
    };
    let comm = LocalBackend;
    let result = sync_forward(&local, &comm, 4).unwrap();

    assert_eq!(result.global_ub_mean, 75.0, "mean must be 300/4 = 75");

    // variance = 500/3, std = sqrt(500/3) ≈ 12.910
    let expected_std = (500.0_f64 / 3.0).sqrt();
    let tolerance = 1e-9;
    assert!(
        (result.global_ub_std - expected_std).abs() < tolerance,
        "std deviation {got} should be ≈ {expected_std}",
        got = result.global_ub_std,
    );

    // ci_95 = 1.96 * std / sqrt(4) = 1.96 * std / 2
    let expected_ci = 1.96_f64 * expected_std / 4.0_f64.sqrt();
    assert!(
        (result.ci_95_half_width - expected_ci).abs() < tolerance,
        "ci_95 {got} should be ≈ {expected_ci}",
        got = result.ci_95_half_width,
    );
}

/// AC: 4 scenarios, costs [60,70,80,90].
///
/// Matches the exact acceptance criterion values: `global_ub_mean` = 75.0,
/// `global_ub_std` > 0.
///
/// The sequential summation of [60, 70, 80, 90] gives:
/// `cost_sum` = 300, `cost_sum_sq` = 23000, N = 4, mean = 75.
/// std = sqrt((23000 - 4*75^2) / 3) = sqrt(500/3) ≈ 12.910.
#[test]
fn acceptance_criterion_ub_mean() {
    let local = ForwardResult {
        scenario_costs: vec![60.0, 70.0, 80.0, 90.0],
        elapsed_ms: 0,
        lp_solves: 0,
        setup_time_ms: 0,
        load_imbalance_ms: 0,
        scheduling_overhead_ms: 0,
        stage_stats: Vec::new(),
    };
    let comm = LocalBackend;
    let result = sync_forward(&local, &comm, 4).unwrap();

    assert_eq!(result.global_ub_mean, 75.0);
    // std = sqrt((23000 - 4 * 75^2) / 3) = sqrt(500/3) ≈ 12.910
    let expected_std = (500.0_f64 / 3.0).sqrt();
    assert!(
        (result.global_ub_std - expected_std).abs() < 1e-9,
        "std deviation {got} should be ≈ {expected_std}",
        got = result.global_ub_std,
    );
}

/// Canonical summation: [1.0, 2.0, 3.0, 4.0] produces identical mean/std
/// regardless of how the vector is split across "ranks".
///
/// Verifies that the `allgatherv` + sequential summation approach produces
/// bit-identical statistics for the full vector `[1, 2, 3, 4]` whether it
/// is presented as a single rank with 4 scenarios or simulated as two
/// ranks with 2 scenarios each.
#[test]
fn canonical_summation_identical_regardless_of_partition() {
    // Single-rank: all 4 costs provided together.
    let single_rank = ForwardResult {
        scenario_costs: vec![1.0, 2.0, 3.0, 4.0],
        elapsed_ms: 0,
        lp_solves: 0,
        setup_time_ms: 0,
        load_imbalance_ms: 0,
        scheduling_overhead_ms: 0,
        stage_stats: Vec::new(),
    };
    let comm = LocalBackend;
    let result_single = sync_forward(&single_rank, &comm, 4).unwrap();

    // Simulate two-rank scenario: rank 0 has [1.0, 2.0], rank 1 has [3.0, 4.0].
    // The mock communicator below pre-fills the recv buf with both ranks' data.
    // We test by constructing the full global buffer manually and verifying
    // that sequential summation produces the same statistics.
    let global_costs = [1.0_f64, 2.0, 3.0, 4.0];
    let global_n = global_costs.len();
    #[allow(clippy::cast_precision_loss)]
    let global_n_f64 = global_n as f64;
    let cost_sum: f64 = global_costs.iter().sum();
    let cost_sum_sq: f64 = global_costs.iter().map(|c| c * c).sum();
    let mean = cost_sum / global_n_f64;
    let variance = (cost_sum_sq - global_n_f64 * mean * mean) / (global_n_f64 - 1.0);
    let expected_std = variance.max(0.0).sqrt();
    let expected_mean = mean;

    assert_eq!(
        result_single.global_ub_mean.to_bits(),
        expected_mean.to_bits(),
        "mean must be bit-identical to sequential summation of [1,2,3,4]"
    );
    assert_eq!(
        result_single.global_ub_std.to_bits(),
        expected_std.to_bits(),
        "std must be bit-identical to sequential summation of [1,2,3,4]"
    );
}

/// AC: Bessel correction edge case — single scenario produces zero variance.
#[test]
fn bessel_correction_single_scenario_zero_std_and_ci() {
    let local = ForwardResult {
        scenario_costs: vec![500.0],
        elapsed_ms: 0,
        lp_solves: 0,
        setup_time_ms: 0,
        load_imbalance_ms: 0,
        scheduling_overhead_ms: 0,
        stage_stats: Vec::new(),
    };
    let comm = LocalBackend;
    let result = sync_forward(&local, &comm, 1).unwrap();

    assert_eq!(
        result.global_ub_std, 0.0,
        "std must be 0.0 for a single scenario (N=1 Bessel correction)"
    );
    assert_eq!(
        result.ci_95_half_width, 0.0,
        "ci_95 must be 0.0 for a single scenario"
    );
}

/// Guard: negative variance from floating-point cancellation → std = 0.0, not NaN.
#[test]
fn negative_variance_guard_produces_zero_std_not_nan() {
    // Construct two identical large values so that the single-pass Bessel
    // formula (sum_sq - N*mean^2) / (N-1) can produce a tiny negative result
    // due to floating-point representation differences.
    //
    // With costs = [v, v]: sum = 2v, mean = v, sum_sq = 2v^2.
    // Variance = (2v^2 - 2v^2) / 1 = 0. In floating-point the exact
    // representation of sum_sq and N*mean^2 may differ by epsilon,
    // potentially yielding a slightly negative result.
    //
    // We synthesise this by using two scenarios whose costs are very close
    // to the representable large value but not exactly equal.
    let v = 1.0e15_f64;
    let local = ForwardResult {
        scenario_costs: vec![v, v],
        elapsed_ms: 0,
        lp_solves: 0,
        setup_time_ms: 0,
        load_imbalance_ms: 0,
        scheduling_overhead_ms: 0,
        stage_stats: Vec::new(),
    };
    let comm = LocalBackend;
    let result = sync_forward(&local, &comm, 2).unwrap();

    assert!(
        !result.global_ub_std.is_nan(),
        "std must not be NaN even when floating-point variance is slightly negative"
    );
    // Both costs are exactly equal so true variance = 0.
    // The max(0, variance).sqrt() guard must clamp any tiny negative value.
    assert_eq!(
        result.global_ub_std, 0.0,
        "std must be 0.0 when variance is zero (or clamps from tiny negative)"
    );
}

// ── Integration tests: sync_forward with LocalBackend ────────────────────

/// Integration: single-rank mode — global UB mean equals local mean.
#[test]
fn sync_forward_local_backend_global_equals_local() {
    // Two scenarios each costing 420.0 → mean = 420.0.
    let local = ForwardResult {
        scenario_costs: vec![420.0, 420.0],
        elapsed_ms: 5,
        lp_solves: 0,
        setup_time_ms: 0,
        load_imbalance_ms: 0,
        scheduling_overhead_ms: 0,
        stage_stats: Vec::new(),
    };
    let comm = LocalBackend;
    let result = sync_forward(&local, &comm, 2).unwrap();

    // In single-rank mode, allgatherv is an identity copy.
    assert_eq!(
        result.global_ub_mean, 420.0,
        "global_ub_mean must equal the arithmetic mean of the cost vector"
    );
}

/// Integration: `sync_time_ms` is a valid non-negative u64.
#[test]
fn sync_forward_sync_time_ms_is_valid_u64() {
    let local = ForwardResult {
        scenario_costs: vec![50.0, 50.0],
        elapsed_ms: 0,
        lp_solves: 0,
        setup_time_ms: 0,
        load_imbalance_ms: 0,
        scheduling_overhead_ms: 0,
        stage_stats: Vec::new(),
    };
    let comm = LocalBackend;
    let result = sync_forward(&local, &comm, 2).unwrap();
    // sync_time_ms is u64 — any value is a valid non-negative u64.
    // We just verify the field exists and doesn't overflow to something absurd.
    let _ = result.sync_time_ms;
}

/// Integration: `CommError` from a failing communicator is wrapped as `SddpError::Communication`.
#[test]
fn sync_forward_comm_error_wraps_as_sddp_communication() {
    use cobre_comm::CommError;

    /// Communicator that always returns `CommError::InvalidCommunicator`.
    struct FailingComm;

    impl Communicator for FailingComm {
        fn allgatherv<T: CommData>(
            &self,
            _send: &[T],
            _recv: &mut [T],
            _counts: &[usize],
            _displs: &[usize],
        ) -> Result<(), CommError> {
            Err(CommError::InvalidCommunicator)
        }

        fn allreduce<T: CommData>(
            &self,
            _send: &[T],
            _recv: &mut [T],
            _op: ReduceOp,
        ) -> Result<(), CommError> {
            Err(CommError::InvalidCommunicator)
        }

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

        fn barrier(&self) -> Result<(), CommError> {
            Err(CommError::InvalidCommunicator)
        }

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

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

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

    let local = ForwardResult {
        scenario_costs: vec![100.0],
        elapsed_ms: 0,
        lp_solves: 0,
        setup_time_ms: 0,
        load_imbalance_ms: 0,
        scheduling_overhead_ms: 0,
        stage_stats: Vec::new(),
    };
    let comm = FailingComm;
    let err = sync_forward(&local, &comm, 1).unwrap_err();

    assert!(
        matches!(err, crate::SddpError::Communication(_)),
        "CommError must be wrapped as SddpError::Communication, got: {err:?}"
    );
}

// ── Unit tests: warm-start basis caching ─────────────────────────────────

/// Helper: run one iteration of `run_forward_pass` with a single scenario
/// and a 3-stage horizon. The workspace is passed mutably; the basis store
/// is returned so callers can inspect per-scenario, per-stage cached bases.
fn run_one_iteration(
    ws: &mut SolverWorkspace<MockSolver>,
    basis_store: &mut BasisStore,
) -> Result<(), crate::SddpError> {
    let indexer = {
        let mut ix = StageIndexer::new(1, 0);
        ix.finalize_for_test();
        ix
    };
    let fcf = FutureCostFunction::new(3, indexer.n_state, 1, 100, &[0; 3]);
    let config = TrainingConfig {
        loop_config: LoopConfig {
            forward_passes: 1,
            max_iterations: 100,
            start_iteration: 0,
            n_fwd_threads: 1,
            max_blocks: 1,
            stopping_rules: StoppingRuleSet {
                rules: vec![StoppingRule::IterationLimit { limit: 100 }],
                mode: StoppingMode::Any,
            },
        },
        cut_management: CutManagementConfig {
            cut_selection: None,
            budget: None,
            cut_activity_tolerance: 0.0,
            warm_start_cuts: 0,
            risk_measures: vec![RiskMeasure::Expectation],
        },
        events: EventConfig {
            event_sender: None,
            checkpoint_interval: None,
            shutdown_flag: None,
            export_states: false,
        },
    };

    let horizon = HorizonMode::Finite { num_stages: 3 };

    let templates = vec![
        minimal_template_1_0(),
        minimal_template_1_0(),
        minimal_template_1_0(),
    ];
    let base_rows = vec![2usize, 2, 2];
    let initial_state = vec![0.0_f64; indexer.n_state];
    let mut records = empty_records(3);
    let stochastic = make_stochastic_context_1_hydro_3_stages();
    let stages = make_stages_3();

    let ctx = StageContext {
        templates: &templates,
        base_rows: &base_rows,
        noise_scale: &[],
        n_hydros: 0,
        n_load_buses: 0,
        load_balance_row_starts: &[],
        load_bus_indices: &[],
        block_counts_per_stage: &[1usize, 1, 1],
        ncs_max_gen: &[],
        ncs_allow_curtailment: &[],
        discount_factors: &[],
        cumulative_discount_factors: &[],
        stage_lag_transitions: &[],
        noise_group_ids: &[],
        downstream_par_order: 0,
    };
    run_forward_pass(
        std::slice::from_mut(ws),
        basis_store,
        &ctx,
        &templates,
        &fcf,
        &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,
            dcs: None,
            noise_key_diag: None,
        },
        &ForwardPassBatch {
            local_forward_passes: config.loop_config.forward_passes as usize,
            total_forward_passes: config.loop_config.forward_passes as usize,
            iteration: 0,
            fwd_offset: 0,
            event_sender: None,
        },
        &mut records,
    )
    .map(|_| ())
}

/// Warm-start invocation: the first iteration calls `solve(None)` (cold start);
/// the second iteration calls `solve(Some(&basis))` (warm start).
///
/// AC: `run_forward_pass` called twice, sharing the same `BasisStore`
/// (1 scenario × 3 stages). After iteration 1: `warm_start_calls` == 0
/// (3 cold-start solves). After iteration 2: `warm_start_calls` > 0 (all
/// 3 stages warm-start from the store populated in iteration 1).
#[test]
fn warm_start_first_iteration_cold_second_iteration_warm() {
    let indexer = {
        let mut ix = StageIndexer::new(1, 0);
        ix.finalize_for_test();
        ix
    };
    let solution = fixed_solution(4, 100.0, indexer.theta, 30.0);
    let solver = MockSolver::always_ok(solution);
    // Single workspace and a shared basis store (1 scenario × 3 stages).
    let mut ws = single_workspace(solver, &indexer);
    let mut basis_store = BasisStore::new(1, 3);

    // First iteration: no cached bases → all cold-start.
    run_one_iteration(&mut ws, &mut basis_store).unwrap();
    assert_eq!(
        ws.solver.inner().warm_start_calls,
        0,
        "first iteration must use cold-start for all stages (warm_start_calls == 0)"
    );

    // After first iteration, all 3 stages for scenario 0 have a cached basis.
    assert!(
        (0..3).all(|t| basis_store.get(0, t).is_some()),
        "basis_store must be fully populated for scenario 0 after the first iteration"
    );

    // Second iteration: cached bases present → all stages warm-start.
    run_one_iteration(&mut ws, &mut basis_store).unwrap();
    assert!(
        ws.solver.inner().warm_start_calls > 0,
        "second iteration must use warm-start for at least one stage \
         (warm_start_calls > 0, got {})",
        ws.solver.inner().warm_start_calls
    );
}

/// Basis invalidation on solver error: when a forward solve returns
/// `SolverError::Infeasible`, the `BasisStore` slot at `(scenario, stage)`
/// must be set to `None` before the error propagates.
///
/// AC: `MockSolver` returns `Infeasible` on call index 4 (second iteration,
/// stage 1 — calls 0-2 = first iteration stages 0,1,2; calls 3,4,5 = second
/// iteration stages 0,1,2). After the error:
/// - `basis_store.get(0, 1)` is `None` (invalidated at the failing stage).
/// - `basis_store.get(0, 0)` is `Some` (stage 0 succeeded in iteration 2).
#[test]
fn basis_invalidated_on_solver_error() {
    let indexer = {
        let mut ix = StageIndexer::new(1, 0);
        ix.finalize_for_test();
        ix
    };
    let solution = fixed_solution(4, 100.0, indexer.theta, 30.0);
    // Call 4 = second iteration, stage 1 (calls 0-2 = first iteration
    // stages 0,1,2; calls 3,4,5 = second iteration stages 0,1,2).
    let solver = MockSolver::infeasible_on(solution, 4);
    // Single workspace and a shared basis store (1 scenario × 3 stages).
    let mut ws = single_workspace(solver, &indexer);
    let mut basis_store = BasisStore::new(1, 3);

    // First iteration: all cold-start, all succeed, populate all 3 stages.
    run_one_iteration(&mut ws, &mut basis_store).unwrap();
    assert!(
        (0..3).all(|t| basis_store.get(0, t).is_some()),
        "basis_store must be fully populated for scenario 0 after iteration 1"
    );

    // Second iteration: stage 0 warm-starts (call 3 OK), stage 1 infeasible (call 4).
    let err = run_one_iteration(&mut ws, &mut basis_store).unwrap_err();
    assert!(
        matches!(err, crate::SddpError::Infeasible { stage: 1, .. }),
        "expected Infeasible at stage 1, got: {err:?}"
    );

    // AC: basis_store slot (0, 1) must be None after the error (invalidated).
    assert!(
        basis_store.get(0, 1).is_none(),
        "basis_store.get(0, 1) must be None after solver error at stage 1"
    );

    // Stage 0 succeeded in iteration 2 — its basis was re-extracted.
    assert!(
        basis_store.get(0, 0).is_some(),
        "basis_store.get(0, 0) must be Some (stage 0 succeeded before error)"
    );
}

// ── New test: parallel cost agreement ────────────────────────────────────

/// AC: with 1-workspace and 4-workspace pools producing the same `cost_sum`.
///
/// Given the same input data, `run_forward_pass` with a single workspace
/// must produce identical `cost_sum` and `cost_sum_sq` values compared to
/// running with 4 workspaces. This verifies the static partitioning
/// produces deterministic results regardless of workspace count.
#[test]
#[allow(clippy::too_many_lines)]
fn test_forward_pass_parallel_cost_agreement() {
    let indexer = {
        let mut ix = StageIndexer::new(1, 0);
        ix.finalize_for_test();
        ix
    };
    let solution = fixed_solution(4, 100.0, indexer.theta, 30.0);
    let stochastic = make_stochastic_context_1_hydro_3_stages();
    let stages = make_stages_3();
    let fcf = FutureCostFunction::new(3, indexer.n_state, 2, 100, &[0; 3]);
    let horizon = HorizonMode::Finite { num_stages: 3 };
    let templates = vec![
        minimal_template_1_0(),
        minimal_template_1_0(),
        minimal_template_1_0(),
    ];
    let base_rows = vec![2usize, 2, 2];
    let initial_state = vec![0.0_f64; indexer.n_state];
    let n_scenarios = 10;

    let ctx = StageContext {
        templates: &templates,
        base_rows: &base_rows,
        noise_scale: &[],
        n_hydros: 0,
        n_load_buses: 0,
        load_balance_row_starts: &[],
        load_bus_indices: &[],
        block_counts_per_stage: &[1usize, 1, 1],
        ncs_max_gen: &[],
        ncs_allow_curtailment: &[],
        discount_factors: &[],
        cumulative_discount_factors: &[],
        stage_lag_transitions: &[],
        noise_group_ids: &[],
        downstream_par_order: 0,
    };

    // Run with 1 workspace.
    let mut ws1 = single_workspace(MockSolver::always_ok(solution.clone()), &indexer);
    let mut records1 = empty_records(n_scenarios * 3);
    let mut basis_store1 = BasisStore::new(n_scenarios, templates.len());
    let result1 = run_forward_pass(
        std::slice::from_mut(&mut ws1),
        &mut basis_store1,
        &ctx,
        &templates,
        &fcf,
        &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,
            dcs: None,
            noise_key_diag: None,
        },
        &ForwardPassBatch {
            local_forward_passes: n_scenarios,
            total_forward_passes: n_scenarios,
            iteration: 0,
            fwd_offset: 0,
            event_sender: None,
        },
        &mut records1,
    )
    .unwrap();

    // Run with 4 workspaces.
    let mut workspaces4: Vec<SolverWorkspace<MockSolver>> = (0..4)
        .map(|_| single_workspace(MockSolver::always_ok(solution.clone()), &indexer))
        .collect();
    let mut records4 = empty_records(n_scenarios * 3);
    let mut basis_store4 = BasisStore::new(n_scenarios, templates.len());
    let result4 = run_forward_pass(
        &mut workspaces4,
        &mut basis_store4,
        &ctx,
        &templates,
        &fcf,
        &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,
            dcs: None,
            noise_key_diag: None,
        },
        &ForwardPassBatch {
            local_forward_passes: n_scenarios,
            total_forward_passes: n_scenarios,
            iteration: 0,
            fwd_offset: 0,
            event_sender: None,
        },
        &mut records4,
    )
    .unwrap();

    assert_eq!(
        result1.scenario_costs.len(),
        result4.scenario_costs.len(),
        "scenario_costs length must be identical for 1 and 4 workspaces"
    );
    // Each scenario cost must be bit-identical regardless of workspace count.
    for (i, (c1, c4)) in result1
        .scenario_costs
        .iter()
        .zip(result4.scenario_costs.iter())
        .enumerate()
    {
        assert_eq!(
            c1.to_bits(),
            c4.to_bits(),
            "scenario_costs[{i}] must be bit-identical: 1-workspace={c1:.17e}, 4-workspace={c4:.17e}"
        );
    }
}

// ── New test: work distribution across 4 workspaces ──────────────────────

/// C1: verify that 10 scenarios distributed across 4 workspaces assign
/// each workspace `floor(10/4)` or `ceil(10/4)` scenarios.
///
/// With `n=10`, `n_workers=4`: `base=2`, `remainder=2`.
/// - Workers 0,1 receive 3 scenarios each → 3 * 3 stages = 9 solve calls.
/// - Workers 2,3 receive 2 scenarios each → 2 * 3 stages = 6 solve calls.
///
/// `MockSolver.statistics().solve_count` now returns `call_count`, so we
/// can verify each workspace performed its assigned number of LP solves.
#[allow(clippy::too_many_lines)]
#[test]
fn test_forward_pass_work_distribution() {
    let indexer = {
        let mut ix = StageIndexer::new(1, 0);
        ix.finalize_for_test();
        ix
    };
    let solution = fixed_solution(4, 100.0, indexer.theta, 30.0);
    let stochastic = make_stochastic_context_1_hydro_3_stages();
    let stages = make_stages_3();
    let fcf = FutureCostFunction::new(3, indexer.n_state, 2, 100, &[0; 3]);
    let horizon = HorizonMode::Finite { num_stages: 3 };
    let num_stages = 3usize;
    let templates = vec![
        minimal_template_1_0(),
        minimal_template_1_0(),
        minimal_template_1_0(),
    ];
    let base_rows = vec![2usize, 2, 2];
    let initial_state = vec![0.0_f64; indexer.n_state];
    let n_scenarios = 10usize;
    let n_workers = 4usize;

    let mut workspaces: Vec<SolverWorkspace<MockSolver>> = (0..n_workers)
        .map(|_| single_workspace(MockSolver::always_ok(solution.clone()), &indexer))
        .collect();
    let mut records = empty_records(n_scenarios * num_stages);
    let mut basis_store = BasisStore::new(n_scenarios, num_stages);

    let ctx = StageContext {
        templates: &templates,
        base_rows: &base_rows,
        noise_scale: &[],
        n_hydros: 0,
        n_load_buses: 0,
        load_balance_row_starts: &[],
        load_bus_indices: &[],
        block_counts_per_stage: &[1usize, 1, 1],
        ncs_max_gen: &[],
        ncs_allow_curtailment: &[],
        discount_factors: &[],
        cumulative_discount_factors: &[],
        stage_lag_transitions: &[],
        noise_group_ids: &[],
        downstream_par_order: 0,
    };
    let _result = run_forward_pass(
        &mut workspaces,
        &mut basis_store,
        &ctx,
        &templates,
        &fcf,
        &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,
            dcs: None,
            noise_key_diag: None,
        },
        &ForwardPassBatch {
            local_forward_passes: n_scenarios,
            total_forward_passes: n_scenarios,
            iteration: 0,
            fwd_offset: 0,
            event_sender: None,
        },
        &mut records,
    )
    .unwrap();

    // Verify each workspace performed the expected number of LP solves.
    // partition(10, 4, w): base=2, remainder=2.
    // Workers 0,1: 3 scenarios × 3 stages = 9 solves each.
    // Workers 2,3: 2 scenarios × 3 stages = 6 solves each.
    for (w, ws) in workspaces.iter().enumerate() {
        let (start_m, end_m) = partition(n_scenarios, n_workers, w);
        let assigned_scenarios = end_m - start_m;
        let expected_solves = assigned_scenarios * num_stages;

        // floor and ceil of n_scenarios / n_workers for boundary check.
        let floor_scenarios = n_scenarios / n_workers;
        let ceil_scenarios = n_scenarios.div_ceil(n_workers);
        assert!(
            assigned_scenarios == floor_scenarios || assigned_scenarios == ceil_scenarios,
            "worker {w} assigned {assigned_scenarios} scenarios, expected {floor_scenarios} or {ceil_scenarios}"
        );

        let actual_solves = usize::try_from(ws.solver.statistics().solve_count)
            .expect("solve_count fits in usize in tests");
        assert_eq!(
            actual_solves, expected_solves,
            "worker {w} (scenarios [{start_m}, {end_m})) performed {actual_solves} solves, expected {expected_solves}"
        );
    }

    // Verify the total solve count equals n_scenarios * num_stages.
    let total_solves: usize = workspaces
        .iter()
        .map(|ws| {
            usize::try_from(ws.solver.statistics().solve_count)
                .expect("solve_count fits in usize in tests")
        })
        .sum();
    assert_eq!(
        total_solves,
        n_scenarios * num_stages,
        "total solve count {total_solves} must equal n_scenarios * num_stages = {}",
        n_scenarios * num_stages
    );
}

// ── Truncation unit tests ────────────────────────────────────────────────

/// Build a `StochasticContext` for 1 hydro and 1 stage with the given
/// `mean_m3s` and `std_m3s`. Used by truncation tests.
#[allow(clippy::too_many_lines)]
fn make_stochastic_1h_1s(mean_m3s: f64, std_m3s: f64) -> StochasticContext {
    use std::collections::BTreeMap;

    use chrono::NaiveDate;
    use cobre_core::entities::hydro::{Hydro, HydroGenerationModel, HydroPenalties};
    use cobre_core::scenario::{
        CorrelationEntity, CorrelationGroup, CorrelationModel, CorrelationProfile, InflowModel,
    };
    use cobre_core::temporal::{
        Block, BlockMode, NoiseMethod, ScenarioSourceConfig, Stage, StageRiskConfig,
        StageStateConfig,
    };
    use cobre_core::{Bus, DeficitSegment, EntityId, SystemBuilder};

    let bus = Bus {
        id: EntityId(0),
        name: "B0".to_string(),
        deficit_segments: vec![DeficitSegment {
            depth_mw: None,
            cost_per_mwh: 1000.0,
        }],
        excess_cost: 0.0,
    };
    let hydro = Hydro {
        id: EntityId(1),
        name: "H1".to_string(),
        bus_id: EntityId(0),
        downstream_id: None,
        entry_stage_id: None,
        exit_stage_id: None,
        min_storage_hm3: 0.0,
        max_storage_hm3: 100.0,
        min_outflow_m3s: 0.0,
        max_outflow_m3s: None,
        generation_model: HydroGenerationModel::ConstantProductivity,
        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 stage = Stage {
        index: 0,
        id: 0,
        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: 3,
            noise_method: NoiseMethod::Saa,
        },
    };
    let inflow_model = InflowModel {
        hydro_id: EntityId(1),
        stage_id: 0,
        mean_m3s,
        std_m3s,
        ar_coefficients: vec![],
        residual_std_ratio: 1.0,
        annual: None,
    };
    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(vec![stage])
        .inflow_models(vec![inflow_model])
        .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()
}

/// Minimal stage template for N=1 hydro, L=0 PAR, with a single water-balance
/// row at position `base_row_idx`.
///
/// This is a three-row template:
/// - Row 0: storage fixing row
/// - Row 1: z-inflow definition row (at N*(1+L) = 1)
/// - Row 2: water-balance row (`base_rows`[t] = 2)
///
/// `row_lower[2]` encodes the deterministic inflow base (ζ * `mean_m3s`).
fn minimal_template_1_0_with_base(base_rhs: f64) -> StageTemplate {
    StageTemplate {
        num_cols: 4,
        num_rows: 3,
        num_nz: 1,
        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, base_rhs],
        row_upper: vec![0.0, 0.0, base_rhs],
        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(),
    }
}

/// Helper that runs `run_forward_pass` with 1 scenario, 1 stage, and returns
/// the `noise_buf` from the workspace after the call.
fn run_single_stage_forward(
    stochastic: &StochasticContext,
    inflow_method: InflowNonNegativityMethod,
    base_rhs: f64,
    noise_scale_val: f64,
) -> Vec<f64> {
    let indexer = {
        let mut ix = StageIndexer::new(1, 0);
        ix.finalize_for_test();
        ix
    };
    let solution = fixed_solution(4, 0.0, indexer.theta, 0.0);
    let solver = MockSolver::always_ok(solution);
    let fcf = FutureCostFunction::new(1, indexer.n_state, 1, 10, &[0; 1]);
    let horizon = HorizonMode::Finite { num_stages: 1 };
    let template = minimal_template_1_0_with_base(base_rhs);
    let templates = vec![template];
    // base_rows[t] = 1: the water-balance row is at row index 1.
    let base_rows = vec![2usize];
    let initial_state = vec![0.0_f64; indexer.n_state];
    let mut records = empty_records(1);
    let mut ws = single_workspace(solver, &indexer);
    let mut basis_store = BasisStore::new(1, 1);
    let noise_scale = vec![noise_scale_val];

    let ctx = StageContext {
        templates: &templates,
        base_rows: &base_rows,
        noise_scale: &noise_scale,
        n_hydros: 1,
        n_load_buses: 0,
        load_balance_row_starts: &[],
        load_bus_indices: &[],
        block_counts_per_stage: &[1usize],
        ncs_max_gen: &[],
        ncs_allow_curtailment: &[],
        discount_factors: &[],
        cumulative_discount_factors: &[],
        stage_lag_transitions: &[],
        noise_group_ids: &[],
        downstream_par_order: 0,
    };
    let stages = vec![Stage {
        index: 0,
        id: 0,
        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: 1,
            noise_method: NoiseMethod::Saa,
        },
    }];
    let _ = run_forward_pass(
        std::slice::from_mut(&mut ws),
        &mut basis_store,
        &ctx,
        &templates,
        &fcf,
        &TrainingContext {
            horizon: &horizon,
            indexer: &indexer,
            inflow_method: &inflow_method,
            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,
            dcs: None,
            noise_key_diag: None,
        },
        &ForwardPassBatch {
            local_forward_passes: 1,
            total_forward_passes: 1,
            iteration: 0,
            fwd_offset: 0,
            event_sender: None,
        },
        &mut records,
    )
    .unwrap();

    ws.scratch.noise_buf.clone()
}

/// AC: truncation clamps negative inflow noise.
///
/// Set up a 1-hydro, 1-stage system where the PAR mean is very large
/// negative (`mean_m3s = -1000.0`) and sigma is small (`std_m3s = 1.0`)
/// so that the deterministic base alone produces a hugely negative inflow.
/// Any sampled noise value will produce a negative inflow without truncation.
///
/// With truncation active: the noise buffer entry must be >= 0.0 because
/// `noise_buf[h] = base_rhs + noise_scale * clamped_eta`
/// = `zeta * (mean + sigma * eta_min)` = `zeta * 0.0` = 0.0 (or near it).
///
/// Concretely: `noise_buf[0] = base_rhs + noise_scale * clamped_eta`
/// where `base_rhs = zeta * (-1000)` and `noise_scale = zeta * 1`.
/// After clamping: `clamped_eta = max(eta, (0 - (-1000)) / 1) = 1000`.
/// So `noise_buf[0] = zeta * (-1000) + zeta * 1000 = 0`.
///
/// The actual PAR inflow = `mean + sigma * clamped_eta = -1000 + 1 * 1000 = 0 >= 0`.
#[test]
fn truncation_clamps_negative_inflow_noise() {
    // Deterministic base = -1000 m³/s (always produces negative inflow).
    // sigma = 1.0. For AR(0): zeta * mean = base_rhs, zeta * sigma = noise_scale.
    // Use zeta = 1.0 for simplicity (noise_scale = sigma).
    let mean_m3s = -1000.0_f64;
    let sigma = 1.0_f64;
    let zeta = 1.0_f64; // simplified for test: treat zeta=1
    let base_rhs = zeta * mean_m3s;
    let noise_scale_val = zeta * sigma;

    let stochastic = make_stochastic_1h_1s(mean_m3s, sigma);

    let noise_buf_truncation = run_single_stage_forward(
        &stochastic,
        InflowNonNegativityMethod::Truncation,
        base_rhs,
        noise_scale_val,
    );

    assert_eq!(noise_buf_truncation.len(), 1, "noise_buf must have 1 entry");
    // After truncation: noise_buf[0] = base_rhs + noise_scale * eta_clamped.
    // eta_clamped = max(eta, eta_min) where eta_min = (0 - mean) / sigma = 1000.
    // noise_buf[0] = -1000 + 1 * 1000 = 0.0 (exactly, no rounding).
    assert!(
        noise_buf_truncation[0] >= 0.0,
        "after truncation, noise_buf[0] must be >= 0 (inflow cannot be negative), got {}",
        noise_buf_truncation[0]
    );
}

/// AC: truncation does not clamp when inflow is positive.
///
/// With a very large positive mean (`mean_m3s = 1000.0`) and small sigma,
/// the PAR inflow is always positive for any sampled noise. The noise buffer
/// must be identical to the no-truncation path.
#[test]
fn truncation_no_clamp_when_inflow_positive() {
    let mean_m3s = 1000.0_f64;
    let sigma = 1.0_f64;
    let zeta = 1.0_f64;
    let base_rhs = zeta * mean_m3s;
    let noise_scale_val = zeta * sigma;

    let stochastic = make_stochastic_1h_1s(mean_m3s, sigma);

    let noise_buf_truncation = run_single_stage_forward(
        &stochastic,
        InflowNonNegativityMethod::Truncation,
        base_rhs,
        noise_scale_val,
    );
    let noise_buf_none = run_single_stage_forward(
        &stochastic,
        InflowNonNegativityMethod::None,
        base_rhs,
        noise_scale_val,
    );

    assert_eq!(noise_buf_truncation.len(), 1);
    assert_eq!(noise_buf_none.len(), 1);
    assert_eq!(
        noise_buf_truncation[0].to_bits(),
        noise_buf_none[0].to_bits(),
        "when inflow is positive, truncation must not alter the noise buffer (expected identical bits)"
    );
}

/// AC: `InflowNonNegativityMethod::None` produces unchanged behavior when
/// truncation code is present in the same function.
///
/// Uses the existing 3-stage fixture to verify no regression.
#[test]
#[allow(clippy::too_many_lines)]
fn none_method_unchanged_with_truncation_code_present() {
    let indexer = {
        let mut ix = StageIndexer::new(1, 0);
        ix.finalize_for_test();
        ix
    };
    let solution = fixed_solution(4, 100.0, indexer.theta, 30.0);
    let solver = MockSolver::always_ok(solution);
    let fcf = FutureCostFunction::new(3, indexer.n_state, 2, 100, &[0; 3]);
    let config = TrainingConfig {
        loop_config: LoopConfig {
            forward_passes: 2,
            max_iterations: 100,
            start_iteration: 0,
            n_fwd_threads: 1,
            max_blocks: 1,
            stopping_rules: StoppingRuleSet {
                rules: vec![StoppingRule::IterationLimit { limit: 100 }],
                mode: StoppingMode::Any,
            },
        },
        cut_management: CutManagementConfig {
            cut_selection: None,
            budget: None,
            cut_activity_tolerance: 0.0,
            warm_start_cuts: 0,
            risk_measures: vec![RiskMeasure::Expectation],
        },
        events: EventConfig {
            event_sender: None,
            checkpoint_interval: None,
            shutdown_flag: None,
            export_states: false,
        },
    };

    let horizon = HorizonMode::Finite { num_stages: 3 };
    let templates = vec![
        minimal_template_1_0(),
        minimal_template_1_0(),
        minimal_template_1_0(),
    ];
    let base_rows = vec![2usize, 2, 2];
    let initial_state = vec![0.0_f64; indexer.n_state];
    let mut records = empty_records(2 * 3);
    let stochastic = make_stochastic_context_1_hydro_3_stages();
    let stages = make_stages_3();
    let mut ws = single_workspace(solver, &indexer);
    let mut basis_store =
        BasisStore::new(config.loop_config.forward_passes as usize, templates.len());

    let ctx = StageContext {
        templates: &templates,
        base_rows: &base_rows,
        noise_scale: &[],
        n_hydros: 0,
        n_load_buses: 0,
        load_balance_row_starts: &[],
        load_bus_indices: &[],
        block_counts_per_stage: &[1usize, 1, 1],
        ncs_max_gen: &[],
        ncs_allow_curtailment: &[],
        discount_factors: &[],
        cumulative_discount_factors: &[],
        stage_lag_transitions: &[],
        noise_group_ids: &[],
        downstream_par_order: 0,
    };
    let result = run_forward_pass(
        std::slice::from_mut(&mut ws),
        &mut basis_store,
        &ctx,
        &templates,
        &fcf,
        &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,
            dcs: None,
            noise_key_diag: None,
        },
        &ForwardPassBatch {
            local_forward_passes: config.loop_config.forward_passes as usize,
            total_forward_passes: config.loop_config.forward_passes as usize,
            iteration: 0,
            fwd_offset: 0,
            event_sender: None,
        },
        &mut records,
    )
    .unwrap();

    // Regression guard: same assertions as `ac_two_scenarios_three_stages_fixed_solution`.
    assert_eq!(result.scenario_costs.len(), 2);
    for (i, record) in records.iter().enumerate() {
        assert_eq!(
            record.stage_cost, 70_000_000.0,
            "none_method: record[{i}].stage_cost should be 70_000_000.0 ((objective - theta) * COST_SCALE_FACTOR)"
        );
    }
}

// ── Load noise test helpers ─────────────────────────────────────────────

/// Build a `StochasticContext` with 1 hydro and 1 stochastic load bus over
/// a single stage.
///
/// `mean_mw` and `std_mw` control the load bus noise model.  An empty
/// correlation model is used so the two noise entities are treated as
/// independent standard normals.
#[allow(clippy::too_many_lines)]
fn make_stochastic_context_1_hydro_1_load_bus(mean_mw: f64, std_mw: f64) -> StochasticContext {
    let bus0 = Bus {
        id: EntityId(0),
        name: "B0".to_string(),
        deficit_segments: vec![DeficitSegment {
            depth_mw: None,
            cost_per_mwh: 1000.0,
        }],
        excess_cost: 0.0,
    };
    let bus1 = Bus {
        id: EntityId(1),
        name: "B1".to_string(),
        deficit_segments: vec![DeficitSegment {
            depth_mw: None,
            cost_per_mwh: 1000.0,
        }],
        excess_cost: 0.0,
    };
    let hydro = Hydro {
        id: EntityId(10),
        name: "H10".to_string(),
        bus_id: EntityId(0),
        downstream_id: None,
        entry_stage_id: None,
        exit_stage_id: None,
        min_storage_hm3: 0.0,
        max_storage_hm3: 100.0,
        min_outflow_m3s: 0.0,
        max_outflow_m3s: None,
        generation_model: HydroGenerationModel::ConstantProductivity,
        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: cobre_core::entities::hydro::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 stage = Stage {
        index: 0,
        id: 0,
        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: 3,
            noise_method: NoiseMethod::Saa,
        },
    };
    let inflow_model = InflowModel {
        hydro_id: EntityId(10),
        stage_id: 0,
        mean_m3s: 100.0,
        std_m3s: 20.0,
        ar_coefficients: vec![],
        residual_std_ratio: 1.0,
        annual: None,
    };
    let load_model = LoadModel {
        bus_id: EntityId(1),
        stage_id: 0,
        mean_mw,
        std_mw,
    };
    // No correlation profile: entities are treated as independent.
    let correlation = CorrelationModel {
        method: "spectral".to_string(),
        profiles: std::collections::BTreeMap::new(),
        schedule: vec![],
    };
    let system = SystemBuilder::new()
        .buses(vec![bus0, bus1])
        .hydros(vec![hydro])
        .stages(vec![stage])
        .inflow_models(vec![inflow_model])
        .load_models(vec![load_model])
        .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()
}

// ── New test: parallel infeasibility propagation ──────────────────────────

/// C3: when one of 4 workspaces returns `SolverError::Infeasible`, the
/// error propagates as `SddpError::Infeasible` with correct stage and
/// scenario indices.
///
/// Worker 1 handles scenarios [3, 6) (3 scenarios) with 3 stages each.
/// Its solver is configured to fail on call 0 (scenario 3, stage 0).
/// Expected: `SddpError::Infeasible { stage: 0, scenario: 3, .. }`.
#[test]
fn test_forward_pass_parallel_infeasibility() {
    let indexer = {
        let mut ix = StageIndexer::new(1, 0);
        ix.finalize_for_test();
        ix
    };
    let solution = fixed_solution(4, 100.0, indexer.theta, 30.0);
    let stochastic = make_stochastic_context_1_hydro_3_stages();
    let stages = make_stages_3();
    let fcf = FutureCostFunction::new(3, indexer.n_state, 2, 100, &[0; 3]);
    let horizon = HorizonMode::Finite { num_stages: 3 };
    let num_stages = 3usize;
    let templates = vec![
        minimal_template_1_0(),
        minimal_template_1_0(),
        minimal_template_1_0(),
    ];
    let base_rows = vec![2usize, 2, 2];
    let initial_state = vec![0.0_f64; indexer.n_state];
    let n_scenarios = 10usize;
    let n_workers = 4usize;

    // Worker 1 handles scenarios [3, 6). Its first solve call (call index 0
    // within that worker) corresponds to scenario 3, stage 0.
    let mut workspaces: Vec<SolverWorkspace<MockSolver>> = (0..n_workers)
        .map(|w| {
            let solver = if w == 1 {
                // Fail on the first solve call this worker makes.
                MockSolver::infeasible_on(solution.clone(), 0)
            } else {
                MockSolver::always_ok(solution.clone())
            };
            single_workspace(solver, &indexer)
        })
        .collect();

    let mut records = empty_records(n_scenarios * num_stages);
    let mut basis_store = BasisStore::new(n_scenarios, num_stages);

    let ctx = StageContext {
        templates: &templates,
        base_rows: &base_rows,
        noise_scale: &[],
        n_hydros: 0,
        n_load_buses: 0,
        load_balance_row_starts: &[],
        load_bus_indices: &[],
        block_counts_per_stage: &[1usize, 1, 1],
        ncs_max_gen: &[],
        ncs_allow_curtailment: &[],
        discount_factors: &[],
        cumulative_discount_factors: &[],
        stage_lag_transitions: &[],
        noise_group_ids: &[],
        downstream_par_order: 0,
    };
    let result = run_forward_pass(
        &mut workspaces,
        &mut basis_store,
        &ctx,
        &templates,
        &fcf,
        &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,
            dcs: None,
            noise_key_diag: None,
        },
        &ForwardPassBatch {
            local_forward_passes: n_scenarios,
            total_forward_passes: n_scenarios,
            iteration: 0,
            fwd_offset: 0,
            event_sender: None,
        },
        &mut records,
    );

    // Worker 1's partition: partition(10, 4, 1) → start_m=3.
    // The first solve in that worker is scenario 3, stage 0.
    match result {
        Err(crate::SddpError::Infeasible {
            stage,
            scenario,
            iteration,
        }) => {
            assert_eq!(
                stage, 0,
                "infeasible stage must be 0 (first stage of worker 1)"
            );
            assert_eq!(
                scenario, 3,
                "infeasible scenario must be 3 (start_m of worker 1)"
            );
            assert_eq!(
                iteration, 0,
                "iteration must be 0 (first training iteration)"
            );
        }
        Err(other) => panic!("expected SddpError::Infeasible, got: {other:?}"),
        Ok(_) => panic!("expected Err(SddpError::Infeasible), got Ok"),
    }
}

// ── Load noise wiring tests ──────────────────────────────────────────────

/// Verify that the load balance row is patched to a positive value when
/// `mean_mw + std_mw * eta > 0`.
///
/// With `mean_mw = 300.0` and `std_mw = 30.0`, any standard-normal draw `eta`
/// satisfying `|eta| <= 10` produces a positive realization, which holds with
/// overwhelming probability.  After a single-scenario forward pass the
/// `load_rhs_buf` must contain a positive value equal to
/// `max(0, 300 + 30 * eta) * block_factor`.  Since no load factors file is
/// supplied, `block_factor = 1.0`, so `load_rhs_buf[0] = max(0, 300 + 30 * eta)`.
#[test]
#[allow(clippy::too_many_lines)]
fn forward_pass_load_noise_positive_realization() {
    let n_load_buses = 1usize;
    let stochastic = make_stochastic_context_1_hydro_1_load_bus(300.0, 30.0);
    let indexer = {
        let mut ix = StageIndexer::new(1, 0);
        ix.finalize_for_test();
        ix
    };
    let patch_buf = crate::lp_builder::PatchBuffer::new(1, 0, n_load_buses, 1, 0, 0);
    let mut ws = SolverWorkspace {
        rank: 0,
        worker_id: 0,
        solver: ProfiledSolver::new(MockSolver::always_ok(fixed_solution(
            4,
            100.0,
            indexer.theta,
            30.0,
        ))),
        patch_buf,
        current_state: Vec::with_capacity(indexer.n_state),
        scratch: crate::workspace::ScratchBuffers {
            noise_buf: Vec::with_capacity(1),
            inflow_m3s_buf: Vec::with_capacity(1),
            lag_matrix_buf: Vec::with_capacity(0),
            par_inflow_buf: Vec::with_capacity(1),
            eta_floor_buf: Vec::with_capacity(1),
            zero_targets_buf: vec![0.0_f64; 1],
            ncs_col_upper_buf: Vec::new(),
            ncs_col_lower_buf: Vec::new(),
            ncs_col_indices_buf: Vec::new(),
            load_rhs_buf: Vec::with_capacity(n_load_buses),
            row_lower_buf: Vec::new(),
            z_inflow_rhs_buf: Vec::new(),
            effective_eta_buf: Vec::new(),
            unscaled_primal: Vec::new(),
            unscaled_dual: Vec::new(),
            lag_accumulator: vec![],
            lag_weight_accum: 0.0,
            downstream_accumulator: Vec::new(),
            downstream_weight_accum: 0.0,
            downstream_completed_lags: Vec::new(),
            downstream_n_completed: 0,
            recon_slot_lookup: Vec::new(),
            trajectory_costs_buf: Vec::new(),
            raw_noise_buf: Vec::new(),
            perm_scratch: Vec::new(),
            anticipated_state_buf: Vec::new(),
        },
        scratch_basis: Basis::new(0, 0),
        backward_accum: BackwardAccumulators::default(),
        worker_timing_buf: cobre_core::WorkerPhaseTimings::default(),
    };

    let templates = vec![minimal_template_1_0_with_base(100.0)];
    let base_rows = vec![2usize];
    let initial_state = vec![0.0_f64; indexer.n_state];
    let mut records = empty_records(1);
    let fcf = FutureCostFunction::new(1, indexer.n_state, 1, 10, &[0; 1]);
    let horizon = HorizonMode::Finite { num_stages: 1 };
    let mut basis_store = BasisStore::new(1, 1);
    let load_balance_row_starts = vec![10usize];
    let load_bus_indices = vec![0usize];
    let block_counts_per_stage = vec![1usize];

    let ctx = StageContext {
        templates: &templates,
        base_rows: &base_rows,
        noise_scale: &[1.0],
        n_hydros: 1,
        n_load_buses,
        load_balance_row_starts: &load_balance_row_starts,
        load_bus_indices: &load_bus_indices,
        block_counts_per_stage: &block_counts_per_stage,
        ncs_max_gen: &[],
        ncs_allow_curtailment: &[],
        discount_factors: &[],
        cumulative_discount_factors: &[],
        stage_lag_transitions: &[],
        noise_group_ids: &[],
        downstream_par_order: 0,
    };
    let _fwd = run_forward_pass(
        std::slice::from_mut(&mut ws),
        &mut basis_store,
        &ctx,
        &templates,
        &fcf,
        &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: &[],
            historical_library: None,
            external_inflow_library: None,
            external_load_library: None,
            external_ncs_library: None,
            recent_accum_seed: &[],
            recent_weight_seed: 0.0,
            dcs: None,
            noise_key_diag: None,
        },
        &ForwardPassBatch {
            local_forward_passes: 1,
            total_forward_passes: 1,
            iteration: 0,
            fwd_offset: 0,
            event_sender: None,
        },
        &mut records,
    )
    .unwrap();

    assert_eq!(
        ws.scratch.load_rhs_buf.len(),
        n_load_buses,
        "load_rhs_buf must have 1 entry (1 load bus x 1 block)"
    );
    assert!(
        ws.scratch.load_rhs_buf[0] > 0.0,
        "load realization must be positive with mean=300, std=30: got {}",
        ws.scratch.load_rhs_buf[0]
    );

    let cat4_start = 1;
    assert_eq!(
        ws.patch_buf.lower[cat4_start], ws.scratch.load_rhs_buf[0],
        "patch_buf lower must equal load_rhs_buf[0]"
    );
    assert_eq!(
        ws.patch_buf.upper[cat4_start], ws.scratch.load_rhs_buf[0],
        "patch_buf upper must equal load_rhs_buf[0] (equality constraint)"
    );
    assert_eq!(
        ws.patch_buf.indices[cat4_start], 10,
        "patch index must be load_balance_row_starts[0] + 0 * n_blks"
    );
}

/// Verify that a load realization that would be negative is clamped to zero.
///
/// With `mean_mw = -1000.0` and `std_mw = 1.0`, any standard-normal draw
/// (bounded to roughly +-5 in practice) produces `mean + std * eta ~= -1000`,
/// which must be clamped to `0.0` before block factor scaling.
#[test]
#[allow(clippy::too_many_lines)]
fn forward_pass_load_noise_clamped_to_zero() {
    let n_load_buses = 1usize;
    let stochastic = make_stochastic_context_1_hydro_1_load_bus(-1000.0, 1.0);
    let indexer = {
        let mut ix = StageIndexer::new(1, 0);
        ix.finalize_for_test();
        ix
    };
    let patch_buf = crate::lp_builder::PatchBuffer::new(1, 0, n_load_buses, 1, 0, 0);
    let mut ws = SolverWorkspace {
        rank: 0,
        worker_id: 0,
        solver: ProfiledSolver::new(MockSolver::always_ok(fixed_solution(
            4,
            100.0,
            indexer.theta,
            30.0,
        ))),
        patch_buf,
        current_state: Vec::with_capacity(indexer.n_state),
        scratch: crate::workspace::ScratchBuffers {
            noise_buf: Vec::with_capacity(1),
            inflow_m3s_buf: Vec::with_capacity(1),
            lag_matrix_buf: Vec::with_capacity(0),
            par_inflow_buf: Vec::with_capacity(1),
            eta_floor_buf: Vec::with_capacity(1),
            zero_targets_buf: vec![0.0_f64; 1],
            ncs_col_upper_buf: Vec::new(),
            ncs_col_lower_buf: Vec::new(),
            ncs_col_indices_buf: Vec::new(),
            load_rhs_buf: Vec::with_capacity(n_load_buses),
            row_lower_buf: Vec::new(),
            z_inflow_rhs_buf: Vec::new(),
            effective_eta_buf: Vec::new(),
            unscaled_primal: Vec::new(),
            unscaled_dual: Vec::new(),
            lag_accumulator: vec![],
            lag_weight_accum: 0.0,
            downstream_accumulator: Vec::new(),
            downstream_weight_accum: 0.0,
            downstream_completed_lags: Vec::new(),
            downstream_n_completed: 0,
            recon_slot_lookup: Vec::new(),
            trajectory_costs_buf: Vec::new(),
            raw_noise_buf: Vec::new(),
            perm_scratch: Vec::new(),
            anticipated_state_buf: Vec::new(),
        },
        scratch_basis: Basis::new(0, 0),
        backward_accum: BackwardAccumulators::default(),
        worker_timing_buf: cobre_core::WorkerPhaseTimings::default(),
    };

    let templates = vec![minimal_template_1_0_with_base(100.0)];
    let base_rows = vec![2usize];
    let initial_state = vec![0.0_f64; indexer.n_state];
    let mut records = empty_records(1);
    let fcf = FutureCostFunction::new(1, indexer.n_state, 1, 10, &[0; 1]);
    let horizon = HorizonMode::Finite { num_stages: 1 };
    let mut basis_store = BasisStore::new(1, 1);
    let load_balance_row_starts = vec![10usize];
    let load_bus_indices = vec![0usize];
    let block_counts_per_stage = vec![1usize];

    let ctx = StageContext {
        templates: &templates,
        base_rows: &base_rows,
        noise_scale: &[1.0],
        n_hydros: 1,
        n_load_buses,
        load_balance_row_starts: &load_balance_row_starts,
        load_bus_indices: &load_bus_indices,
        block_counts_per_stage: &block_counts_per_stage,
        ncs_max_gen: &[],
        ncs_allow_curtailment: &[],
        discount_factors: &[],
        cumulative_discount_factors: &[],
        stage_lag_transitions: &[],
        noise_group_ids: &[],
        downstream_par_order: 0,
    };
    let _fwd = run_forward_pass(
        std::slice::from_mut(&mut ws),
        &mut basis_store,
        &ctx,
        &templates,
        &fcf,
        &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: &[],
            historical_library: None,
            external_inflow_library: None,
            external_load_library: None,
            external_ncs_library: None,
            recent_accum_seed: &[],
            recent_weight_seed: 0.0,
            dcs: None,
            noise_key_diag: None,
        },
        &ForwardPassBatch {
            local_forward_passes: 1,
            total_forward_passes: 1,
            iteration: 0,
            fwd_offset: 0,
            event_sender: None,
        },
        &mut records,
    )
    .unwrap();

    assert_eq!(
        ws.scratch.load_rhs_buf.len(),
        n_load_buses,
        "load_rhs_buf must have 1 entry (1 load bus x 1 block)"
    );
    assert_eq!(
        ws.scratch.load_rhs_buf[0], 0.0,
        "realization with mean=-1000 must be clamped to 0.0, got {}",
        ws.scratch.load_rhs_buf[0]
    );

    let cat4_start = 1;
    assert_eq!(
        ws.patch_buf.lower[cat4_start], 0.0,
        "patch lower must be 0.0 (clamped)"
    );
    assert_eq!(
        ws.patch_buf.upper[cat4_start], 0.0,
        "patch upper must be 0.0 (clamped)"
    );
}

/// Verify that when `n_load_buses == 0` no load patches are applied and
/// `forward_patch_count()` equals `N`.
///
/// With N=1 hydro, L=0 PAR order, and no load buses the patch count must be
/// exactly `1`.
#[test]
fn forward_pass_no_load_buses_unchanged() {
    // Use the existing 1-hydro-3-stage context that has no load buses.
    let stochastic = make_stochastic_context_1_hydro_3_stages();
    let stages = make_stages_3();
    let indexer = {
        let mut ix = StageIndexer::new(1, 0);
        ix.finalize_for_test();
        ix
    };
    let solution = fixed_solution(4, 100.0, indexer.theta, 30.0);
    let mut ws = single_workspace(MockSolver::always_ok(solution), &indexer);

    let templates = vec![
        minimal_template_1_0(),
        minimal_template_1_0(),
        minimal_template_1_0(),
    ];
    let base_rows = vec![2usize, 2, 2];
    let initial_state = vec![0.0_f64; indexer.n_state];
    let mut records = empty_records(3); // 1 scenario * 3 stages
    let fcf = FutureCostFunction::new(3, indexer.n_state, 1, 10, &[0; 3]);
    let horizon = HorizonMode::Finite { num_stages: 3 };
    let mut basis_store = BasisStore::new(1, 3);

    let ctx = StageContext {
        templates: &templates,
        base_rows: &base_rows,
        noise_scale: &[], // noise_scale empty when n_hydros=0
        n_hydros: 0,      // skip inflow noise loop (minimal_template_1_0 has 1 row)
        n_load_buses: 0,  // no load patches
        load_balance_row_starts: &[],
        load_bus_indices: &[],
        block_counts_per_stage: &[1, 1, 1],
        ncs_max_gen: &[],
        ncs_allow_curtailment: &[],
        discount_factors: &[],
        cumulative_discount_factors: &[],
        stage_lag_transitions: &[],
        noise_group_ids: &[],
        downstream_par_order: 0,
    };
    let _fwd = run_forward_pass(
        std::slice::from_mut(&mut ws),
        &mut basis_store,
        &ctx,
        &templates,
        &fcf,
        &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,
            dcs: None,
            noise_key_diag: None,
        },
        &ForwardPassBatch {
            local_forward_passes: 1,
            total_forward_passes: 1,
            iteration: 0,
            fwd_offset: 0,
            event_sender: None,
        },
        &mut records,
    )
    .unwrap();

    // With n_load_buses=0, active_load_patches stays 0.
    // forward_patch_count = N = 1.
    // The PatchBuffer was constructed for 1 hydro (single_workspace uses indexer.hydro_count=1).
    assert_eq!(
        ws.patch_buf.forward_patch_count(),
        1,
        "forward_patch_count must be N=1 when n_load_buses=0, got {}",
        ws.patch_buf.forward_patch_count()
    );
    // load_rhs_buf must remain empty (never pushed to).
    assert!(
        ws.scratch.load_rhs_buf.is_empty(),
        "load_rhs_buf must be empty when n_load_buses=0"
    );
}

// ── Tests for build_delta_cut_row_batch_into ─────────────────────────

fn empty_delta_batch() -> RowBatch {
    RowBatch {
        num_rows: 0,
        row_starts: Vec::new(),
        col_indices: Vec::new(),
        values: Vec::new(),
        row_lower: Vec::new(),
        row_upper: Vec::new(),
    }
}

#[test]
fn test_build_delta_empty_pool() {
    // Empty pool → num_rows == 0, row_starts == [0], col_indices empty.
    let fcf = FutureCostFunction::new(2, 1, 1, 10, &[0; 2]);
    let indexer = {
        let mut ix = StageIndexer::new(1, 0);
        ix.finalize_for_test();
        ix
    };
    let mut batch = empty_delta_batch();

    build_delta_cut_row_batch_into(&mut batch, &fcf, 0, &indexer, &[], 1);

    assert_eq!(batch.num_rows, 0);
    assert_eq!(batch.row_starts, vec![0_i32]);
    assert!(batch.col_indices.is_empty());
    assert!(batch.values.is_empty());
    assert!(batch.row_lower.is_empty());
    assert!(batch.row_upper.is_empty());
}

#[test]
fn test_build_delta_single_iteration_filter() {
    // Pool has cuts at iterations 1, 2, 3; calling with current_iteration=2
    // emits only the iteration-2 cut.
    //
    // FCF: 2 stages, 1 state dimension, 1 forward pass, 10 max iterations,
    // 0 warm-start cuts per stage.
    let mut fcf = FutureCostFunction::new(2, 1, 1, 10, &[0; 2]);
    // iteration=1, fwd_idx=0: slot = 0 + 1*1 + 0 = 1
    fcf.add_cut(0, 1, 0, 10.0, &[1.0]);
    // iteration=2, fwd_idx=0: slot = 0 + 2*1 + 0 = 2
    fcf.add_cut(0, 2, 0, 20.0, &[2.0]);
    // iteration=3, fwd_idx=0: slot = 0 + 3*1 + 0 = 3
    fcf.add_cut(0, 3, 0, 30.0, &[3.0]);

    let indexer = {
        let mut ix = StageIndexer::new(1, 0);
        ix.finalize_for_test();
        ix
    };
    let mut batch = empty_delta_batch();

    build_delta_cut_row_batch_into(&mut batch, &fcf, 0, &indexer, &[], 2);

    assert_eq!(batch.num_rows, 1);
    assert_eq!(batch.row_lower, vec![20.0]);
    assert_eq!(batch.row_starts, vec![0_i32, 2_i32]);
    // The cut emitted must carry iteration-2's coefficient (-2.0).
    assert_eq!(batch.values[0], -2.0);
}

#[test]
fn test_build_delta_skips_deactivated_cuts() {
    // Pool has cuts at iteration 1, some deactivated; only active
    // iteration-1 cuts are emitted.
    let mut fcf = FutureCostFunction::new(2, 1, 2, 10, &[0; 2]);
    // iteration=1, fwd_idx=0: slot = 0 + 1*2 + 0 = 2
    fcf.add_cut(0, 1, 0, 10.0, &[1.0]);
    // iteration=1, fwd_idx=1: slot = 0 + 1*2 + 1 = 3
    fcf.add_cut(0, 1, 1, 20.0, &[2.0]);

    // Deactivate slot 2 (the first iteration-1 cut).
    fcf.pools[0].deactivate(&[2]);

    let indexer = {
        let mut ix = StageIndexer::new(1, 0);
        ix.finalize_for_test();
        ix
    };
    let mut batch = empty_delta_batch();

    build_delta_cut_row_batch_into(&mut batch, &fcf, 0, &indexer, &[], 1);

    // Only slot 3 (intercept=20.0) should appear.
    assert_eq!(batch.num_rows, 1);
    assert_eq!(batch.row_lower, vec![20.0]);
}

#[test]
fn test_build_delta_excludes_warm_start_cuts() {
    // Pool seeded with a warm-start cut AND one training iteration cut.
    // Delta call with current_iteration=1 must exclude the warm-start row.
    use cobre_io::OwnedPolicyCutRecord;

    let warm_record = OwnedPolicyCutRecord {
        cut_id: 0,
        slot_index: 0,
        coefficients: vec![5.0],
        intercept: 99.0,
        iteration: 0,
        forward_pass_index: 0,
        is_active: true,
    };
    let mut pool = crate::cut::pool::CutPool::new_with_warm_start(1, 2, 10, &[warm_record]);
    // Now add a training cut at iteration=1, fwd_idx=0:
    // slot = warm_start_count(1) + 1*2 + 0 = 3
    pool.add_cut(1, 0, 7.0, &[1.0]);

    // Build an FCF with 2 stages (n_state=1, 2 fwd passes, 10 max iters).
    let mut fcf = FutureCostFunction::new(2, 1, 2, 10, &[0; 2]);
    fcf.pools[0] = pool;

    let indexer = {
        let mut ix = StageIndexer::new(1, 0);
        ix.finalize_for_test();
        ix
    };
    let mut batch = empty_delta_batch();

    build_delta_cut_row_batch_into(&mut batch, &fcf, 0, &indexer, &[], 1);

    // Warm-start cut (intercept=99.0) must be excluded; training cut
    // (intercept=7.0) must be present.
    assert_eq!(batch.num_rows, 1);
    assert_eq!(batch.row_lower, vec![7.0]);
}

#[test]
fn test_build_delta_matches_full_batch_when_pool_has_only_current_iter() {
    // When the pool contains only cuts from current_iteration, delta and
    // full builders must produce byte-identical output.
    let mut fcf = FutureCostFunction::new(2, 1, 2, 10, &[0; 2]);
    // iteration=1, fwd_idx=0: slot = 1*2+0 = 2
    fcf.add_cut(0, 1, 0, 10.0, &[1.0]);
    // iteration=1, fwd_idx=1: slot = 1*2+1 = 3
    fcf.add_cut(0, 1, 1, 20.0, &[3.0]);

    let indexer = {
        let mut ix = StageIndexer::new(1, 0);
        ix.finalize_for_test();
        ix
    };

    let mut batch_full = empty_delta_batch();
    build_cut_row_batch_into(&mut batch_full, &fcf, 0, &indexer, &[]);

    let mut batch_delta = empty_delta_batch();
    build_delta_cut_row_batch_into(&mut batch_delta, &fcf, 0, &indexer, &[], 1);

    assert_eq!(batch_delta.num_rows, batch_full.num_rows);
    assert_eq!(batch_delta.row_starts, batch_full.row_starts);
    assert_eq!(batch_delta.col_indices, batch_full.col_indices);
    assert_eq!(batch_delta.values, batch_full.values);
    assert_eq!(batch_delta.row_lower, batch_full.row_lower);
    assert_eq!(batch_delta.row_upper, batch_full.row_upper);
}

#[test]
fn test_build_delta_sparse_path() {
    // StageIndexer with non-empty nonzero_state_indices (sparse path).
    // Verify that the emitted col_indices for the cut contain exactly
    // nonzero_state_indices.len() + 1 entries (mask entries plus theta).
    //
    // StageIndexer::new(n_hydro, n_lag): n_state = n_hydro * (1 + n_lag)
    // With n_hydro=2, n_lag=0: n_state=2, no lags, sparse mask is empty.
    // We need a lag to get a nonzero_state_indices mask.
    // With n_hydro=1, n_lag=1: n_state=2, nonzero_state_indices=[0,1] (len=2)
    // for a cut that touches both state components.
    //
    // Actually we verify against the existing build_cut_row_batch_into for
    // correctness, which already tests the sparse path thoroughly.
    // Here we just verify col_indices.len() == mask.len() + 1 per row.

    // n_hydro=1, n_lag=1: n_state=2 (vol + lag).
    // nonzero_state_indices should be non-empty (check via indexer).
    let indexer = {
        let mut ix = StageIndexer::new(1, 1);
        ix.finalize_for_test();
        ix
    };
    // nonzero_state_indices is the mask for non-trivially-zero state dims.
    let mask_len = indexer.nonzero_state_indices.len();

    // Only proceed if this indexer actually uses the sparse path.
    if mask_len == 0 {
        // Sparse path not active for this indexer; skip the assertion.
        return;
    }

    let mut fcf = FutureCostFunction::new(2, indexer.n_state, 1, 10, &[0; 2]);
    fcf.add_cut(0, 1, 0, 5.0, &vec![1.0; indexer.n_state]);

    let mut batch = empty_delta_batch();
    build_delta_cut_row_batch_into(&mut batch, &fcf, 0, &indexer, &[], 1);

    assert_eq!(batch.num_rows, 1);
    // Each row: mask_len state entries + 1 theta entry.
    assert_eq!(batch.col_indices.len(), mask_len + 1);
}

#[test]
fn test_build_delta_reuses_out_buffer() {
    // Call twice; second call must produce correct output even when `batch`
    // had stale data from the first call.
    let mut fcf = FutureCostFunction::new(2, 1, 1, 10, &[0; 2]);
    fcf.add_cut(0, 1, 0, 11.0, &[1.0]);
    fcf.add_cut(0, 2, 0, 22.0, &[2.0]);

    let indexer = {
        let mut ix = StageIndexer::new(1, 0);
        ix.finalize_for_test();
        ix
    };
    let mut batch = empty_delta_batch();

    // First call: iteration 1 → should yield the iteration-1 cut.
    build_delta_cut_row_batch_into(&mut batch, &fcf, 0, &indexer, &[], 1);
    assert_eq!(batch.num_rows, 1);
    assert_eq!(batch.row_lower, vec![11.0]);

    // Second call: iteration 2 → stale data from first call must be gone.
    build_delta_cut_row_batch_into(&mut batch, &fcf, 0, &indexer, &[], 2);
    assert_eq!(batch.num_rows, 1);
    assert_eq!(batch.row_lower, vec![22.0]);
    assert_eq!(batch.row_starts.len(), 2); // [0, 2]
}

#[test]
fn test_build_delta_clears_row_starts() {
    // batch.row_starts[0] must be 0 regardless of prior state.
    let mut fcf = FutureCostFunction::new(2, 1, 1, 10, &[0; 2]);
    fcf.add_cut(0, 1, 0, 5.0, &[1.0]);

    let indexer = {
        let mut ix = StageIndexer::new(1, 0);
        ix.finalize_for_test();
        ix
    };

    // Pre-populate batch with garbage.
    let mut batch = RowBatch {
        num_rows: 5,
        row_starts: vec![0_i32, 2, 4, 6, 8, 10],
        col_indices: vec![0_i32; 10],
        values: vec![99.0_f64; 10],
        row_lower: vec![0.0_f64; 5],
        row_upper: vec![0.0_f64; 5],
    };

    build_delta_cut_row_batch_into(&mut batch, &fcf, 0, &indexer, &[], 1);

    assert_eq!(batch.row_starts[0], 0_i32);
    assert_eq!(batch.num_rows, 1);
    // Prior garbage must be gone.
    assert_eq!(batch.row_starts.len(), 2);
}

/// `build_delta_cut_row_batch_into` with a pool containing only warm-start
/// cuts must emit zero rows regardless of the requested iteration.
#[test]
fn build_delta_cut_row_batch_into_skips_warm_start_slots() {
    use crate::cut::pool::CutPool;
    use cobre_io::OwnedPolicyCutRecord;

    // One warm-start cut at slot 0.
    let ws_record = OwnedPolicyCutRecord {
        cut_id: 0,
        slot_index: 0,
        coefficients: vec![1.0],
        intercept: 99.0,
        iteration: 0,
        forward_pass_index: 0,
        is_active: true,
    };
    let mut pool = CutPool::new_with_warm_start(1, 1, 10, &[ws_record]);
    // One iteration-1 cut at slot 1 (warm_start_count=1, so slot = 1+1*1+0 = 2,
    // but new_with_warm_start sets warm_start_count=1 → slot = 1+1*1+0 = 2).
    pool.add_cut(1, 0, 7.0, &[1.0]);

    let mut fcf = FutureCostFunction::new(2, 1, 1, 10, &[0; 2]);
    fcf.pools[0] = pool;

    let indexer = {
        let mut ix = StageIndexer::new(1, 0);
        ix.finalize_for_test();
        ix
    };
    let mut batch = RowBatch {
        num_rows: 0,
        row_starts: Vec::new(),
        col_indices: Vec::new(),
        values: Vec::new(),
        row_lower: Vec::new(),
        row_upper: Vec::new(),
    };

    build_delta_cut_row_batch_into(&mut batch, &fcf, 0, &indexer, &[], 1);

    // Warm-start slot must be excluded; only the iteration-1 cut appears.
    assert_eq!(batch.num_rows, 1);
    assert_eq!(batch.row_lower[0], 7.0);
}

// -----------------------------------------------------------------------
// Forward DCS integration tests (real ActiveSolver)
// -----------------------------------------------------------------------
//
// These exercise the forward DCS branch in `run_forward_stage`: load the
// cut-free base, patch the pinned incoming state, solve the cut pool lazily,
// and extract the primal. The LP shapes mirror the backward DCS fixture
// (a coupling row ties outgoing storage col0 to the pinned incoming storage
// col2, and cuts constrain theta against col0), so the primal/objective are
// determinate at the pinned state.
mod dcs_forward {
    use cobre_solver::{ActiveSolver, SolverInterface, StageTemplate};

    use super::super::{StageKey, run_forward_stage};
    use crate::context::{StageContext, TrainingContext};
    use crate::cut::FutureCostFunction;
    use crate::cut_selection::CutMetadata;
    use crate::dcs::DcsParams;
    use crate::horizon_mode::HorizonMode;
    use crate::indexer::StageIndexer;
    use crate::inflow_method::InflowNonNegativityMethod;
    use crate::lp_builder::{COST_SCALE_FACTOR, PatchBuffer};
    use crate::trajectory::TrajectoryRecord;
    use crate::workspace::{BasisStore, SolverWorkspace, WorkspaceSizing};

    const X_HAT: f64 = 2.0;

    /// Cut-free base template for `StageIndexer::new(1, 0)`:
    /// cols `[storage_out=0, z_inflow=1, storage_in=2, theta=3]`, one coupling
    /// row `storage_out - storage_in = 0`. Minimise theta. The incoming
    /// state (col 2) is pinned to `x_hat` by the patch; the coupling row ties
    /// col0 to it; cuts constrain theta against col0.
    fn fwd_core_template() -> StageTemplate {
        StageTemplate {
            num_cols: 4,
            num_rows: 1,
            num_nz: 2,
            col_starts: vec![0_i32, 1, 1, 2, 2],
            row_indices: vec![0_i32, 0],
            values: vec![1.0, -1.0],
            col_lower: vec![0.0, 0.0, 0.0, -1.0e6],
            col_upper: vec![f64::INFINITY, f64::INFINITY, f64::INFINITY, 1.0e6],
            objective: vec![0.0, 0.0, 0.0, 1.0],
            row_lower: vec![0.0],
            row_upper: vec![0.0],
            n_state: 1,
            n_transfer: 0,
            n_dual_relevant: 1,
            n_hydro: 1,
            max_par_order: 0,
            col_scale: Vec::new(),
            row_scale: Vec::new(),
        }
    }

    /// All-cuts baked template: the cut-free base plus the three pool cuts
    /// baked as structural rows (rows 1..4), in pool slot order:
    ///   slot 0: -0*col0 + theta >= 1 ; slot 1: -2*col0 + theta >= 0 ;
    ///   slot 2: -0*col0 + theta >= 3.
    /// `num_rows = 4 = template_num_rows(1) + 3 cuts`.
    fn fwd_all_cuts_baked() -> StageTemplate {
        // CSC by column. col0 entries: coupling (row0,+1), slot1 cut (row2,-2).
        // col3 (theta): rows 1,2,3 each +1. col2: coupling (row0,-1).
        StageTemplate {
            num_cols: 4,
            num_rows: 4,
            num_nz: 6,
            // col0: rows [0,2] vals [1,-2]; col1: none; col2: row[0] val[-1];
            // col3: rows [1,2,3] vals [1,1,1].
            col_starts: vec![0_i32, 2, 2, 3, 6],
            row_indices: vec![0_i32, 2, 0, 1, 2, 3],
            values: vec![1.0, -2.0, -1.0, 1.0, 1.0, 1.0],
            col_lower: vec![0.0, 0.0, 0.0, -1.0e6],
            col_upper: vec![f64::INFINITY, f64::INFINITY, f64::INFINITY, 1.0e6],
            objective: vec![0.0, 0.0, 0.0, 1.0],
            // row0 coupling (=0); rows1..3 cuts (>= intercept).
            row_lower: vec![0.0, 1.0, 0.0, 3.0],
            row_upper: vec![0.0, f64::INFINITY, f64::INFINITY, f64::INFINITY],
            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(),
        }
    }

    /// Baked template carrying a single DOMINATING spurious cut
    /// (`-5*col0 + theta >= 0`, floor 10 at `x_hat = 2`, NOT in the pool),
    /// used to prove the DCS path loads the cut-free base and ignores this
    /// row.
    fn fwd_baked_dominating_cut() -> StageTemplate {
        StageTemplate {
            num_cols: 4,
            num_rows: 2,
            num_nz: 4,
            col_starts: vec![0_i32, 2, 2, 3, 4],
            row_indices: vec![0_i32, 1, 0, 1],
            values: vec![1.0, -5.0, -1.0, 1.0],
            col_lower: vec![0.0, 0.0, 0.0, -1.0e6],
            col_upper: vec![f64::INFINITY, f64::INFINITY, f64::INFINITY, 1.0e6],
            objective: vec![0.0, 0.0, 0.0, 1.0],
            row_lower: vec![0.0, 0.0],
            row_upper: vec![0.0, f64::INFINITY],
            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(),
        }
    }

    /// Pool of three cuts on the incoming-storage state, seeded so the DCS
    /// initial set omits the binding slot 1 (stale `last_active_iter`).
    fn fwd_pool() -> FutureCostFunction {
        let mut fcf = FutureCostFunction::new(1, 1, 8, 10, &[0]);
        fcf.add_cut(0, 0, 0, 1.0, &[0.0]);
        fcf.add_cut(0, 0, 1, 0.0, &[2.0]); // binding: floor 2*x_hat = 4
        fcf.add_cut(0, 0, 2, 3.0, &[0.0]);
        let meta = |generated: u64, last: u64| CutMetadata {
            iteration_generated: generated,
            forward_pass_index: 0,
            active_count: 0,
            last_active_iter: last,
        };
        fcf.pools[0].metadata[0] = meta(1, 5);
        fcf.pools[0].metadata[1] = meta(1, 1); // stale → outside k2=2 window at iter 5
        fcf.pools[0].metadata[2] = meta(1, 5);
        fcf
    }

    fn fwd_active_workspace() -> SolverWorkspace<ActiveSolver> {
        let sizing = WorkspaceSizing {
            hydro_count: 1,
            max_par_order: 0,
            n_load_buses: 0,
            max_blocks: 0,
            downstream_par_order: 0,
            max_openings: 1,
            initial_pool_capacity: 16,
            n_state: 1,
            max_local_fwd: 1,
            total_forward_passes: 1,
            noise_dim: 1,
            n_anticipated: 0,
            k_max: 0,
        };
        let solver = ActiveSolver::new().expect("ActiveSolver::new()");
        SolverWorkspace::new(0, 0, solver, PatchBuffer::new(1, 0, 0, 0, 0, 0), 1, sizing)
    }

    fn dcs_params(start_iteration: u64) -> DcsParams {
        DcsParams {
            k1: None,
            k2: 2,
            nadic: 10,
            epsilon_viol: 1e-10,
            start_iteration,
            max_inner_iterations: 50,
        }
    }

    /// Run one forward stage (stage 0 of a 2-stage horizon, so theta is not
    /// terminal-zeroed) with the given `dcs` option and `baked` template,
    /// returning `(stage_cost, advanced_state, scoring_time_seconds)`. The
    /// cut-free base is `ctx.templates[0]`; on the baked path the
    /// caller-equivalent `load_model(baked)` is performed here (mirroring
    /// `run_forward_worker`).
    ///
    /// The production `run_forward_worker` filter is reproduced verbatim:
    /// `dcs.filter(|p| p.is_active(iteration))`. When `iteration <
    /// start_iteration` this collapses `dcs` to `None`, so the baked path is
    /// taken — exactly as the worker does. The returned
    /// `scoring_time_seconds` (read from the workspace's lazy-solve scratch)
    /// is `0.0` iff the lazy path was never entered, giving callers a faithful
    /// witness for which branch ran.
    fn run_one_forward_stage(
        dcs: Option<DcsParams>,
        baked: &StageTemplate,
        iteration: u64,
    ) -> (f64, Vec<f64>, f64) {
        // Mirror run_forward_worker's per-pass gate: DCS is `Some` only when
        // configured AND active at this iteration. This is the production
        // suppression site; reproducing it here (rather than passing `dcs`
        // straight through) is what makes the inactive-iteration assertion a
        // real witness instead of a coincidence.
        let dcs = dcs.filter(|p| p.is_active(iteration));
        let indexer = {
            let mut ix = StageIndexer::new(1, 0);
            ix.finalize_for_test();
            ix
        };
        let core = fwd_core_template();
        let templates = vec![core.clone(), core.clone()];
        let base_rows = vec![0_usize, 0_usize];
        let stochastic = super::make_stochastic_context_1_hydro_3_stages();
        let horizon = HorizonMode::Finite { num_stages: 2 };
        let fcf = fwd_pool();

        let mut ws = fwd_active_workspace();
        ws.current_state.clear();
        ws.current_state.push(X_HAT);
        let mut basis_store = BasisStore::new(1, 2);

        // Discount factor 0 at this stage makes stage_cost = objective * SCALE
        // = theta * SCALE (no other objective term), so the observable is
        // directly sensitive to the converged theta — letting the
        // dominating-baked-cut test distinguish theta=4 (correct) from
        // theta=10 (a wrong baked load).
        let discount_factors = [0.0_f64, 0.0];
        let ctx = StageContext {
            templates: &templates,
            base_rows: &base_rows,
            noise_scale: &[],
            n_hydros: 0,
            n_load_buses: 0,
            load_balance_row_starts: &[],
            load_bus_indices: &[],
            block_counts_per_stage: &[1usize, 1],
            ncs_max_gen: &[],
            ncs_allow_curtailment: &[],
            discount_factors: &discount_factors,
            cumulative_discount_factors: &[],
            stage_lag_transitions: &[],
            noise_group_ids: &[],
            downstream_par_order: 0,
        };
        let training_ctx = TrainingContext {
            horizon: &horizon,
            indexer: &indexer,
            inflow_method: &InflowNonNegativityMethod::None,
            stochastic: &stochastic,
            initial_state: &[],
            inflow_scheme: cobre_core::scenario::SamplingScheme::InSample,
            load_scheme: cobre_core::scenario::SamplingScheme::InSample,
            ncs_scheme: cobre_core::scenario::SamplingScheme::InSample,
            stages: &[],
            historical_library: None,
            external_inflow_library: None,
            external_load_library: None,
            external_ncs_library: None,
            recent_accum_seed: &[],
            recent_weight_seed: 0.0,
            dcs,
            noise_key_diag: None,
        };

        // Baked path: mirror run_forward_worker's per-scenario baked load.
        // DCS path: run_forward_stage loads the cut-free base itself.
        if dcs.is_none() {
            ws.solver.load_model(baked);
        }

        let mut records = vec![TrajectoryRecord {
            primal: Vec::new(),
            dual: Vec::new(),
            stage_cost: 0.0,
            state: Vec::new(),
        }];
        let key = StageKey {
            t: 0,
            m: 0,
            local_m: 0,
            num_stages: 2,
            iteration,
            raw_noise: &[],
            basis_row_capacity: baked.num_rows,
            terminal_has_boundary_cuts: false,
            pool: &fcf.pools[0],
            dcs,
        };
        let mut slices = basis_store.split_workers_mut(1);
        let stage_cost = run_forward_stage(
            &mut ws,
            &mut slices[0],
            &ctx,
            &training_ctx,
            &key,
            &mut records,
        )
        .expect("forward stage solve must succeed");
        // The lazy-solve scratch accumulates scoring wall time only when
        // `lazy_solve_preloaded` runs (the DCS branch). It stays exactly
        // `0.0` on the baked path, so it witnesses which branch executed.
        let scoring_time_seconds = ws.backward_accum.dcs_solve.scoring_time_seconds;
        (stage_cost, records[0].state.clone(), scoring_time_seconds)
    }

    /// AC1: DCS branch (binding cut omitted from the seed) yields the same
    /// stage cost and advanced state as the baked all-cuts path within 1e-9.
    #[test]
    fn forward_dcs_exact_matches_all_cuts() {
        let all_cuts = fwd_all_cuts_baked();
        // iteration 5 >= start_iteration 2 → DCS active; the filter is a pass-through.
        let (baked_cost, baked_state, baked_scoring) = run_one_forward_stage(None, &all_cuts, 5);
        let (dcs_cost, dcs_state, dcs_scoring) =
            run_one_forward_stage(Some(dcs_params(2)), &all_cuts, 5);

        // Baked path never scores; the active DCS path does at least one pass.
        assert_eq!(
            baked_scoring, 0.0,
            "baked path must not enter the lazy solve"
        );
        assert!(
            dcs_scoring > 0.0,
            "active DCS path must enter the lazy solve (scoring time accumulated)"
        );

        assert!(
            (baked_cost - dcs_cost).abs() < 1e-9,
            "stage cost: baked {baked_cost} vs DCS {dcs_cost}"
        );
        // The binding cut floor is 4 at x_hat=2; minimise-theta gives theta=4.
        // With discount 0, stage_cost = objective * SCALE = theta * SCALE, so
        // both paths land on 4 * COST_SCALE_FACTOR.
        assert!((dcs_cost - 4.0 * COST_SCALE_FACTOR).abs() < 1e-3);
        assert_eq!(baked_state.len(), dcs_state.len());
        for (b, d) in baked_state.iter().zip(&dcs_state) {
            assert!((b - d).abs() < 1e-9, "state: baked {b} vs DCS {d}");
        }
        // Sanity: advanced storage state equals the pinned x_hat (coupling
        // row forces storage_out = storage_in = x_hat).
        assert!((dcs_state[0] - X_HAT).abs() < 1e-9);
    }

    /// AC2: a baked template with a DOMINATING embedded cut (floor 10,
    /// gradient 5, NOT in the pool) must NOT change the DCS result — proving
    /// the cut-free `ctx.templates[t]` is loaded, not `params.baked[t]`. If
    /// the DCS path erroneously loaded the dominating baked template, the
    /// advanced state / cost would reflect theta=10, diverging from all-cuts.
    #[test]
    fn forward_dcs_baked_cuts_present_uses_cut_free_core() {
        let all_cuts = fwd_all_cuts_baked();
        let dominating = fwd_baked_dominating_cut();
        let (allcuts_cost, allcuts_state, _) = run_one_forward_stage(None, &all_cuts, 5);
        // DCS path is handed the dominating baked template, but must ignore it
        // and load the cut-free base, recovering the all-cuts result.
        let (dcs_cost, dcs_state, _) = run_one_forward_stage(Some(dcs_params(2)), &dominating, 5);

        assert!(
            (allcuts_cost - dcs_cost).abs() < 1e-9,
            "stage cost: all-cuts {allcuts_cost} vs DCS {dcs_cost} (DCS must \
             ignore the dominating baked cut)"
        );
        assert_eq!(allcuts_state.len(), dcs_state.len());
        for (a, d) in allcuts_state.iter().zip(&dcs_state) {
            assert!(
                (a - d).abs() < 1e-9,
                "state: all-cuts {a} vs DCS {d} (DCS must load the cut-free base)"
            );
        }
    }

    /// AC3: the `run_forward_worker` `is_active` filter actually suppresses
    /// DCS before `start_iteration` and lets it through at/after it.
    ///
    /// This is a real witness, not a coincidence: the suppression is observed
    /// directly via the `scoring_time_seconds` counter, which is `0.0` iff the
    /// lazy solve was never entered. `run_one_forward_stage` reproduces the
    /// production filter (`dcs.filter(|p| p.is_active(iteration))`) verbatim.
    ///
    /// - `start_iteration = 4`, `iteration = 1` (inactive): the filter
    ///   collapses `dcs` to `None`, so the baked path runs — proven by
    ///   `scoring_time_seconds == 0.0` AND a bit-for-bit match with the
    ///   `dcs = None` baked run.
    /// - `iteration = 4` (active): the filter lets DCS through — proven by
    ///   `scoring_time_seconds > 0.0`.
    ///
    /// The `DcsParams::is_active` boundary is also asserted directly so the
    /// filter's flip point is pinned independently of the stage solve.
    #[test]
    fn forward_dcs_inactive_before_start_iteration() {
        // Boundary semantics of the filter predicate itself.
        let params = dcs_params(4);
        assert!(
            !params.is_active(1),
            "iteration 1 < start_iteration 4 must be inactive"
        );
        assert!(
            params.is_active(4),
            "iteration 4 == start_iteration 4 must be active"
        );

        let all_cuts = fwd_all_cuts_baked();

        // Inactive iteration: filter suppresses DCS → baked path. The
        // scoring counter must be exactly 0.0 (lazy solve never entered),
        // and the result must match the dcs=None baked run bit-for-bit.
        let (baked_cost, baked_state, baked_scoring) = run_one_forward_stage(None, &all_cuts, 1);
        let (early_cost, early_state, early_scoring) =
            run_one_forward_stage(Some(dcs_params(4)), &all_cuts, 1);
        assert_eq!(
            baked_scoring, 0.0,
            "dcs=None baked run must not score (sanity)"
        );
        assert_eq!(
            early_scoring, 0.0,
            "iteration 1 < start_iteration 4: filter must suppress DCS, so \
             the baked path runs and no lazy scoring occurs"
        );
        assert_eq!(baked_cost.to_bits(), early_cost.to_bits());
        assert_eq!(baked_state.len(), early_state.len());
        for (b, e) in baked_state.iter().zip(&early_state) {
            assert_eq!(b.to_bits(), e.to_bits());
        }

        // Active iteration (== start_iteration): filter lets DCS through, so
        // the lazy solve runs and accumulates scoring time.
        let (_active_cost, _active_state, active_scoring) =
            run_one_forward_stage(Some(dcs_params(4)), &all_cuts, 4);
        assert!(
            active_scoring > 0.0,
            "iteration 4 >= start_iteration 4: filter must let DCS through, \
             so the lazy solve runs and scoring time accumulates"
        );
    }
}