muse2/

output.rs

//! The module responsible for writing output data to disk.
use crate::agent::AgentID;
use crate::asset::{Asset, AssetID, AssetRef};
use crate::commodity::CommodityID;
use crate::process::ProcessID;
use crate::region::RegionID;
use crate::simulation::investment::appraisal::AppraisalOutput;
use crate::simulation::optimisation::{FlowMap, Solution};
use crate::simulation::prices::PriceMap;
use crate::time_slice::{TimeSliceID, TimeSliceLevel, TimeSliceSelection};
use crate::units::{Activity, Capacity, Flow, Money, MoneyPerActivity, MoneyPerFlow};
use anyhow::{Context, Result, ensure};
use csv;
use indexmap::IndexMap;
use serde::{Deserialize, Serialize};
use std::fs;
use std::fs::File;
use std::path::{Path, PathBuf};

pub mod metadata;
use metadata::write_metadata;

/// The output file name for commodity flows
const COMMODITY_FLOWS_FILE_NAME: &str = "commodity_flows.csv";

/// The output file name for commodity prices
const COMMODITY_PRICES_FILE_NAME: &str = "commodity_prices.csv";

/// The output file name for assets
const ASSETS_FILE_NAME: &str = "assets.csv";

/// The output file name for asset capacities
const ASSET_CAPACITIES_FILE_NAME: &str = "asset_capacities.csv";

/// Debug output file for asset dispatch
const ACTIVITY_ASSET_DISPATCH: &str = "debug_dispatch_assets.csv";

/// The output file name for commodity balance duals
const COMMODITY_BALANCE_DUALS_FILE_NAME: &str = "debug_commodity_balance_duals.csv";

/// The output file name for unmet demand values
const UNMET_DEMAND_FILE_NAME: &str = "debug_unmet_demand.csv";

/// The output file name for extra solver output values
const SOLVER_VALUES_FILE_NAME: &str = "debug_solver.csv";

/// The output file name for appraisal results
const APPRAISAL_RESULTS_FILE_NAME: &str = "debug_appraisal_results.csv";

/// The output file name for appraisal time slice results
const APPRAISAL_RESULTS_TIME_SLICE_FILE_NAME: &str = "debug_appraisal_results_time_slices.csv";

/// Get the default output directory for the model
pub fn get_output_dir(model_dir: &Path, results_root: PathBuf) -> Result<PathBuf> {
    // Get the model name from the dir path. This ends up being convoluted because we need to check
    // for all possible errors. Ugh.
    let model_dir = model_dir
        .canonicalize() // canonicalise in case the user has specified "."
        .context("Could not resolve path to model")?;

    let model_name = model_dir
        .file_name()
        .context("Model cannot be in root folder")?
        .to_str()
        .context("Invalid chars in model dir name")?;

    // Construct path
    Ok([results_root, model_name.into()].iter().collect())
}

/// Get the default output directory for commodity flow graphs for the model
pub fn get_graphs_dir(model_dir: &Path, graph_results_root: PathBuf) -> Result<PathBuf> {
    let model_dir = model_dir
        .canonicalize() // canonicalise in case the user has specified "."
        .context("Could not resolve path to model")?;
    let model_name = model_dir
        .file_name()
        .context("Model cannot be in root folder")?
        .to_str()
        .context("Invalid chars in model dir name")?;
    Ok([graph_results_root, model_name.into()].iter().collect())
}

/// Create a new output directory for the model, optionally overwriting existing data
///
/// # Arguments
///
/// * `output_dir` - The output directory to create/overwrite
/// * `allow_overwrite` - Whether to delete and recreate the folder if it is non-empty
///
/// # Returns
///
/// True if the output dir contained existing data that was deleted, false if not, or an error.
pub fn create_output_directory(output_dir: &Path, allow_overwrite: bool) -> Result<bool> {
    // If the folder already exists, then delete it
    let overwrite = if let Ok(mut it) = fs::read_dir(output_dir) {
        if it.next().is_none() {
            // Folder exists and is empty: nothing to do
            return Ok(false);
        }

        ensure!(
            allow_overwrite,
            "Output folder already exists and is not empty. \
            Please delete the folder or pass the --overwrite command-line option."
        );

        fs::remove_dir_all(output_dir).context("Could not delete folder")?;
        true
    } else {
        false
    };

    // Try to create the directory, with parents
    fs::create_dir_all(output_dir)?;

    Ok(overwrite)
}

/// Copy input files to output directory
pub fn copy_input_files(model_dir: &Path, output_dir: &Path, model_name: &str) -> Result<()> {
    // Get the model name from the dir path.
    let mut input_copy_dir = output_dir.to_path_buf();
    input_copy_dir.extend(["input", model_name]);

    fs::create_dir_all(&input_copy_dir).context("Could not create input copy directory")?;

    for entry in fs::read_dir(model_dir)? {
        let entry = entry?;
        let path = entry.path();
        if path.is_file() {
            let file_name = path.file_name().unwrap();
            fs::copy(&path, input_copy_dir.join(file_name))?;
        }
    }
    Ok(())
}

/// Represents a row in the assets output CSV file.
#[derive(Serialize, Deserialize, Debug, PartialEq)]
struct AssetRow {
    asset_id: AssetID,
    process_id: ProcessID,
    region_id: RegionID,
    agent_id: AgentID,
    commission_year: u32,
}

impl AssetRow {
    /// Create a new [`AssetRow`] for the given asset
    fn new(asset: &Asset) -> Self {
        Self {
            asset_id: asset.id().unwrap(),
            process_id: asset.process_id().clone(),
            region_id: asset.region_id().clone(),
            agent_id: asset.agent_id().unwrap().clone(),
            commission_year: asset.commission_year(),
        }
    }
}

/// Represents a row in the asset capacities output CSV file.
#[derive(Serialize, Deserialize, Debug, PartialEq)]
struct AssetCapacityRow {
    milestone_year: u32,
    asset_id: AssetID,
    capacity: Capacity,
    num_tranches: u32,
    mothballed_capacity: Capacity,
    mothballed_tranches: u32,
}

/// Represents the flow-related data in a row of the commodity flows CSV file.
#[derive(Serialize, Deserialize, Debug, PartialEq)]
struct CommodityFlowRow {
    milestone_year: u32,
    asset_id: AssetID,
    commodity_id: CommodityID,
    time_slice: TimeSliceID,
    flow: Flow,
}

/// Represents a row in the commodity prices CSV file
#[derive(Serialize, Deserialize, Debug, PartialEq)]
struct CommodityPriceRow {
    milestone_year: u32,
    commodity_id: CommodityID,
    region_id: RegionID,
    time_slice: TimeSliceID,
    price: MoneyPerFlow,
}

/// Represents the activity in a row of the dispatch CSV file
#[derive(Serialize, Deserialize, Debug, PartialEq)]
struct DispatchRow {
    milestone_year: u32,
    run_description: String,
    asset_id: Option<AssetID>,
    process_id: ProcessID,
    region_id: RegionID,
    time_slice: TimeSliceID,
    activity: Option<Activity>,
    activity_dual: Option<MoneyPerActivity>,
    column_dual: Option<MoneyPerActivity>,
}

/// Represents the commodity balance duals data in a row of the commodity balance duals CSV file
#[derive(Serialize, Deserialize, Debug, PartialEq)]
struct CommodityBalanceDualsRow {
    milestone_year: u32,
    run_description: String,
    commodity_id: CommodityID,
    region_id: RegionID,
    time_slice: TimeSliceID,
    value: MoneyPerFlow,
}

/// Represents the unmet demand data in a row of the unmet demand CSV file
#[derive(Serialize, Deserialize, Debug, PartialEq)]
struct UnmetDemandRow {
    milestone_year: u32,
    run_description: String,
    commodity_id: CommodityID,
    region_id: RegionID,
    time_slice: TimeSliceID,
    value: Flow,
}

/// Represents solver output values
#[derive(Serialize, Deserialize, Debug, PartialEq)]
struct SolverValuesRow {
    milestone_year: u32,
    run_description: String,
    objective_value: Money,
}

/// Represents the appraisal results in a row of the appraisal results CSV file
#[derive(Serialize, Deserialize, Debug, PartialEq)]
struct AppraisalResultsRow {
    milestone_year: u32,
    run_description: String,
    asset_id: Option<AssetID>,
    process_id: ProcessID,
    region_id: RegionID,
    capacity: Capacity,
    metric: Option<f64>,
}

/// Represents the appraisal results in a row of the appraisal results CSV file
#[derive(Serialize, Deserialize, Debug, PartialEq)]
struct AppraisalResultsTimeSliceRow {
    milestone_year: u32,
    run_description: String,
    asset_id: Option<AssetID>,
    process_id: ProcessID,
    region_id: RegionID,
    time_slice: TimeSliceID,
    activity: Activity,
    activity_coefficient: MoneyPerActivity,
    time_slice_level: TimeSliceLevel,
    demand_for_selection: Flow,
    unmet_demand_for_selection: Flow,
}

/// For writing extra debug information about the model
struct DebugDataWriter {
    context: Option<String>,
    unmet_demand_file_path: PathBuf,
    commodity_balance_duals_writer: csv::Writer<File>,
    unmet_demand_writer: Option<csv::Writer<File>>,
    solver_values_writer: csv::Writer<File>,
    appraisal_results_writer: csv::Writer<File>,
    appraisal_results_time_slice_writer: csv::Writer<File>,
    dispatch_asset_writer: csv::Writer<File>,
}

impl DebugDataWriter {
    /// Open CSV files to write debug info to
    ///
    /// # Arguments
    ///
    /// * `output_path` - Folder where files will be saved
    fn create(output_path: &Path) -> Result<Self> {
        let new_writer = |file_name| {
            let file_path = output_path.join(file_name);
            csv::Writer::from_path(file_path)
        };

        Ok(Self {
            context: None,
            unmet_demand_file_path: output_path.join(UNMET_DEMAND_FILE_NAME),
            commodity_balance_duals_writer: new_writer(COMMODITY_BALANCE_DUALS_FILE_NAME)?,
            unmet_demand_writer: None,
            solver_values_writer: new_writer(SOLVER_VALUES_FILE_NAME)?,
            appraisal_results_writer: new_writer(APPRAISAL_RESULTS_FILE_NAME)?,
            appraisal_results_time_slice_writer: new_writer(
                APPRAISAL_RESULTS_TIME_SLICE_FILE_NAME,
            )?,
            dispatch_asset_writer: new_writer(ACTIVITY_ASSET_DISPATCH)?,
        })
    }

    /// Prepend the current context to the run description
    fn with_context(&self, run_description: &str) -> String {
        if let Some(context) = &self.context {
            format!("{context}; {run_description}")
        } else {
            run_description.to_string()
        }
    }

    /// Write debug info about the dispatch optimisation
    fn write_dispatch_debug_info(
        &mut self,
        milestone_year: u32,
        run_description: &str,
        solution: &Solution,
    ) -> Result<()> {
        self.write_dispatch(
            milestone_year,
            run_description,
            solution.iter_activity(),
            solution.iter_activity_duals(),
            solution.iter_column_duals(),
        )?;
        self.write_commodity_balance_duals(
            milestone_year,
            run_description,
            solution.iter_commodity_balance_duals(),
        )?;
        self.write_unmet_demand(
            milestone_year,
            run_description,
            solution.iter_unmet_demand(),
        )?;
        self.write_solver_values(milestone_year, run_description, solution.objective_value)?;
        Ok(())
    }

    // Write activity to file
    fn write_dispatch<'a, I, J, K>(
        &mut self,
        milestone_year: u32,
        run_description: &str,
        iter_activity: I,
        iter_activity_duals: J,
        iter_column_duals: K,
    ) -> Result<()>
    where
        I: Iterator<Item = (&'a AssetRef, &'a TimeSliceID, Activity)>,
        J: Iterator<Item = (&'a AssetRef, &'a TimeSliceID, MoneyPerActivity)>,
        K: Iterator<Item = (&'a AssetRef, &'a TimeSliceID, MoneyPerActivity)>,
    {
        // To account for different order of entries or missing ones, we first compile data in hash map
        type CompiledActivityData = (
            Option<Activity>,
            Option<MoneyPerActivity>,
            Option<MoneyPerActivity>,
        );
        let mut map: IndexMap<(&AssetRef, &TimeSliceID), CompiledActivityData> = IndexMap::new();

        // For the activities
        for (asset, time_slice, activity) in iter_activity {
            map.entry((asset, time_slice)).or_default().0 = Some(activity);
        }
        // The activity duals
        for (asset, time_slice, activity_dual) in iter_activity_duals {
            map.entry((asset, time_slice)).or_default().1 = Some(activity_dual);
        }
        // And the column duals
        for (asset, time_slice, column_dual) in iter_column_duals {
            map.entry((asset, time_slice)).or_default().2 = Some(column_dual);
        }

        for ((asset, time_slice), (activity, activity_dual, column_dual)) in map {
            let row = DispatchRow {
                milestone_year,
                run_description: self.with_context(run_description),
                asset_id: asset.id(),
                process_id: asset.process_id().clone(),
                region_id: asset.region_id().clone(),
                time_slice: time_slice.clone(),
                activity,
                activity_dual,
                column_dual,
            };
            self.dispatch_asset_writer.serialize(row)?;
        }

        Ok(())
    }

    /// Write commodity balance duals to file
    fn write_commodity_balance_duals<'a, I>(
        &mut self,
        milestone_year: u32,
        run_description: &str,
        iter: I,
    ) -> Result<()>
    where
        I: Iterator<Item = (&'a CommodityID, &'a RegionID, &'a TimeSliceID, MoneyPerFlow)>,
    {
        for (commodity_id, region_id, time_slice, value) in iter {
            let row = CommodityBalanceDualsRow {
                milestone_year,
                run_description: self.with_context(run_description),
                commodity_id: commodity_id.clone(),
                region_id: region_id.clone(),
                time_slice: time_slice.clone(),
                value,
            };
            self.commodity_balance_duals_writer.serialize(row)?;
        }

        Ok(())
    }

    /// Write unmet demand values to file
    fn write_unmet_demand<'a, I>(
        &mut self,
        milestone_year: u32,
        run_description: &str,
        iter: I,
    ) -> Result<()>
    where
        I: Iterator<Item = (&'a CommodityID, &'a RegionID, &'a TimeSliceID, Flow)>,
    {
        let mut rows = iter.peekable();

        if rows.peek().is_none() {
            return Ok(());
        }

        // If the unmet demand writer already exist, we panic, as it should not happen
        assert!(
            self.unmet_demand_writer.is_none(),
            "Unmet demand file already exists!"
        );

        let run_description = self.with_context(run_description);
        let writer = self
            .unmet_demand_writer
            .insert(csv::Writer::from_path(&self.unmet_demand_file_path)?);
        for (commodity_id, region_id, time_slice, value) in rows {
            let row = UnmetDemandRow {
                milestone_year,
                run_description: run_description.clone(),
                commodity_id: commodity_id.clone(),
                region_id: region_id.clone(),
                time_slice: time_slice.clone(),
                value,
            };
            writer.serialize(row)?;
        }

        Ok(())
    }

    /// Write additional solver output values to file
    fn write_solver_values(
        &mut self,
        milestone_year: u32,
        run_description: &str,
        objective_value: Money,
    ) -> Result<()> {
        let row = SolverValuesRow {
            milestone_year,
            run_description: self.with_context(run_description),
            objective_value,
        };
        self.solver_values_writer.serialize(row)?;
        self.solver_values_writer.flush()?;

        Ok(())
    }

    /// Write appraisal results to file
    fn write_appraisal_results(
        &mut self,
        milestone_year: u32,
        run_description: &str,
        appraisal_results: &[AppraisalOutput],
    ) -> Result<()> {
        for result in appraisal_results {
            let row = AppraisalResultsRow {
                milestone_year,
                run_description: self.with_context(run_description),
                asset_id: result.asset.id(),
                process_id: result.asset.process_id().clone(),
                region_id: result.asset.region_id().clone(),
                capacity: result.asset.total_capacity(),
                metric: result.metric.as_ref().map(|m| m.value()),
            };
            self.appraisal_results_writer.serialize(row)?;
        }

        Ok(())
    }

    /// Write appraisal results to file
    fn write_appraisal_time_slice_results(
        &mut self,
        milestone_year: u32,
        run_description: &str,
        appraisal_results: &[AppraisalOutput],
        demand: &IndexMap<TimeSliceSelection, Flow>,
        time_slice_level: TimeSliceLevel,
    ) -> Result<()> {
        for result in appraisal_results {
            for (time_slice, activity) in &result.activity {
                let activity_coefficient = result.coefficients.activity_coefficients[time_slice];
                // Map the individual time slice back to its containing selection so we can look
                // up selection-level demand and unmet demand.
                let selection = match time_slice_level {
                    TimeSliceLevel::Annual => TimeSliceSelection::Annual,
                    TimeSliceLevel::Season => TimeSliceSelection::Season(time_slice.season.clone()),
                    TimeSliceLevel::DayNight => TimeSliceSelection::Single(time_slice.clone()),
                };
                let demand = demand[&selection];
                let unmet_demand = result.unmet_demand[&selection];
                let row = AppraisalResultsTimeSliceRow {
                    milestone_year,
                    run_description: self.with_context(run_description),
                    asset_id: result.asset.id(),
                    process_id: result.asset.process_id().clone(),
                    region_id: result.asset.region_id().clone(),
                    time_slice: time_slice.clone(),
                    activity: *activity,
                    activity_coefficient,
                    time_slice_level,
                    demand_for_selection: demand,
                    unmet_demand_for_selection: unmet_demand,
                };
                self.appraisal_results_time_slice_writer.serialize(row)?;
            }
        }

        Ok(())
    }

    /// Flush the underlying streams
    fn flush(&mut self) -> Result<()> {
        if let Some(wrt) = &mut self.unmet_demand_writer {
            wrt.flush()?;
        }
        self.commodity_balance_duals_writer.flush()?;
        self.solver_values_writer.flush()?;
        self.appraisal_results_writer.flush()?;
        self.appraisal_results_time_slice_writer.flush()?;
        self.dispatch_asset_writer.flush()?;

        Ok(())
    }
}

/// An object for writing output data to file
pub struct DataWriter {
    assets: csv::Writer<File>,
    asset_capacities: csv::Writer<File>,
    flows: csv::Writer<File>,
    prices: csv::Writer<File>,
    debug: Option<DebugDataWriter>,
}

impl DataWriter {
    /// Open CSV files to write output data to
    ///
    /// # Arguments
    ///
    /// * `output_path` - Folder where files will be saved
    /// * `model_path` - Path to input model
    /// * `save_debug_info` - Whether to include extra CSV files for debugging model
    pub fn create(output_path: &Path, model_path: &Path, save_debug_info: bool) -> Result<Self> {
        write_metadata(output_path, model_path).context("Failed to save metadata")?;

        let new_writer = |file_name| {
            let file_path = output_path.join(file_name);
            csv::Writer::from_path(file_path)
        };

        let debug_writer = if save_debug_info {
            // Create debug CSV files
            Some(DebugDataWriter::create(output_path)?)
        } else {
            None
        };

        Ok(Self {
            assets: new_writer(ASSETS_FILE_NAME)?,
            asset_capacities: new_writer(ASSET_CAPACITIES_FILE_NAME)?,
            flows: new_writer(COMMODITY_FLOWS_FILE_NAME)?,
            prices: new_writer(COMMODITY_PRICES_FILE_NAME)?,
            debug: debug_writer,
        })
    }

    /// Write debug info about the dispatch optimisation
    pub fn write_dispatch_debug_info(
        &mut self,
        milestone_year: u32,
        run_description: &str,
        solution: &Solution,
    ) -> Result<()> {
        if let Some(wtr) = &mut self.debug {
            wtr.write_dispatch_debug_info(milestone_year, run_description, solution)?;
        }

        Ok(())
    }

    /// Write debug info about the investment appraisal
    pub fn write_appraisal_debug_info(
        &mut self,
        milestone_year: u32,
        run_description: &str,
        appraisal_results: &[AppraisalOutput],
        demand: &IndexMap<TimeSliceSelection, Flow>,
        time_slice_level: TimeSliceLevel,
    ) -> Result<()> {
        if let Some(wtr) = &mut self.debug {
            wtr.write_appraisal_results(milestone_year, run_description, appraisal_results)?;
            wtr.write_appraisal_time_slice_results(
                milestone_year,
                run_description,
                appraisal_results,
                demand,
                time_slice_level,
            )?;
        }

        Ok(())
    }

    /// Append newly commissioned asset definitions to the assets CSV file
    pub fn write_assets(&mut self, assets: &[AssetRef]) -> Result<()> {
        for asset in assets {
            self.assets.serialize(AssetRow::new(asset))?;
        }

        Ok(())
    }

    /// Write asset capacities for the current milestone year to a CSV file
    pub fn write_asset_capacities(
        &mut self,
        milestone_year: u32,
        assets: &[AssetRef],
    ) -> Result<()> {
        for asset in assets {
            let row = AssetCapacityRow {
                milestone_year,
                asset_id: asset.id().unwrap(),
                capacity: asset.total_capacity(),
                num_tranches: asset.capacity().num_tranches(),
                mothballed_capacity: asset.mothballed_capacity(),
                mothballed_tranches: asset.get_num_mothballed_tranches(),
            };
            self.asset_capacities.serialize(row)?;
        }

        Ok(())
    }

    /// Write commodity flows to a CSV file
    pub fn write_flows(&mut self, milestone_year: u32, flow_map: &FlowMap) -> Result<()> {
        for ((asset, commodity_id, time_slice), flow) in flow_map {
            let row = CommodityFlowRow {
                milestone_year,
                asset_id: asset.id().unwrap(),
                commodity_id: commodity_id.clone(),
                time_slice: time_slice.clone(),
                flow: *flow,
            };
            self.flows.serialize(row)?;
        }

        Ok(())
    }

    /// Write commodity prices to a CSV file
    pub fn write_prices(&mut self, milestone_year: u32, prices: &PriceMap) -> Result<()> {
        for (commodity_id, region_id, time_slice, price) in prices.iter() {
            let row = CommodityPriceRow {
                milestone_year,
                commodity_id: commodity_id.clone(),
                region_id: region_id.clone(),
                time_slice: time_slice.clone(),
                price,
            };
            self.prices.serialize(row)?;
        }

        Ok(())
    }

    /// Flush the underlying streams
    pub fn flush(&mut self) -> Result<()> {
        self.assets.flush()?;
        self.asset_capacities.flush()?;
        self.flows.flush()?;
        self.prices.flush()?;
        if let Some(wtr) = &mut self.debug {
            wtr.flush()?;
        }

        Ok(())
    }

    /// Add context to the debug writer
    pub fn set_debug_context(&mut self, context: String) {
        if let Some(wtr) = &mut self.debug {
            wtr.context = Some(context);
        }
    }

    /// Clear context from the debug writer
    pub fn clear_debug_context(&mut self) {
        if let Some(wtr) = &mut self.debug {
            wtr.context = None;
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::asset::AssetPool;
    use crate::fixture::{
        appraisal_output, asset, assets, commodity_id, multi_tranche_asset, region_id, time_slice,
    };
    use crate::simulation::investment::appraisal::AppraisalOutput;
    use crate::time_slice::{TimeSliceID, TimeSliceLevel, TimeSliceSelection};
    use indexmap::indexmap;
    use itertools::{Itertools, assert_equal};
    use rstest::rstest;
    use std::iter;
    use tempfile::tempdir;

    #[rstest]
    fn write_assets(assets: AssetPool) {
        let dir = tempdir().unwrap();

        // Write an asset
        {
            let mut writer = DataWriter::create(dir.path(), dir.path(), false).unwrap();
            writer.write_assets(&assets).unwrap();
            writer.flush().unwrap();
        }

        // Read back and compare
        let asset = assets.iter().next().unwrap();
        let expected = AssetRow::new(asset);
        let records: Vec<AssetRow> = csv::Reader::from_path(dir.path().join(ASSETS_FILE_NAME))
            .unwrap()
            .into_deserialize()
            .try_collect()
            .unwrap();
        assert_equal(records, iter::once(expected));
    }

    #[rstest]
    fn write_asset_capacities(assets: AssetPool) {
        let milestone_year = 2020;
        let dir = tempdir().unwrap();

        // Write asset capacities
        {
            let mut writer = DataWriter::create(dir.path(), dir.path(), false).unwrap();
            writer
                .write_asset_capacities(milestone_year, &assets)
                .unwrap();
            writer.flush().unwrap();
        }

        // Read back and compare
        let asset = assets.iter().next().unwrap();
        let expected = AssetCapacityRow {
            milestone_year,
            asset_id: asset.id().unwrap(),
            capacity: asset.total_capacity(),
            num_tranches: 1,
            mothballed_capacity: Capacity(0.0),
            mothballed_tranches: 0,
        };
        let records: Vec<AssetCapacityRow> =
            csv::Reader::from_path(dir.path().join(ASSET_CAPACITIES_FILE_NAME))
                .unwrap()
                .into_deserialize()
                .try_collect()
                .unwrap();
        assert_equal(records, iter::once(expected));
    }

    #[rstest]
    fn write_asset_capacities_with_mothballed_tranches(multi_tranche_asset: Asset) {
        let milestone_year = 2020;
        let dir = tempdir().unwrap();
        let mut assets = AssetPool::new();
        assets.commission_new(2010, &mut vec![multi_tranche_asset.into()]);
        let asset = assets
            .iter()
            .next()
            .unwrap()
            .clone()
            .with_mothballed_tranches(1, Some(milestone_year));

        {
            let mut writer = DataWriter::create(dir.path(), dir.path(), false).unwrap();
            writer
                .write_asset_capacities(milestone_year, std::slice::from_ref(&asset))
                .unwrap();
            writer.flush().unwrap();
        }

        let expected = AssetCapacityRow {
            milestone_year,
            asset_id: asset.id().unwrap(),
            capacity: Capacity(12.0),
            num_tranches: 3,
            mothballed_capacity: Capacity(4.0),
            mothballed_tranches: 1,
        };
        let records: Vec<AssetCapacityRow> =
            csv::Reader::from_path(dir.path().join(ASSET_CAPACITIES_FILE_NAME))
                .unwrap()
                .into_deserialize()
                .try_collect()
                .unwrap();
        assert_equal(records, iter::once(expected));
    }

    #[rstest]
    fn write_flows(assets: AssetPool, commodity_id: CommodityID, time_slice: TimeSliceID) {
        let milestone_year = 2020;
        let asset = assets.iter().next().unwrap();
        let flow_map = indexmap! {
            (asset.clone(), commodity_id.clone(), time_slice.clone()) => Flow(42.0)
        };

        // Write a flow
        let dir = tempdir().unwrap();
        {
            let mut writer = DataWriter::create(dir.path(), dir.path(), false).unwrap();
            writer.write_flows(milestone_year, &flow_map).unwrap();
            writer.flush().unwrap();
        }

        // Read back and compare
        let expected = CommodityFlowRow {
            milestone_year,
            asset_id: asset.id().unwrap(),
            commodity_id,
            time_slice,
            flow: Flow(42.0),
        };
        let records: Vec<CommodityFlowRow> =
            csv::Reader::from_path(dir.path().join(COMMODITY_FLOWS_FILE_NAME))
                .unwrap()
                .into_deserialize()
                .try_collect()
                .unwrap();
        assert_equal(records, iter::once(expected));
    }

    #[rstest]
    fn write_prices(commodity_id: CommodityID, region_id: RegionID, time_slice: TimeSliceID) {
        let milestone_year = 2020;
        let price = MoneyPerFlow(42.0);
        let mut prices = PriceMap::default();
        prices.insert(&commodity_id, &region_id, &time_slice, price);

        let dir = tempdir().unwrap();

        // Write a price
        {
            let mut writer = DataWriter::create(dir.path(), dir.path(), false).unwrap();
            writer.write_prices(milestone_year, &prices).unwrap();
            writer.flush().unwrap();
        }

        // Read back and compare
        let expected = CommodityPriceRow {
            milestone_year,
            commodity_id,
            region_id,
            time_slice,
            price,
        };
        let records: Vec<CommodityPriceRow> =
            csv::Reader::from_path(dir.path().join(COMMODITY_PRICES_FILE_NAME))
                .unwrap()
                .into_deserialize()
                .try_collect()
                .unwrap();
        assert_equal(records, iter::once(expected));
    }

    #[rstest]
    fn write_commodity_balance_duals(
        commodity_id: CommodityID,
        region_id: RegionID,
        time_slice: TimeSliceID,
    ) {
        let milestone_year = 2020;
        let run_description = "test_run".to_string();
        let value = MoneyPerFlow(0.5);
        let dir = tempdir().unwrap();

        // Write commodity balance dual
        {
            let mut writer = DebugDataWriter::create(dir.path()).unwrap();
            writer
                .write_commodity_balance_duals(
                    milestone_year,
                    &run_description,
                    iter::once((&commodity_id, &region_id, &time_slice, value)),
                )
                .unwrap();
            writer.flush().unwrap();
        }

        // Read back and compare
        let expected = CommodityBalanceDualsRow {
            milestone_year,
            run_description,
            commodity_id,
            region_id,
            time_slice,
            value,
        };
        let records: Vec<CommodityBalanceDualsRow> =
            csv::Reader::from_path(dir.path().join(COMMODITY_BALANCE_DUALS_FILE_NAME))
                .unwrap()
                .into_deserialize()
                .try_collect()
                .unwrap();
        assert_equal(records, iter::once(expected));
    }

    #[rstest]
    fn write_unmet_demand(commodity_id: CommodityID, region_id: RegionID, time_slice: TimeSliceID) {
        let milestone_year = 2020;
        let run_description = "test_run".to_string();
        let value = Flow(0.5);
        let dir = tempdir().unwrap();

        // Write unmet demand
        {
            let mut writer = DebugDataWriter::create(dir.path()).unwrap();
            writer
                .write_unmet_demand(
                    milestone_year,
                    &run_description,
                    iter::once((&commodity_id, &region_id, &time_slice, value)),
                )
                .unwrap();
            writer.flush().unwrap();
        }

        // Read back and compare
        let expected = UnmetDemandRow {
            milestone_year,
            run_description,
            commodity_id,
            region_id,
            time_slice,
            value,
        };
        let records: Vec<UnmetDemandRow> =
            csv::Reader::from_path(dir.path().join(UNMET_DEMAND_FILE_NAME))
                .unwrap()
                .into_deserialize()
                .try_collect()
                .unwrap();
        assert_equal(records, iter::once(expected));
    }

    #[rstest]
    fn write_dispatch(assets: AssetPool, time_slice: TimeSliceID) {
        let milestone_year = 2020;
        let run_description = "test_run".to_string();
        let activity = Activity(100.5);
        let activity_dual = MoneyPerActivity(-1.5);
        let column_dual = MoneyPerActivity(5.0);
        let dir = tempdir().unwrap();
        let asset = assets.iter().next().unwrap();

        // Write activity
        {
            let mut writer = DebugDataWriter::create(dir.path()).unwrap();
            writer
                .write_dispatch(
                    milestone_year,
                    &run_description,
                    iter::once((asset, &time_slice, activity)),
                    iter::once((asset, &time_slice, activity_dual)),
                    iter::once((asset, &time_slice, column_dual)),
                )
                .unwrap();
            writer.flush().unwrap();
        }

        // Read back and compare
        let expected = DispatchRow {
            milestone_year,
            run_description,
            asset_id: asset.id(),
            process_id: asset.process_id().clone(),
            region_id: asset.region_id().clone(),
            time_slice,
            activity: Some(activity),
            activity_dual: Some(activity_dual),
            column_dual: Some(column_dual),
        };
        let records: Vec<DispatchRow> =
            csv::Reader::from_path(dir.path().join(ACTIVITY_ASSET_DISPATCH))
                .unwrap()
                .into_deserialize()
                .try_collect()
                .unwrap();
        assert_equal(records, iter::once(expected));
    }

    #[rstest]
    fn write_dispatch_with_missing_keys(assets: AssetPool, time_slice: TimeSliceID) {
        let milestone_year = 2020;
        let run_description = "test_run".to_string();
        let activity = Activity(100.5);
        let dir = tempdir().unwrap();
        let asset = assets.iter().next().unwrap();

        // Write activity
        {
            let mut writer = DebugDataWriter::create(dir.path()).unwrap();
            writer
                .write_dispatch(
                    milestone_year,
                    &run_description,
                    iter::once((asset, &time_slice, activity)),
                    iter::empty::<(&AssetRef, &TimeSliceID, MoneyPerActivity)>(),
                    iter::empty::<(&AssetRef, &TimeSliceID, MoneyPerActivity)>(),
                )
                .unwrap();
            writer.flush().unwrap();
        }

        // Read back and compare
        let expected = DispatchRow {
            milestone_year,
            run_description,
            asset_id: asset.id(),
            process_id: asset.process_id().clone(),
            region_id: asset.region_id().clone(),
            time_slice,
            activity: Some(activity),
            activity_dual: None,
            column_dual: None,
        };
        let records: Vec<DispatchRow> =
            csv::Reader::from_path(dir.path().join(ACTIVITY_ASSET_DISPATCH))
                .unwrap()
                .into_deserialize()
                .try_collect()
                .unwrap();
        assert_equal(records, iter::once(expected));
    }

    #[rstest]
    fn write_solver_values() {
        let milestone_year = 2020;
        let run_description = "test_run".to_string();
        let objective_value = Money(1234.56);
        let dir = tempdir().unwrap();

        // Write solver values
        {
            let mut writer = DebugDataWriter::create(dir.path()).unwrap();
            writer
                .write_solver_values(milestone_year, &run_description, objective_value)
                .unwrap();
            writer.flush().unwrap();
        }

        // Read back and compare
        let expected = SolverValuesRow {
            milestone_year,
            run_description,
            objective_value,
        };
        let records: Vec<SolverValuesRow> =
            csv::Reader::from_path(dir.path().join(SOLVER_VALUES_FILE_NAME))
                .unwrap()
                .into_deserialize()
                .try_collect()
                .unwrap();
        assert_equal(records, iter::once(expected));
    }

    #[rstest]
    fn write_appraisal_results(asset: Asset, appraisal_output: AppraisalOutput) {
        let milestone_year = 2020;
        let run_description = "test_run".to_string();
        let dir = tempdir().unwrap();

        // Write appraisal results
        {
            let mut writer = DebugDataWriter::create(dir.path()).unwrap();
            writer
                .write_appraisal_results(milestone_year, &run_description, &[appraisal_output])
                .unwrap();
            writer.flush().unwrap();
        }

        // Read back and compare
        let expected = AppraisalResultsRow {
            milestone_year,
            run_description,
            asset_id: None,
            process_id: asset.process_id().clone(),
            region_id: asset.region_id().clone(),
            capacity: Capacity(2.0),
            metric: Some(4.14),
        };
        let records: Vec<AppraisalResultsRow> =
            csv::Reader::from_path(dir.path().join(APPRAISAL_RESULTS_FILE_NAME))
                .unwrap()
                .into_deserialize()
                .try_collect()
                .unwrap();
        assert_equal(records, iter::once(expected));
    }

    #[rstest]
    fn write_appraisal_time_slice_results(
        asset: Asset,
        appraisal_output: AppraisalOutput,
        time_slice: TimeSliceID,
    ) {
        let milestone_year = 2020;
        let run_description = "test_run".to_string();
        let dir = tempdir().unwrap();
        let demand = indexmap! {TimeSliceSelection::Single(time_slice.clone()) => Flow(100.0) };

        // Write appraisal time slice results
        {
            let mut writer = DebugDataWriter::create(dir.path()).unwrap();
            writer
                .write_appraisal_time_slice_results(
                    milestone_year,
                    &run_description,
                    &[appraisal_output],
                    &demand,
                    TimeSliceLevel::DayNight,
                )
                .unwrap();
            writer.flush().unwrap();
        }

        // Read back and compare
        let expected = AppraisalResultsTimeSliceRow {
            milestone_year,
            run_description,
            asset_id: None,
            process_id: asset.process_id().clone(),
            region_id: asset.region_id().clone(),
            time_slice: time_slice.clone(),
            activity: Activity(10.0),
            activity_coefficient: MoneyPerActivity(0.5),
            time_slice_level: TimeSliceLevel::DayNight,
            demand_for_selection: Flow(100.0),
            unmet_demand_for_selection: Flow(5.0),
        };
        let records: Vec<AppraisalResultsTimeSliceRow> =
            csv::Reader::from_path(dir.path().join(APPRAISAL_RESULTS_TIME_SLICE_FILE_NAME))
                .unwrap()
                .into_deserialize()
                .try_collect()
                .unwrap();
        assert_equal(records, iter::once(expected));
    }

    #[test]
    fn create_output_directory_new_directory() {
        let temp_dir = tempdir().unwrap();
        let output_dir = temp_dir.path().join("new_output");

        // Create a new directory should succeed and return false (no overwrite)
        let result = create_output_directory(&output_dir, false).unwrap();
        assert!(!result);
        assert!(output_dir.exists());
        assert!(output_dir.is_dir());
    }

    #[test]
    fn create_output_directory_existing_empty_directory() {
        let temp_dir = tempdir().unwrap();
        let output_dir = temp_dir.path().join("empty_output");

        // Create the directory first
        fs::create_dir(&output_dir).unwrap();

        // Creating again should succeed and return false (no overwrite needed)
        let result = create_output_directory(&output_dir, false).unwrap();
        assert!(!result);
        assert!(output_dir.exists());
        assert!(output_dir.is_dir());
    }

    #[test]
    fn create_output_directory_existing_with_files_no_overwrite() {
        let temp_dir = tempdir().unwrap();
        let output_dir = temp_dir.path().join("output_with_files");

        // Create directory with a file
        fs::create_dir(&output_dir).unwrap();
        fs::write(output_dir.join("existing_file.txt"), "some content").unwrap();

        // Should fail when allow_overwrite is false
        let result = create_output_directory(&output_dir, false);
        assert!(result.is_err());
        assert!(
            result
                .unwrap_err()
                .to_string()
                .contains("Output folder already exists")
        );
    }

    #[test]
    fn create_output_directory_existing_with_files_allow_overwrite() {
        let temp_dir = tempdir().unwrap();
        let output_dir = temp_dir.path().join("output_with_files");

        // Create directory with a file
        fs::create_dir(&output_dir).unwrap();
        let file_path = output_dir.join("existing_file.txt");
        fs::write(&file_path, "some content").unwrap();

        // Should succeed when allow_overwrite is true and return true (overwrite occurred)
        let result = create_output_directory(&output_dir, true).unwrap();
        assert!(result);
        assert!(output_dir.exists());
        assert!(output_dir.is_dir());
        assert!(!file_path.exists()); // File should be gone
    }

    #[test]
    fn create_output_directory_nested_path() {
        let temp_dir = tempdir().unwrap();
        let output_dir = temp_dir.path().join("nested").join("path").join("output");

        // Should create nested directories and return false (no overwrite)
        let result = create_output_directory(&output_dir, false).unwrap();
        assert!(!result);
        assert!(output_dir.exists());
        assert!(output_dir.is_dir());
    }

    #[test]
    fn create_output_directory_existing_subdirs_with_files_allow_overwrite() {
        let temp_dir = tempdir().unwrap();
        let output_dir = temp_dir.path().join("output_with_subdirs");

        // Create directory structure with files
        fs::create_dir_all(output_dir.join("subdir")).unwrap();
        fs::write(output_dir.join("file1.txt"), "content1").unwrap();
        fs::write(output_dir.join("subdir").join("file2.txt"), "content2").unwrap();

        // Should succeed when allow_overwrite is true and return true (overwrite occurred)
        let result = create_output_directory(&output_dir, true).unwrap();
        assert!(result);
        assert!(output_dir.exists());
        assert!(output_dir.is_dir());
        // All previous content should be gone
        assert!(!output_dir.join("file1.txt").exists());
        assert!(!output_dir.join("subdir").exists());
    }
}