//! Integration test for warm-start training.
//!
//! Trains a policy, saves it, then warm-starts a new training run from the
//! saved cuts. Verifies that the warm-start FCF has non-zero `warm_start_count`
//! and that training produces correct slot layouts.

#![allow(
    clippy::unwrap_used,
    clippy::expect_used,
    clippy::panic,
    clippy::float_cmp,
    clippy::cast_precision_loss,
    clippy::cast_possible_truncation
)]

use std::path::Path;

use cobre_comm::{CommData, CommError, Communicator, ReduceOp};
use cobre_core::scenario::ScenarioSource;
use cobre_io::output::policy::{read_policy_checkpoint, write_policy_checkpoint};
use cobre_sddp::{
    FutureCostFunction, StudySetup, hydro_models::prepare_hydro_models, setup::prepare_stochastic,
};
use cobre_solver::highs::HighsSolver;

/// Single-rank communicator stub for testing.
struct StubComm;

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

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

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

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

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

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

/// Return the path to the d01-thermal-dispatch example case.
fn d01_case_dir() -> std::path::PathBuf {
    Path::new(env!("CARGO_MANIFEST_DIR"))
        .parent()
        .unwrap()
        .parent()
        .unwrap()
        .join("examples/deterministic/d01-thermal-dispatch")
}

/// Write a policy checkpoint to `policy_dir` from the given setup and training result.
///
/// Convenience helper to avoid duplicating the 20-line metadata + write block
/// in every test that exercises the checkpoint round-trip.
fn write_test_checkpoint(
    policy_dir: &Path,
    setup: &StudySetup,
    result: &cobre_sddp::TrainingResult,
    seed: u64,
) {
    use cobre_sddp::policy_export::{
        build_active_indices, build_stage_basis_records, build_stage_cut_records,
        build_stage_cuts_payloads, convert_basis_cache,
    };
    let fcf = setup.fcf();
    let stage_records = build_stage_cut_records(fcf);
    let stage_active_indices = build_active_indices(&stage_records);
    let stage_cuts = build_stage_cuts_payloads(fcf, &stage_records, &stage_active_indices);
    let (basis_col, basis_row) = convert_basis_cache(result);
    let stage_bases = build_stage_basis_records(fcf, result, &basis_col, &basis_row);
    let metadata = cobre_io::PolicyCheckpointMetadata {
        cobre_version: env!("CARGO_PKG_VERSION").to_string(),
        created_at: "2026-03-29T00:00:00Z".to_string(),
        completed_iterations: result.iterations as u32,
        final_lower_bound: result.final_lb,
        best_upper_bound: Some(result.final_ub),
        state_dimension: fcf.state_dimension as u32,
        num_stages: fcf.pools.len() as u32,
        max_iterations: setup.max_iterations() as u32,
        forward_passes: setup.forward_passes(),
        warm_start_cuts: 0,
        rng_seed: seed,
        total_visited_states: 0,
    };
    write_policy_checkpoint(policy_dir, &stage_cuts, &stage_bases, &metadata, &[])
        .expect("write checkpoint");
}

/// Build a `StudySetup` for the given case directory and config, using seed 42.
fn build_setup(case_dir: &Path, config: &cobre_io::Config) -> StudySetup {
    let system = cobre_io::load_case(case_dir).expect("load_case");
    let prep = prepare_stochastic(system, case_dir, config, 42, &ScenarioSource::default())
        .expect("prepare_stochastic");
    let hydro_models = prepare_hydro_models(&prep.system, case_dir).expect("prepare_hydro_models");
    StudySetup::new(&prep.system, config, prep.stochastic, hydro_models).expect("StudySetup::new")
}

/// Train for 5 iterations, save a checkpoint, resume from iteration 5 up to 10,
/// verify the resumed run reports 10 total iterations and a non-worse lower bound.
#[test]
fn resume_training_from_checkpoint() {
    let case_dir = d01_case_dir();
    let config_path = case_dir.join("config.json");
    let config_full = cobre_io::parse_config(&config_path).expect("config must parse");

    // Build a 5-iteration config for phase 1 by overriding the stopping rules.
    let mut config_phase1 = config_full.clone();
    config_phase1.training.stopping_rules =
        Some(vec![cobre_io::config::StoppingRuleConfig::IterationLimit {
            limit: 5,
        }]);

    // Phase 1: Train for exactly 5 iterations.
    let mut setup_phase1 = build_setup(&case_dir, &config_phase1);
    let comm = StubComm;
    let mut solver_phase1 = HighsSolver::new().expect("HighsSolver");
    let outcome_phase1 = setup_phase1
        .train(&mut solver_phase1, &comm, 1, HighsSolver::new, None, None)
        .expect("train phase1");
    assert!(outcome_phase1.error.is_none());
    let result_phase1 = outcome_phase1.result;
    assert_eq!(
        result_phase1.iterations, 5,
        "phase 1 must complete exactly 5 iterations"
    );
    let lb_phase1 = result_phase1.final_lb;

    // Save the checkpoint from phase 1.
    let tmpdir = tempfile::tempdir().expect("tempdir");
    let policy_dir = tmpdir.path().join("policy");
    write_test_checkpoint(&policy_dir, &setup_phase1, &result_phase1, 42);

    // Phase 2: Resume from the checkpoint using the full 10-iteration config.
    let checkpoint = read_policy_checkpoint(&policy_dir).expect("read checkpoint");
    let mut setup_phase2 = build_setup(&case_dir, &config_full);

    // Replace the FCF with a warm-start version loaded from the checkpoint, then
    // set the start iteration so the training loop begins at iteration 6.
    let warm_fcf = FutureCostFunction::new_with_warm_start(
        &checkpoint.stage_cuts,
        setup_phase2.forward_passes(),
        setup_phase2.max_iterations().saturating_add(1),
    )
    .expect("warm-start FCF");
    setup_phase2.replace_fcf(warm_fcf);
    setup_phase2.set_start_iteration(u64::from(checkpoint.metadata.completed_iterations));

    let mut solver_phase2 = HighsSolver::new().expect("HighsSolver");
    let outcome_phase2 = setup_phase2
        .train(&mut solver_phase2, &comm, 1, HighsSolver::new, None, None)
        .expect("train phase2");
    assert!(outcome_phase2.error.is_none());
    let result_phase2 = outcome_phase2.result;

    // The resumed run must report the absolute iteration count (10), not the delta (5).
    assert_eq!(
        result_phase2.iterations, 10,
        "resumed run must report 10 total iterations (not 5 delta)"
    );

    // The resumed run starts with all cuts from phase 1, so its final LB must be
    // at least as good as phase 1's final LB.
    assert!(
        result_phase2.final_lb >= lb_phase1 - 1e-6,
        "resumed LB ({}) must be >= phase-1 LB ({})",
        result_phase2.final_lb,
        lb_phase1
    );
}

/// Train fresh, save policy, warm-start train, verify improvement and cut counts.
#[test]
fn warm_start_training_preserves_cuts_and_trains_further() {
    let case_dir = d01_case_dir();
    let config_path = case_dir.join("config.json");
    let config = cobre_io::parse_config(&config_path).expect("config must parse");

    // Phase 1: Fresh training.
    let mut setup_fresh = build_setup(&case_dir, &config);
    let comm = StubComm;
    let mut solver = HighsSolver::new().expect("HighsSolver");
    let fresh_outcome = setup_fresh
        .train(&mut solver, &comm, 1, HighsSolver::new, None, None)
        .expect("train");
    assert!(fresh_outcome.error.is_none());
    let fresh_result = fresh_outcome.result;
    let fresh_lb = fresh_result.final_lb;
    let fresh_active = setup_fresh.fcf().total_active_cuts();

    // Write policy checkpoint.
    let tmpdir = tempfile::tempdir().expect("tempdir");
    let policy_dir = tmpdir.path().join("policy");
    write_test_checkpoint(&policy_dir, &setup_fresh, &fresh_result, 42);

    // Phase 2: Warm-start training.
    let checkpoint = read_policy_checkpoint(&policy_dir).expect("read checkpoint");
    let mut setup_warm = build_setup(&case_dir, &config);

    // Build warm-start FCF and replace the fresh one.
    // Use max_iterations + 1 for capacity, matching the original FCF constructor
    // in from_broadcast_params (which does saturating_add(1)).
    let warm_fcf = FutureCostFunction::new_with_warm_start(
        &checkpoint.stage_cuts,
        setup_warm.forward_passes(),
        setup_warm.max_iterations().saturating_add(1),
    )
    .expect("warm-start FCF");
    setup_warm.replace_fcf(warm_fcf);

    // Verify warm-start state before training.
    let warm_start_count = setup_warm.fcf().pools[0].warm_start_count;
    assert!(warm_start_count > 0, "warm_start_count should be > 0");
    assert_eq!(
        setup_warm.fcf().total_active_cuts(),
        fresh_active,
        "warm-start FCF should have same active cuts as fresh training"
    );

    // Train warm-start.
    let mut solver_warm = HighsSolver::new().expect("HighsSolver");
    let warm_outcome = setup_warm
        .train(&mut solver_warm, &comm, 1, HighsSolver::new, None, None)
        .expect("warm-start train");
    assert!(warm_outcome.error.is_none());
    let warm_result = warm_outcome.result;

    // Verify: warm-start final LB should be >= fresh final LB
    // (warm-start starts with all cuts from fresh, so it can only improve).
    assert!(
        warm_result.final_lb >= fresh_lb - 1e-6,
        "warm-start LB ({}) should be >= fresh LB ({})",
        warm_result.final_lb,
        fresh_lb
    );

    // Verify total cuts = warm_start_count + new training cuts.
    let total_active_after = setup_warm.fcf().total_active_cuts();
    assert!(
        total_active_after > fresh_active,
        "warm-start training should produce more total cuts"
    );
}
