lasprs 0.14.2

//! Loading and storing CPS files from a Measurement.

type Result<T> = std::result::Result<T, MeasurementError>;
use super::*;
use crate::{
    ps::CPSSettings,
    tools::{get_current_timestamp, h5::*},
    *,
};
use hdf5_metno::{Error, File};
use itertools::Itertools;
use ndarray::ArrayView3;
use snafu::prelude::*;
use std::hash::{DefaultHasher, Hash, Hasher};

/// Group name for cached cross-power spectrum (cps)
const CACHED_CPS_GROUPNAME: &str = "cached_cps";

// HDF dataset base for cached cross-power spectrum (cps)
const CACHED_CPS_DATASETNAME_BASE: &str = "cached_cps/cps_";

// Maximum number of cached CPS datasets
const MAX_NUM_CACHED_CPS: usize = 99;

fn get_dset_name_for_index(index: usize) -> String {
    format!("{}{:02}", CACHED_CPS_DATASETNAME_BASE, index)
}

// Read all current stored CPS metadata in file. Only returns valid sets.
//
// # Returns
//
// A vector of tuples containing the index, dataset name, timestamp, and
// settings hash.
fn read_current_cps_meta(file: &File) -> Vec<(usize, String, Flt, u64)> {
    let mut res = Vec::with_capacity(MAX_NUM_CACHED_CPS);
    for i in 0..MAX_NUM_CACHED_CPS {
        let dset_name = get_dset_name_for_index(i);
        if let Ok(dataset) = file.dataset(&dset_name)
            && let Ok(timestamp) = read_h5_attr_scalar(&dataset, "timestamp")
            && let Ok(hash) = read_h5_attr_scalar(&dataset, "settings_hash")
        {
            res.push((i, dset_name, timestamp, hash));
        }
    }
    res
}

/// Store the computed CPS data in the measurement file
///
/// Args:
///     file: The HDF5 file to store the CPS data in.
///     settings: The settings used to compute the CPS data.
///     cps: The computed CPS data (WITHOUT SENSITIVITY CORRECTION APPLIED)
pub fn store_CPS(file: &File, settings: &CPSSettings, cps: &Array3<Cflt>) -> Result<()> {
    // If this function runs for the first time, it might be the case that the
    // group does not exist yet. In that case, create it.
    let file_members = file
        .member_names()
        .map_err(H5Error::from)
        .context(H5FileProblemSnafu {
            operation: "reading file members",
        })?;
    if !file_members
        .iter()
        .map(|a| a.as_str())
        .any(|m| m == CACHED_CPS_GROUPNAME)
    {
        file.create_group(CACHED_CPS_GROUPNAME)
            .map_err(H5Error::from)
            .context(H5FileProblemSnafu {
                operation: "creating cached CPS group in file",
            })?;
    }

    // Create hash of settings
    let mut hasher = DefaultHasher::new();
    settings.hash(&mut hasher);
    let settings_hash = hasher.finish();

    // Serialize settings to JSON
    let settings_json = serde_json::to_string(settings).expect("Failed to serialize CPS settings");

    // Get current timestamp
    let current_timestamp = get_current_timestamp();

    // Look for existing dataset with same settings hash, and write new CPS in
    // it if it is found.
    let current_meta = read_current_cps_meta(file);

    let target_index = if current_meta.len() == MAX_NUM_CACHED_CPS {
        // Find oldest dataset index, if the cache is full
        let mut oldest_index = 0;
        let mut oldest_timestamp = current_timestamp;
        for (i, (_, _, timestamp, _)) in current_meta.iter().enumerate() {
            if *timestamp < oldest_timestamp {
                oldest_timestamp = *timestamp;
                oldest_index = i;
            }
        }
        oldest_index
    } else if let Some(i) = current_meta
        .iter()
        .map(|a| a.3)
        .position(|a| a == settings_hash)
    {
        // Return position where the settings hash is found
        i
    } else {
        current_meta.len()
    };

    let dset_name = get_dset_name_for_index(target_index);
    if current_meta.iter().map(|a| &a.1).contains(&dset_name) {
        // Found existing dataset with same settings, update it
        file.unlink(&dset_name)
            .map_err(H5Error::from)
            .context(H5FileProblemSnafu {
                operation: "unlinking old CPS dataset",
            })?;
    }

    // Create new dataset
    let dataset = file
        .new_dataset_builder()
        .with_data(cps)
        .create(dset_name.as_str())
        .map_err(H5Error::from)
        .context(H5FileProblemSnafu {
            operation: "creating new CPS dataset",
        })?;

    // Set attributes
    write_h5_attr_scalar(&dataset, "settings_hash", settings_hash)?;
    write_h5_attr_string(&dataset, "settings_json", &settings_json)?;
    write_h5_attr_scalar(&dataset, "timestamp", current_timestamp)?;

    Ok(())
}

/// Load CPS data from the measurement file based on settings. The resulting
/// data is WITHOUT SENSITIVITY CORRECTION APPLIED.
///
/// # Args
///
/// - `file`: The HDF5 file to load the CPS data from.
/// - `settings`: The settings used to compute the CPS data.
pub fn load_CPS(file: &File, settings: &CPSSettings) -> Result<Option<ndarray::Array3<Cflt>>> {
    // Create hash of settings to match against stored data
    let mut hasher = DefaultHasher::new();
    settings.hash(&mut hasher);
    let settings_hash = hasher.finish();

    // Read all current CPS metadata
    let current_meta = read_current_cps_meta(file);

    // Find dataset with matching settings hash
    for (index, dset_name, _timestamp, hash) in current_meta.into_iter() {
        if hash == settings_hash {
            let dset =
                file.dataset(&dset_name)
                    .map_err(H5Error::from)
                    .context(H5FileProblemSnafu {
                        operation: "reading CPS dataset",
                    })?;
            let settings_json = read_h5_attr_string(&dset, "settings_json")?;
            let this_settings: CPSSettings = serde_json::from_str(&settings_json)
                .map_err(JSONError::from)
                .context(ParseMetaSnafu {
                    field: "settings_json",
                })?;
            if this_settings != *settings {
                // Has collision, skip this dataset
                continue;
            }

            let dset_name = get_dset_name_for_index(index);
            let dataset =
                file.dataset(&dset_name)
                    .map_err(H5Error::from)
                    .context(H5FileProblemSnafu {
                        operation: "reading CPS dataset",
                    })?;

            // Read the CPS data
            let cps_data: ndarray::Array3<Cflt> =
                dataset
                    .read()
                    .map_err(H5Error::from)
                    .context(H5FileProblemSnafu {
                        operation: "reading CPS data",
                    })?;
            return Ok(Some(cps_data));
        }
    }

    // No matching dataset found
    Ok(None)
}

/// Clear all cached CPS datasets from the file
pub fn clear_all_cached_cps(file: &File) -> Result<()> {
    file.unlink(CACHED_CPS_GROUPNAME)
        .map_err(H5Error::from)
        .context(H5FileProblemSnafu {
            operation: "unlinking cached CPS group",
        })?;
    Ok(())
}

/// Remove a specific cached CPS dataset by settings hash
#[expect(dead_code)]
pub fn remove_cached_cps(file: &File, settings: &CPSSettings) -> Result<bool> {
    // Create hash of settings to match against stored data
    let mut hasher = DefaultHasher::new();
    settings.hash(&mut hasher);
    let settings_hash = hasher.finish();

    // Read all current CPS metadata
    let current_meta = read_current_cps_meta(file);

    // Find and remove dataset with matching settings hash
    for (_, dset_name, _, hash) in &current_meta {
        if *hash == settings_hash {
            let dset =
                file.dataset(dset_name)
                    .map_err(H5Error::from)
                    .context(H5FileProblemSnafu {
                        operation: "reading CPS dataset",
                    })?;
            let settings_json = read_h5_attr_string(&dset, "settings_json")?;
            // Check settings, to see if we have no hash collision
            let this_settings: CPSSettings = serde_json::from_str(&settings_json)
                .map_err(JSONError::from)
                .context(ParseMetaSnafu {
                    field: "settings_json",
                })?;
            if *settings == this_settings {
                file.unlink(dset_name)
                    .map_err(H5Error::from)
                    .context(H5FileProblemSnafu {
                        operation: "unlinking CPS dataset",
                    })?;
            }
            return Ok(true);
        }
    }

    // No matching dataset found
    Ok(false)
}