muse2/asset/

pool.rs

//! Defines a data structure for representing the current active pool of assets.
use super::{AssetID, AssetRef, AssetState, UserAsset};
use itertools::Itertools;
use log::warn;

/// The active pool of [`super::Asset`]s
#[derive(Default, derive_more::Deref)]
pub struct AssetPool {
    /// The pool of active assets, sorted by ID
    #[deref]
    assets: Vec<AssetRef>,
    /// Next available asset ID number
    next_id: u32,
}

impl AssetPool {
    /// Create a new empty [`AssetPool`]
    pub fn new() -> Self {
        Self::default()
    }

    /// Commission new assets for the specified milestone year from the input data.
    ///
    /// Returns the newly commissioned assets.
    pub fn commission_new(&mut self, year: u32, user_assets: &mut Vec<UserAsset>) -> &[AssetRef] {
        let start = self.assets.len();
        let to_commission = user_assets.extract_if(.., |asset| asset.commission_year <= year);

        for asset in to_commission {
            // Ignore assets that have already been decommissioned
            if asset.max_decommission_year() <= year {
                warn!(
                    "User asset '{}' with commission year {} with maximum decommission year {} \
                    was decommissioned before start of the simulation",
                    asset.process_id(),
                    asset.commission_year,
                    asset.max_decommission_year
                );
                continue;
            }

            self.commission(asset.into());
        }

        &self.assets[start..]
    }

    /// Commission the specified asset
    fn commission(&mut self, mut asset: AssetRef) {
        asset.make_mut().commission(AssetID(self.next_id));
        self.next_id += 1;
        self.assets.push(asset);
    }

    /// Decommission old assets for the specified milestone year
    pub fn decommission_old(&mut self, year: u32) {
        self.assets
            .extract_if(.., |asset| asset.max_decommission_year() <= year)
            .for_each(|asset| {
                asset.decommission("end of life");
            });
    }

    /// Decommission mothballed assets if mothballed long enough
    pub fn decommission_mothballed(&mut self, year: u32, mothball_years: u32) {
        // Empty the Vec and reconstruct it with only the remaining units of the remaining assets
        // after decommissioning. This sadly means we always allocate a new Vec, but modifying the
        // Vec in place leads to uglier code and unnecessary deep clones of assets.
        self.assets = std::mem::take(&mut self.assets)
            .into_iter()
            .filter_map(|asset| asset.with_decommission_mothballed(year, mothball_years))
            .collect();
    }

    /// Mothball the specified assets if they are no longer in the active pool and put them back
    /// again.
    ///
    /// # Arguments
    ///
    /// * `assets` - Assets to possibly mothball
    /// * `year` - Mothball year
    ///
    /// # Panics
    ///
    /// Panics if any of the provided assets was never commissioned.
    pub fn mothball_unretained<I>(&mut self, assets: I, year: u32)
    where
        I: IntoIterator<Item = AssetRef>,
    {
        for old_asset in assets {
            let id = old_asset
                .id()
                .expect("Cannot mothball asset that has not been commissioned");

            // Note that we cannot use a binary search here, as `self.assets` may have become
            // unsorted by new assets added below
            if let Some(new_asset) = self
                .assets
                .iter_mut()
                .find(|asset| asset.id().unwrap() == id)
            {
                // At least some of the asset's units have made it back into the pool. Increase the
                // capacity back to what it was before, with the unselected units set as mothballed.
                let num_mothballed = old_asset
                    .num_units()
                    .checked_sub(new_asset.num_units())
                    .expect("Number of units has increased");
                *new_asset = old_asset.with_mothballed_units(num_mothballed, Some(year));
            } else {
                // None of this asset's units were selected. We mothball _all_ units and return to
                // the pool.
                let num_mothballed = old_asset.num_units();
                self.assets
                    .push(old_asset.with_mothballed_units(num_mothballed, Some(year)));
            }
        }
        self.assets.sort();
    }

    /// Get an asset with the specified ID.
    ///
    /// # Returns
    ///
    /// An [`AssetRef`] if found, else `None`. The asset may not be found if it has already been
    /// decommissioned.
    pub fn get(&self, id: AssetID) -> Option<&AssetRef> {
        // Assets are sorted by ID
        let idx = self
            .assets
            .binary_search_by(|asset| match &asset.state {
                AssetState::Commissioned { id: asset_id, .. } => asset_id.cmp(&id),
                _ => panic!("Active pool should only contain commissioned assets"),
            })
            .ok()?;

        Some(&self.assets[idx])
    }

    /// Return current active pool and clear
    pub fn take(&mut self) -> Vec<AssetRef> {
        std::mem::take(&mut self.assets)
    }

    /// Extend the active pool with Commissioned or Ready assets.
    ///
    /// Returns the newly commissioned assets (those that were in `Ready` state on entry).
    pub fn extend<I>(&mut self, assets: I) -> &[AssetRef]
    where
        I: IntoIterator<Item = AssetRef>,
    {
        let first_new_id = self.next_id;

        // Check all assets are either Commissioned or Ready, and, if the latter,
        // then commission them
        for asset in assets {
            match &asset.state {
                AssetState::Commissioned { .. } => {
                    self.assets.push(asset);
                }
                AssetState::Ready { .. } => {
                    self.commission(asset);
                }
                AssetState::Candidate => panic!(
                    "Cannot extend asset pool with asset in state {}. Only assets in \
                    Commissioned or Ready states are allowed.",
                    asset.state
                ),
            }
        }

        // New assets may not have been sorted, but we need them sorted by ID
        self.assets.sort();

        // Sanity check: all assets should be unique
        debug_assert_eq!(self.assets.iter().unique().count(), self.assets.len());

        // Newly commissioned assets have IDs >= first_new_id. Since assets are sorted by ID,
        // they are at the tail of the slice.
        let new_start = self.assets.partition_point(|a| match &a.state {
            AssetState::Commissioned { id, .. } => id.0 < first_new_id,
            _ => panic!("Active pool should only contain commissioned assets"),
        });
        &self.assets[new_start..]
    }
}

#[cfg(test)]
mod tests {
    use super::super::Asset;
    use super::*;
    use crate::asset::MothballEvent;
    use crate::fixture::{asset, asset_divisible, process, process_parameter_map};
    use crate::process::{Process, ProcessParameter};
    use crate::units::{
        Capacity, Dimensionless, MoneyPerActivity, MoneyPerCapacity, MoneyPerCapacityPerYear,
    };
    use itertools::{Itertools, assert_equal};
    use rstest::{fixture, rstest};
    use std::iter;
    use std::rc::Rc;

    #[fixture]
    fn user_assets(mut process: Process) -> Vec<UserAsset> {
        // Update process parameters (lifetime = 20 years)
        let process_param = ProcessParameter {
            capital_cost: MoneyPerCapacity(5.0),
            fixed_operating_cost: MoneyPerCapacityPerYear(2.0),
            variable_operating_cost: MoneyPerActivity(1.0),
            lifetime: 20,
            discount_rate: Dimensionless(0.9),
        };
        let process_parameter_map = process_parameter_map(process.regions.clone(), process_param);
        process.parameters = process_parameter_map;

        let rc_process = Rc::new(process);
        [2020, 2010]
            .map(|year| {
                UserAsset::new(
                    "agent1".into(),
                    Rc::clone(&rc_process),
                    "GBR".into(),
                    Capacity(1.0),
                    year,
                    None,
                )
                .unwrap()
            })
            .into_iter()
            .collect_vec()
    }

    #[rstest]
    fn asset_pool_new() {
        assert!(AssetPool::new().assets.is_empty());
    }

    #[rstest]
    fn asset_pool_commission_new1(mut user_assets: Vec<UserAsset>) {
        // Asset to be commissioned in this year
        let mut asset_pool = AssetPool::new();
        asset_pool.commission_new(2010, &mut user_assets);
        assert_equal(asset_pool.iter(), iter::once(&asset_pool.assets[0]));
    }

    #[rstest]
    fn asset_pool_commission_new2(mut user_assets: Vec<UserAsset>) {
        // Commission year has passed
        let mut asset_pool = AssetPool::new();
        asset_pool.commission_new(2011, &mut user_assets);
        assert_equal(asset_pool.iter(), iter::once(&asset_pool.assets[0]));
    }

    #[rstest]
    fn asset_pool_commission_new3(mut user_assets: Vec<UserAsset>) {
        // Nothing to commission for this year
        let mut asset_pool = AssetPool::new();
        asset_pool.commission_new(2000, &mut user_assets);
        assert!(asset_pool.iter().next().is_none()); // no active assets
    }

    #[rstest]
    fn asset_pool_commission_new_divisible(asset_divisible: Asset) {
        let commission_year = asset_divisible.commission_year;
        let mut asset_pool = AssetPool::new();
        let mut user_assets = vec![asset_divisible.into()];
        assert!(asset_pool.assets.is_empty());
        asset_pool.commission_new(commission_year, &mut user_assets);
        assert!(user_assets.is_empty());
        assert_eq!(asset_pool.assets.len(), 1);
        assert_eq!(asset_pool.next_id, 1);
    }

    #[rstest]
    fn asset_pool_commission_already_decommissioned(asset: Asset) {
        let year = asset.max_decommission_year();
        let mut asset_pool = AssetPool::new();
        assert!(asset_pool.assets.is_empty());
        asset_pool.commission_new(year, &mut vec![asset.into()]);
        assert!(asset_pool.assets.is_empty());
    }

    #[rstest]
    fn asset_pool_decommission_old(mut user_assets: Vec<UserAsset>) {
        let mut asset_pool = AssetPool::new();
        asset_pool.commission_new(2020, &mut user_assets);
        assert!(user_assets.is_empty());
        assert_eq!(asset_pool.assets.len(), 2);

        // should decommission first asset (lifetime == 5)
        asset_pool.decommission_old(2030);
        assert_eq!(asset_pool.assets.len(), 1);
        assert_eq!(asset_pool.assets[0].commission_year, 2020);

        // nothing to decommission
        asset_pool.decommission_old(2032);
        assert_eq!(asset_pool.assets.len(), 1);
        assert_eq!(asset_pool.assets[0].commission_year, 2020);

        // should decommission second asset
        asset_pool.decommission_old(2040);
        assert!(asset_pool.assets.is_empty());
    }

    #[rstest]
    fn asset_pool_get(mut user_assets: Vec<UserAsset>) {
        let mut asset_pool = AssetPool::new();
        asset_pool.commission_new(2020, &mut user_assets);
        assert_eq!(asset_pool.get(AssetID(0)), Some(&asset_pool.assets[0]));
        assert_eq!(asset_pool.get(AssetID(1)), Some(&asset_pool.assets[1]));
    }

    #[rstest]
    fn asset_pool_extend_empty(mut user_assets: Vec<UserAsset>) {
        // Start with commissioned assets
        let mut asset_pool = AssetPool::new();
        asset_pool.commission_new(2020, &mut user_assets);
        let original_count = asset_pool.assets.len();

        // Extend with empty iterator
        asset_pool.extend(Vec::<AssetRef>::new());

        assert_eq!(asset_pool.assets.len(), original_count);
    }

    #[rstest]
    fn asset_pool_extend_existing_assets(mut user_assets: Vec<UserAsset>) {
        // Start with some commissioned assets
        let mut asset_pool = AssetPool::new();
        asset_pool.commission_new(2020, &mut user_assets);
        assert_eq!(asset_pool.assets.len(), 2);
        let existing_assets = asset_pool.take();

        // Extend with the same assets (should maintain their IDs)
        asset_pool.extend(existing_assets.clone());

        assert_eq!(asset_pool.assets.len(), 2);
        assert_eq!(asset_pool.assets[0].id(), Some(AssetID(0)));
        assert_eq!(asset_pool.assets[1].id(), Some(AssetID(1)));
    }

    #[rstest]
    fn asset_pool_extend_new_assets(mut user_assets: Vec<UserAsset>, process: Process) {
        // Start with some commissioned assets
        let mut asset_pool = AssetPool::new();
        asset_pool.commission_new(2020, &mut user_assets);
        let original_count = asset_pool.assets.len();

        // Create new non-commissioned assets
        let process_rc = Rc::new(process);
        let new_assets = vec![
            Asset::new_ready(
                "agent2".into(),
                Rc::clone(&process_rc),
                "GBR".into(),
                Capacity(1.5),
                2015,
            )
            .unwrap()
            .into(),
            Asset::new_ready(
                "agent3".into(),
                Rc::clone(&process_rc),
                "GBR".into(),
                Capacity(2.5),
                2020,
            )
            .unwrap()
            .into(),
        ];

        asset_pool.extend(new_assets);

        assert_eq!(asset_pool.assets.len(), original_count + 2);
        // New assets should get IDs 2 and 3
        assert_eq!(asset_pool.assets[original_count].id(), Some(AssetID(2)));
        assert_eq!(asset_pool.assets[original_count + 1].id(), Some(AssetID(3)));
        assert_eq!(
            asset_pool.assets[original_count].agent_id(),
            Some(&"agent2".into())
        );
        assert_eq!(
            asset_pool.assets[original_count + 1].agent_id(),
            Some(&"agent3".into())
        );
    }

    #[rstest]
    fn asset_pool_extend_mixed_assets(mut user_assets: Vec<UserAsset>, process: Process) {
        // Start with some commissioned assets
        let mut asset_pool = AssetPool::new();
        asset_pool.commission_new(2020, &mut user_assets);

        // Create a new non-commissioned asset
        let new_asset = Asset::new_ready(
            "agent_new".into(),
            process.into(),
            "GBR".into(),
            Capacity(3.0),
            2015,
        )
        .unwrap()
        .into();

        // Extend with just the new asset (not mixing with existing to avoid duplicates)
        asset_pool.extend(vec![new_asset]);

        assert_eq!(asset_pool.assets.len(), 3);
        // Check that we have the original assets plus the new one
        assert!(asset_pool.assets.iter().any(|a| a.id() == Some(AssetID(0))));
        assert!(asset_pool.assets.iter().any(|a| a.id() == Some(AssetID(1))));
        assert!(asset_pool.assets.iter().any(|a| a.id() == Some(AssetID(2))));
        // Check that the new asset has the correct agent
        assert!(
            asset_pool
                .assets
                .iter()
                .any(|a| a.agent_id() == Some(&"agent_new".into()))
        );
    }

    #[rstest]
    fn asset_pool_extend_maintains_sort_order(mut user_assets: Vec<UserAsset>, process: Process) {
        // Start with some commissioned assets
        let mut asset_pool = AssetPool::new();
        asset_pool.commission_new(2020, &mut user_assets);

        // Create new assets that would be out of order if added at the end
        let process_rc = Rc::new(process);
        let new_assets = vec![
            Asset::new_ready(
                "agent_high_id".into(),
                Rc::clone(&process_rc),
                "GBR".into(),
                Capacity(1.0),
                2010,
            )
            .unwrap()
            .into(),
            Asset::new_ready(
                "agent_low_id".into(),
                Rc::clone(&process_rc),
                "GBR".into(),
                Capacity(1.0),
                2015,
            )
            .unwrap()
            .into(),
        ];

        asset_pool.extend(new_assets);

        // Check that assets are sorted by ID
        let ids: Vec<u32> = asset_pool.iter().map(|a| a.id().unwrap().0).collect();
        assert_equal(ids, 0..4);
    }

    #[rstest]
    fn asset_pool_extend_no_duplicates_expected(mut user_assets: Vec<UserAsset>) {
        // Start with some commissioned assets
        let mut asset_pool = AssetPool::new();
        asset_pool.commission_new(2020, &mut user_assets);
        let original_count = asset_pool.assets.len();

        // The extend method expects unique assets - adding duplicates would violate
        // the debug assertion, so this test verifies the normal case
        asset_pool.extend(Vec::new());

        assert_eq!(asset_pool.assets.len(), original_count);
        // Verify all assets are still unique (this is what the debug_assert checks)
        assert_eq!(
            asset_pool.assets.iter().unique().count(),
            asset_pool.assets.len()
        );
    }

    #[rstest]
    fn asset_pool_extend_increments_next_id(mut user_assets: Vec<UserAsset>, process: Process) {
        // Start with some commissioned assets
        let mut asset_pool = AssetPool::new();
        asset_pool.commission_new(2020, &mut user_assets);
        assert_eq!(asset_pool.next_id, 2); // Should be 2 after commissioning 2 assets

        // Create new non-commissioned assets
        let process_rc = Rc::new(process);
        let new_assets = vec![
            Asset::new_ready(
                "agent1".into(),
                Rc::clone(&process_rc),
                "GBR".into(),
                Capacity(1.0),
                2015,
            )
            .unwrap()
            .into(),
            Asset::new_ready(
                "agent2".into(),
                Rc::clone(&process_rc),
                "GBR".into(),
                Capacity(1.0),
                2020,
            )
            .unwrap()
            .into(),
        ];

        asset_pool.extend(new_assets);

        // next_id should have incremented for each new asset
        assert_eq!(asset_pool.next_id, 4);
        assert_eq!(asset_pool.assets[2].id(), Some(AssetID(2)));
        assert_eq!(asset_pool.assets[3].id(), Some(AssetID(3)));
    }

    #[rstest]
    fn asset_pool_mothball_unretained(mut user_assets: Vec<UserAsset>) {
        // Commission some assets
        let mut asset_pool = AssetPool::new();
        asset_pool.commission_new(2020, &mut user_assets);
        assert_eq!(asset_pool.assets.len(), 2);

        // Remove one asset from the active pool (simulating it being removed elsewhere)
        let removed_asset = asset_pool.assets.remove(0);
        assert_eq!(asset_pool.assets.len(), 1);

        // Try to mothball both the removed asset (not in active) and an active asset
        let assets_to_check = vec![removed_asset.clone(), asset_pool.assets[0].clone()];
        asset_pool.mothball_unretained(assets_to_check, 2025);

        // Only the removed asset should be mothballed (since it's not in active pool)
        assert_eq!(asset_pool.assets.len(), 2); // And should be back into the pool
        assert_equal(
            asset_pool.assets[0].get_mothball_events().unwrap().iter(),
            &[MothballEvent {
                year: 2025,
                num_units: 1,
            }],
        );
    }

    #[rstest]
    fn asset_pool_decommission_unused(mut user_assets: Vec<UserAsset>) {
        // Commission some assets
        let mut asset_pool = AssetPool::new();
        asset_pool.commission_new(2020, &mut user_assets);
        assert_eq!(asset_pool.assets.len(), 2);

        // Make an asset unused for a few years
        let mothball_years: u32 = 10;
        asset_pool.assets[0] = asset_pool[0]
            .clone()
            .with_mothballed_units(1, Some(2025 - mothball_years));

        assert_equal(
            asset_pool.assets[0].get_mothball_events().unwrap().iter(),
            &[MothballEvent {
                year: 2025 - mothball_years,
                num_units: 1,
            }],
        );

        // Decommission unused assets
        asset_pool.decommission_mothballed(2025, mothball_years);

        // Only the removed asset should be decommissioned (since it's not in active pool)
        assert_eq!(asset_pool.assets.len(), 1); // Active pool unchanged
    }

    #[rstest]
    fn asset_pool_decommission_if_not_active_none_active(mut user_assets: Vec<UserAsset>) {
        // Commission some assets
        let mut asset_pool = AssetPool::new();
        asset_pool.commission_new(2020, &mut user_assets);
        let all_assets = asset_pool.assets.clone();

        // Clear the active pool (simulating all assets being removed)
        asset_pool.assets.clear();

        // Try to mothball the assets that are no longer active
        asset_pool.mothball_unretained(all_assets.clone(), 2025);

        // All assets should be mothballed
        assert_eq!(asset_pool.assets.len(), 2);
        assert_eq!(asset_pool.assets[0].id(), all_assets[0].id());
        assert_equal(
            asset_pool.assets[0].get_mothball_events().unwrap().iter(),
            &[MothballEvent {
                year: 2025,
                num_units: 1,
            }],
        );
        assert_eq!(asset_pool.assets[1].id(), all_assets[1].id());
        assert_equal(
            asset_pool.assets[1].get_mothball_events().unwrap().iter(),
            &[MothballEvent {
                year: 2025,
                num_units: 1,
            }],
        );
    }

    #[rstest]
    #[should_panic(expected = "Cannot mothball asset that has not been commissioned")]
    fn asset_pool_decommission_if_not_active_non_commissioned_asset(process: Process) {
        // Create a non-commissioned asset
        let non_commissioned_asset = Asset::new_ready(
            "agent_new".into(),
            process.into(),
            "GBR".into(),
            Capacity(1.0),
            2015,
        )
        .unwrap()
        .into();

        // This should panic because the asset was never commissioned
        let mut asset_pool = AssetPool::new();
        asset_pool.mothball_unretained(vec![non_commissioned_asset], 2025);
    }

    /// A commissioned divisible asset with three units.
    #[fixture]
    fn commissioned_divisible(mut asset_divisible: Asset) -> AssetRef {
        asset_divisible.commission(AssetID(0));
        assert_eq!(asset_divisible.num_units(), 3);
        AssetRef::from(asset_divisible)
    }

    #[rstest]
    fn asset_pool_mothball_unretained_partial(commissioned_divisible: AssetRef) {
        // The full asset has three units; only two of them were retained in the pool
        let full = commissioned_divisible;
        let retained = full.clone().with_subset_of_units(2);

        let mut asset_pool = AssetPool::new();
        asset_pool.assets.push(retained);

        asset_pool.mothball_unretained(vec![full], 2025);

        // The asset is restored to its full capacity, with the unretained unit mothballed
        assert_eq!(asset_pool.assets.len(), 1);
        let asset = &asset_pool.assets[0];
        assert_eq!(asset.num_units(), 3);
        assert_eq!(asset.get_num_mothballed_units(), 1);
        assert_equal(
            asset.get_mothball_events().unwrap().iter(),
            &[MothballEvent {
                year: 2025,
                num_units: 1,
            }],
        );
    }

    #[rstest]
    fn asset_pool_decommission_mothballed_partial(commissioned_divisible: AssetRef) {
        // Mothball one unit in 2010 and one in 2020, leaving one active
        let asset = commissioned_divisible
            .with_mothballed_units(1, Some(2010))
            .with_mothballed_units(2, Some(2020));

        let mut asset_pool = AssetPool::new();
        asset_pool.assets.push(asset);

        // Threshold of 2015: only the unit mothballed in 2010 is old enough to decommission
        asset_pool.decommission_mothballed(2025, 10);

        assert_eq!(asset_pool.assets.len(), 1);
        let asset = &asset_pool.assets[0];
        assert_eq!(asset.num_units(), 2);
        assert_eq!(asset.get_num_mothballed_units(), 1);
        assert_equal(
            asset.get_mothball_events().unwrap().iter(),
            &[MothballEvent {
                year: 2020,
                num_units: 1,
            }],
        );
    }

    #[rstest]
    fn asset_pool_decommission_mothballed_removes_fully_mothballed(
        commissioned_divisible: AssetRef,
    ) {
        // All three units mothballed long enough ago: the whole asset is removed from the pool
        let asset = commissioned_divisible.with_mothballed_units(3, Some(2010));

        let mut asset_pool = AssetPool::new();
        asset_pool.assets.push(asset);

        asset_pool.decommission_mothballed(2025, 10);

        assert!(asset_pool.assets.is_empty());
    }
}