cobre-io 0.15.0

//! Parsers for scenario data files in the `scenarios/` subdirectory.
//!
//! Scenario files carry the stochastic inputs used by multi-stage optimization
//! solvers. All files are optional; when absent the `load_*` wrapper returns an
//! empty `Vec` (or default type).
//!
//! Cross-reference validation (entity-ID existence) and dimensional/semantic
//! constraints (lag contiguity, AR coefficient count, block count matching) are
//! deferred to Layer 3/5.

pub mod annual_component;
pub mod ar_coefficients;
pub mod assembly;
pub mod correlation;
pub mod estimation;
pub mod external;
pub mod inflow_history;
pub mod inflow_stats;
pub mod load_factors;
pub mod load_stats;
pub mod noise_openings;
pub mod non_controllable_factors;
pub mod non_controllable_stats;
pub mod residual_derivation;

pub use annual_component::{InflowAnnualComponentRow, parse_inflow_annual_component};
pub use ar_coefficients::{InflowArCoefficientRow, parse_inflow_ar_coefficients};
pub use assembly::{assemble_inflow_models, assemble_load_models};
pub use correlation::parse_correlation;
pub use estimation::{EstimationError, EstimationPath, estimate_from_history};
pub use external::{
    ExternalLoadRow, ExternalNcsRow, ExternalScenarioRow, parse_external_inflow_scenarios,
    parse_external_load_scenarios, parse_external_ncs_scenarios,
};
pub use inflow_history::{InflowHistoryRow, parse_inflow_history};
pub use inflow_stats::{InflowSeasonalStatsRow, parse_inflow_seasonal_stats};
pub use load_factors::{BlockFactor, LoadFactorEntry, parse_load_factors};
pub use load_stats::{LoadSeasonalStatsRow, parse_load_seasonal_stats};
pub use noise_openings::{
    NoiseOpeningRow, assemble_opening_tree, parse_noise_openings, validate_noise_openings,
};
pub use non_controllable_factors::{NcsFactorEntry, parse_non_controllable_factors};
pub use non_controllable_stats::parse_ncs_stats;
pub use residual_derivation::{populate_derived_residual_ratios, resolve_stage_seasons};

use cobre_core::scenario::{CorrelationModel, InflowModel, LoadModel, NcsModel};

use crate::LoadError;
use crate::stages::parse_stages;
use crate::validation::structural::FileManifest;
use std::path::Path;

/// Load `scenarios/inflow_seasonal_stats.parquet`, returning an empty `Vec` when absent.
///
/// # Errors
///
/// Propagates [`LoadError`] from the parser when `path` is `Some`.
pub fn load_inflow_seasonal_stats(
    path: Option<&Path>,
) -> Result<Vec<InflowSeasonalStatsRow>, LoadError> {
    match path {
        None => Ok(Vec::new()),
        Some(p) => parse_inflow_seasonal_stats(p),
    }
}

/// Load `scenarios/inflow_ar_coefficients.parquet`, returning an empty `Vec` when absent.
///
/// # Errors
///
/// Propagates [`LoadError`] from the parser when `path` is `Some`.
pub fn load_inflow_ar_coefficients(
    path: Option<&Path>,
) -> Result<Vec<InflowArCoefficientRow>, LoadError> {
    match path {
        None => Ok(Vec::new()),
        Some(p) => parse_inflow_ar_coefficients(p),
    }
}

/// Load `scenarios/inflow_annual_component.parquet`, returning an empty `Vec` when absent.
///
/// # Errors
///
/// Propagates [`LoadError`] from the parser when `path` is `Some`.
pub fn load_inflow_annual_component(
    path: Option<&Path>,
) -> Result<Vec<InflowAnnualComponentRow>, LoadError> {
    match path {
        None => Ok(Vec::new()),
        Some(p) => parse_inflow_annual_component(p),
    }
}

/// Load `scenarios/inflow_history.parquet`, returning an empty `Vec` when absent.
///
/// # Errors
///
/// Propagates [`LoadError`] from the parser when `path` is `Some`.
pub fn load_inflow_history(path: Option<&Path>) -> Result<Vec<InflowHistoryRow>, LoadError> {
    match path {
        None => Ok(Vec::new()),
        Some(p) => parse_inflow_history(p),
    }
}

/// Load `scenarios/load_seasonal_stats.parquet`, returning an empty `Vec` when absent.
///
/// # Errors
///
/// Propagates [`LoadError`] from the parser when `path` is `Some`.
pub fn load_load_seasonal_stats(
    path: Option<&Path>,
) -> Result<Vec<LoadSeasonalStatsRow>, LoadError> {
    match path {
        None => Ok(Vec::new()),
        Some(p) => parse_load_seasonal_stats(p),
    }
}

/// Load `scenarios/load_factors.json`, returning an empty `Vec` when absent.
///
/// # Errors
///
/// Propagates [`LoadError`] from the parser when `path` is `Some`.
pub fn load_load_factors(path: Option<&Path>) -> Result<Vec<LoadFactorEntry>, LoadError> {
    match path {
        None => Ok(Vec::new()),
        Some(p) => parse_load_factors(p),
    }
}

/// Load `scenarios/non_controllable_factors.json`, returning an empty `Vec` when absent.
///
/// # Errors
///
/// Propagates [`LoadError`] from the parser when `path` is `Some`.
pub fn load_non_controllable_factors(
    path: Option<&Path>,
) -> Result<Vec<NcsFactorEntry>, LoadError> {
    match path {
        None => Ok(Vec::new()),
        Some(p) => parse_non_controllable_factors(p),
    }
}

/// Load `scenarios/non_controllable_stats.parquet`, returning an empty `Vec` when absent.
///
/// # Errors
///
/// Propagates [`LoadError`] from the parser when `path` is `Some`.
pub fn load_ncs_stats(path: Option<&Path>) -> Result<Vec<NcsModel>, LoadError> {
    match path {
        None => Ok(Vec::new()),
        Some(p) => parse_ncs_stats(p),
    }
}

/// Load `scenarios/correlation.json`, returning a default model when absent.
///
/// Unlike the Parquet loaders, this returns `Ok(CorrelationModel::default())` for
/// `None` rather than an empty collection, since the target type is a structured model.
///
/// # Errors
///
/// Propagates [`LoadError`] from the parser when `path` is `Some`.
pub fn load_correlation(path: Option<&Path>) -> Result<CorrelationModel, LoadError> {
    match path {
        None => Ok(CorrelationModel::default()),
        Some(p) => parse_correlation(p),
    }
}

/// Load `scenarios/external_inflow_scenarios.parquet`, returning an empty `Vec` when absent.
///
/// # Errors
///
/// Propagates [`LoadError`] from the parser when `path` is `Some`.
pub fn load_external_inflow_scenarios(
    path: Option<&Path>,
) -> Result<Vec<ExternalScenarioRow>, LoadError> {
    match path {
        None => Ok(Vec::new()),
        Some(p) => parse_external_inflow_scenarios(p),
    }
}

/// Load `scenarios/external_load_scenarios.parquet`, returning an empty `Vec` when absent.
///
/// # Errors
///
/// Propagates [`LoadError`] from the parser when `path` is `Some`.
pub fn load_external_load_scenarios(
    path: Option<&Path>,
) -> Result<Vec<ExternalLoadRow>, LoadError> {
    match path {
        None => Ok(Vec::new()),
        Some(p) => parse_external_load_scenarios(p),
    }
}

/// Load `scenarios/external_ncs_scenarios.parquet`, returning an empty `Vec` when absent.
///
/// # Errors
///
/// Propagates [`LoadError`] from the parser when `path` is `Some`.
pub fn load_external_ncs_scenarios(path: Option<&Path>) -> Result<Vec<ExternalNcsRow>, LoadError> {
    match path {
        None => Ok(Vec::new()),
        Some(p) => parse_external_ncs_scenarios(p),
    }
}

/// Load `scenarios/noise_openings.parquet`, returning an empty `Vec` when absent.
///
/// # Errors
///
/// Propagates [`LoadError`] from the parser when `path` is `Some`.
pub fn load_noise_openings(path: Option<&Path>) -> Result<Vec<NoiseOpeningRow>, LoadError> {
    match path {
        None => Ok(Vec::new()),
        Some(p) => parse_noise_openings(p),
    }
}

// ── ScenarioData ──────────────────────────────────────────────────────────────

/// All assembled scenario pipeline data for one case directory.
///
/// Produced by [`load_scenarios`] after loading and assembling every scenario
/// file present in the manifest. Each field maps directly to the corresponding
/// field in [`cobre_core::System`].
///
/// Optional collections are empty `Vec`s when the corresponding file is absent
/// from the manifest. [`correlation`] is [`CorrelationModel::default()`] when
/// `scenarios/correlation.json` is absent.
///
/// # Examples
///
/// ```
/// use cobre_io::ScenarioData;
/// use cobre_core::scenario::CorrelationModel;
///
/// // Default-like state produced when all manifest flags are false.
/// let data = ScenarioData {
///     inflow_models: vec![],
///     load_models: vec![],
///     ncs_models: vec![],
///     correlation: CorrelationModel::default(),
///     inflow_history: vec![],
///     external_scenarios: vec![],
///     external_load_scenarios: vec![],
///     external_ncs_scenarios: vec![],
///     load_factors: vec![],
/// };
/// assert!(data.inflow_models.is_empty());
/// assert!(data.correlation.profiles.is_empty());
/// ```
#[derive(Debug, Clone)]
pub struct ScenarioData {
    /// PAR(p) inflow models, sorted by `(hydro_id, stage_id)`.
    pub inflow_models: Vec<InflowModel>,
    /// Load seasonal statistics, sorted by `(bus_id, stage_id)`.
    pub load_models: Vec<LoadModel>,
    /// NCS availability noise models, sorted by `(ncs_id, stage_id)`.
    pub ncs_models: Vec<NcsModel>,
    /// Correlation model (profiles + schedule).
    pub correlation: CorrelationModel,
    /// Inflow history rows, sorted by `(hydro_id, start_date)`.
    pub inflow_history: Vec<InflowHistoryRow>,
    /// External inflow scenario rows, sorted by `(stage_id, scenario_id, hydro_id)`.
    pub external_scenarios: Vec<ExternalScenarioRow>,
    /// External load scenario rows, sorted by `(stage_id, scenario_id, bus_id)`.
    pub external_load_scenarios: Vec<ExternalLoadRow>,
    /// External NCS scenario rows, sorted by `(stage_id, scenario_id, ncs_id)`.
    pub external_ncs_scenarios: Vec<ExternalNcsRow>,
    /// Load factor entries, sorted by `(bus_id, stage_id)`.
    pub load_factors: Vec<LoadFactorEntry>,
}

// ── load_scenarios ────────────────────────────────────────────────────────────

/// Load every scenario file present in `manifest` and assemble [`ScenarioData`].
///
/// File paths are constructed as `case_root.join("scenarios/<filename>")`.
///
/// `inflow_models[].residual_std_ratio` is closure-derived, never read from a
/// file (see [`populate_derived_residual_ratios`]). When `manifest.stages_json`
/// is set, this function parses `case_root.join("stages.json")` for stage/season
/// context and runs the derivation over the assembled models, mirroring the
/// four production call sites (`pipeline.rs`, `scenarios::estimation`). When
/// `manifest.stages_json` is unset — this function has no other source of
/// stage/season context — derivation is skipped and every model keeps
/// [`assemble_inflow_models`]'s placeholder `residual_std_ratio = 1.0`
/// unresolved; a caller relying on the derived value in that case must call
/// [`populate_derived_residual_ratios`] itself once stage data is available.
///
/// # Errors
///
/// | Condition                                                | Error variant              |
/// |-----------------------------------------------------------|----------------------------|
/// | Any file read or parse failure                           | Propagated from parser     |
/// | AR coefficient rows without matching stats row           | [`LoadError::SchemaError`] |
/// | AR coefficient rows exist for unknown (hydro, stage)     | [`LoadError::SchemaError`] |
/// | `stages.json` present but malformed                      | Propagated from `parse_stages` |
/// | An order-bearing model's stage has no resolvable season (`stages.json` present) | [`LoadError::ConstraintError`] |
/// | The `residual_std_ratio` closure is singular for a hydro (`stages.json` present) | [`LoadError::ConstraintError`] |
///
/// # Examples
///
/// ```no_run
/// use std::path::Path;
/// use cobre_io::validation::structural::FileManifest;
/// use cobre_io::scenarios::load_scenarios;
///
/// let manifest = FileManifest::default(); // all flags false
/// let data = load_scenarios(Path::new("/case"), &manifest)
///     .expect("empty manifest always succeeds");
/// assert!(data.inflow_models.is_empty());
/// assert!(data.correlation.profiles.is_empty());
/// ```
pub fn load_scenarios(
    case_root: &Path,
    manifest: &FileManifest,
) -> Result<ScenarioData, LoadError> {
    let scenarios_dir = case_root.join("scenarios");

    let raw_stats = load_inflow_seasonal_stats(
        manifest
            .scenarios_inflow_seasonal_stats_parquet
            .then(|| scenarios_dir.join("inflow_seasonal_stats.parquet"))
            .as_deref(),
    )?;
    let raw_coefficients = load_inflow_ar_coefficients(
        manifest
            .scenarios_inflow_ar_coefficients_parquet
            .then(|| scenarios_dir.join("inflow_ar_coefficients.parquet"))
            .as_deref(),
    )?;
    let raw_annual_components = load_inflow_annual_component(
        manifest
            .scenarios_inflow_annual_component_parquet
            .then(|| scenarios_dir.join("inflow_annual_component.parquet"))
            .as_deref(),
    )?;
    let inflow_history = load_inflow_history(
        manifest
            .scenarios_inflow_history_parquet
            .then(|| scenarios_dir.join("inflow_history.parquet"))
            .as_deref(),
    )?;
    let raw_load_stats = load_load_seasonal_stats(
        manifest
            .scenarios_load_seasonal_stats_parquet
            .then(|| scenarios_dir.join("load_seasonal_stats.parquet"))
            .as_deref(),
    )?;
    let load_factors = load_load_factors(
        manifest
            .scenarios_load_factors_json
            .then(|| scenarios_dir.join("load_factors.json"))
            .as_deref(),
    )?;
    let correlation = load_correlation(
        manifest
            .scenarios_correlation_json
            .then(|| scenarios_dir.join("correlation.json"))
            .as_deref(),
    )?;
    let external_scenarios = load_external_inflow_scenarios(
        manifest
            .scenarios_external_inflow_scenarios_parquet
            .then(|| scenarios_dir.join("external_inflow_scenarios.parquet"))
            .as_deref(),
    )?;
    let external_load_scenarios = load_external_load_scenarios(
        manifest
            .scenarios_external_load_scenarios_parquet
            .then(|| scenarios_dir.join("external_load_scenarios.parquet"))
            .as_deref(),
    )?;
    let external_ncs_scenarios = load_external_ncs_scenarios(
        manifest
            .scenarios_external_ncs_scenarios_parquet
            .then(|| scenarios_dir.join("external_ncs_scenarios.parquet"))
            .as_deref(),
    )?;
    let ncs_models = load_ncs_stats(
        manifest
            .scenarios_non_controllable_stats_parquet
            .then(|| scenarios_dir.join("non_controllable_stats.parquet"))
            .as_deref(),
    )?;

    let mut inflow_models =
        assemble_inflow_models(raw_stats, raw_coefficients, raw_annual_components)?;
    if manifest.stages_json {
        let stages_data = parse_stages(&case_root.join("stages.json"))?;
        let (stage_to_season, n_seasons) = resolve_stage_seasons(
            &stages_data.stages,
            stages_data.policy_graph.season_map.as_ref(),
        );
        populate_derived_residual_ratios(&mut inflow_models, &stage_to_season, n_seasons)?;
    }
    let load_models = assemble_load_models(raw_load_stats);

    Ok(ScenarioData {
        inflow_models,
        load_models,
        ncs_models,
        correlation,
        inflow_history,
        external_scenarios,
        external_load_scenarios,
        external_ncs_scenarios,
        load_factors,
    })
}

// ── Tests ─────────────────────────────────────────────────────────────────────

#[cfg(test)]
#[allow(
    clippy::doc_markdown,
    clippy::unwrap_used,
    clippy::expect_used,
    clippy::panic,
    clippy::too_many_lines
)]
mod tests {
    use super::*;
    use crate::validation::structural::FileManifest;
    use tempfile::TempDir;

    #[test]
    fn test_load_inflow_seasonal_stats_none_returns_empty() {
        let result = load_inflow_seasonal_stats(None).unwrap();
        assert!(result.is_empty());
    }

    #[test]
    fn test_load_inflow_ar_coefficients_none_returns_empty() {
        let result = load_inflow_ar_coefficients(None).unwrap();
        assert!(result.is_empty());
    }

    #[test]
    fn test_load_inflow_history_none_returns_empty() {
        let result = load_inflow_history(None).unwrap();
        assert!(result.is_empty());
    }

    #[test]
    fn test_load_load_seasonal_stats_none_returns_empty() {
        let result = load_load_seasonal_stats(None).unwrap();
        assert!(result.is_empty());
    }

    #[test]
    fn test_load_load_factors_none_returns_empty() {
        let result = load_load_factors(None).unwrap();
        assert!(result.is_empty());
    }

    #[test]
    fn test_load_correlation_none_returns_default() {
        let result = load_correlation(None).unwrap();
        assert!(result.profiles.is_empty());
        assert!(result.schedule.is_empty());
    }

    #[test]
    fn test_load_external_inflow_scenarios_none_returns_empty() {
        let result = load_external_inflow_scenarios(None).unwrap();
        assert!(result.is_empty());
    }

    #[test]
    fn test_load_external_load_scenarios_none_returns_empty() {
        let result = load_external_load_scenarios(None).unwrap();
        assert!(result.is_empty());
    }

    #[test]
    fn test_load_external_ncs_scenarios_none_returns_empty() {
        let result = load_external_ncs_scenarios(None).unwrap();
        assert!(result.is_empty());
    }

    #[test]
    fn test_load_scenarios_all_flags_false_returns_empty() {
        let dir = TempDir::new().unwrap();
        let manifest = FileManifest::default();

        let data =
            load_scenarios(dir.path(), &manifest).expect("empty manifest should always succeed");

        assert!(
            data.inflow_models.is_empty(),
            "inflow_models should be empty when no scenario files present"
        );
        assert!(
            data.load_models.is_empty(),
            "load_models should be empty when no scenario files present"
        );
        assert!(
            data.correlation.profiles.is_empty(),
            "correlation.profiles should be empty (default) when correlation.json absent"
        );
        assert!(
            data.correlation.schedule.is_empty(),
            "correlation.schedule should be empty (default) when correlation.json absent"
        );
        assert!(
            data.inflow_history.is_empty(),
            "inflow_history should be empty when file absent"
        );
        assert!(
            data.external_scenarios.is_empty(),
            "external_scenarios should be empty when file absent"
        );
        assert!(
            data.external_load_scenarios.is_empty(),
            "external_load_scenarios should be empty when file absent"
        );
        assert!(
            data.external_ncs_scenarios.is_empty(),
            "external_ncs_scenarios should be empty when file absent"
        );
        assert!(
            data.load_factors.is_empty(),
            "load_factors should be empty when file absent"
        );
    }

    #[test]
    fn test_load_noise_openings_none_returns_empty() {
        let result = load_noise_openings(None).unwrap();
        assert!(result.is_empty());
    }

    #[test]
    fn test_load_scenarios_no_annual_file_yields_none_field() {
        let dir = TempDir::new().unwrap();
        let manifest = FileManifest::default();

        let data =
            load_scenarios(dir.path(), &manifest).expect("empty manifest should always succeed");

        // No inflow stats means no models: the assertion holds vacuously.
        assert!(
            data.inflow_models.iter().all(|m| m.annual.is_none()),
            "every inflow model should have annual == None when annual component file is absent"
        );
    }

    /// With `stages.json` present (`manifest.stages_json = true`), an
    /// order-bearing model's `residual_std_ratio` must come out closure-derived
    /// (`sqrt(1 - psi^2)` for this uniform-AR(1) fixture, per the closure's
    /// exact order-1 decoupling), not the [`assemble_inflow_models`] placeholder
    /// `1.0` — proving the derivation pass is actually wired into
    /// `load_scenarios`, not merely reachable.
    #[test]
    fn test_load_scenarios_derives_residual_std_ratio_when_stages_present() {
        use crate::output::stochastic::{
            write_inflow_ar_coefficients, write_inflow_seasonal_stats,
        };

        let dir = TempDir::new().unwrap();
        let scenarios_dir = dir.path().join("scenarios");
        std::fs::create_dir_all(&scenarios_dir).unwrap();

        std::fs::write(
            dir.path().join("stages.json"),
            r#"{
                "season_definitions": {
                    "cycle_type": "monthly",
                    "seasons": [{ "id": 0, "month_start": 1, "label": "January" }]
                },
                "policy_graph": {
                    "type": "finite_horizon",
                    "annual_discount_rate": 0.0,
                    "transitions": []
                },
                "stages": [{
                    "id": 0, "start_date": "2024-01-01", "end_date": "2024-02-01",
                    "season_id": 0,
                    "blocks": [{ "id": 0, "name": "FLAT", "hours": 744.0 }],
                    "num_openings": 10
                }]
            }"#,
        )
        .unwrap();

        let psi = 0.5_f64;
        write_inflow_seasonal_stats(
            &scenarios_dir.join("inflow_seasonal_stats.parquet"),
            &[InflowSeasonalStatsRow {
                hydro_id: cobre_core::EntityId::from(1),
                stage_id: 0,
                mean_m3s: 100.0,
                std_m3s: 20.0,
            }],
        )
        .unwrap();
        write_inflow_ar_coefficients(
            &scenarios_dir.join("inflow_ar_coefficients.parquet"),
            &[InflowArCoefficientRow {
                hydro_id: cobre_core::EntityId::from(1),
                stage_id: 0,
                lag: 1,
                coefficient: psi,
            }],
        )
        .unwrap();

        let manifest = FileManifest {
            stages_json: true,
            scenarios_inflow_seasonal_stats_parquet: true,
            scenarios_inflow_ar_coefficients_parquet: true,
            ..FileManifest::default()
        };

        let data = load_scenarios(dir.path(), &manifest).expect("load_scenarios must succeed");

        assert_eq!(data.inflow_models.len(), 1);
        let expected = (1.0 - psi * psi).sqrt();
        let got = data.inflow_models[0].residual_std_ratio;
        assert!(
            (got - expected).abs() < 1e-9,
            "residual_std_ratio must be the closure-derived value {expected}, not the \
             assembly placeholder 1.0; got {got}"
        );
    }
}