data-beans 0.12.0

use crate::hdf5_io::*;
use crate::sparse_io::*;

use hdf5::types::{H5Type, VarLenUnicode};
use legume_numeric::matrix::common_io::read_lines;
use log::info;
use std::collections::HashMap;

/// gzip level for the numeric `X` datasets. gzip (DEFLATE) is chosen over the
/// backend's Blosc because h5py / anndata read it without any filter plugin.
const GZIP_LEVEL: u8 = 4;

/// The obs column `from-h5ad` names multi-donor cells by: `barcode@donor`.
const DONOR_COLUMN: &str = "donor_id";

/// Upper bound on the chunk length (elements) for the `X` datasets.
const MAX_CHUNK_ELEMS: usize = 1 << 20;

/// The metadata [`to_h5ad`] attaches.
#[derive(Clone, Debug, Default)]
pub struct H5adExport {
    /// Cell-metadata TSV (gzip is fine) whose first column is the cell key,
    /// matched against the column names; the other columns become `obs`
    /// columns. A cell named `barcode@donor` also matches the row keyed
    /// `barcode` whose `donor_id` column is `donor`, or else the row keyed
    /// `barcode`. Unmatched cells get `NA`, as do cells whose key is listed
    /// on rows that differ. The `*.cell_metadata.tsv.gz` that
    /// [`crate::import::from_h5ad`] writes round-trips this way.
    pub obs: Option<Box<str>>,
    /// Feature-metadata TSV whose first column is the feature key, matched
    /// against the row names; the other columns become `var` columns, `NA`
    /// where unmatched or ambiguous.
    pub var: Option<Box<str>>,
}

/// Where `to-h5ad` writes for the backend `data_file`: `output`, given one,
/// with `.h5ad` appended when it lacks it, else `data_file`'s stem plus
/// `.h5ad` (`sim.zarr.zip` → `sim.h5ad`).
pub fn h5ad_path(data_file: &str, output: Option<&str>) -> String {
    match output {
        Some(o) if o.ends_with(".h5ad") => o.to_string(),
        Some(o) => format!("{}.h5ad", o),
        None => format!("{}.h5ad", strip_backend_suffix(data_file)),
    }
}

/// Export `data` to the AnnData `.h5ad` file `output`, readable by scanpy /
/// anndata. This is the inverse of [`crate::import::from_h5ad`].
///
/// data-beans stores the matrix `M` as (features x cells) with `/by_column`
/// holding the CSC of `M` (column pointers over cells, row indices over
/// features). AnnData's `X` is (cells x features) = `Mᵀ`, stored as CSR.
/// Because CSR(`Mᵀ`) is byte-for-byte identical to CSC(`M`), the `/by_column`
/// arrays copy straight into `X` as a `csr_matrix` of shape `[n_cells,
/// n_features]` — no re-sort, exactly undoing the transpose `from_h5ad` did.
/// `data` is preloaded by column for it (and stays so), and the whole CSC
/// has to fit the preload budget.
pub fn to_h5ad(
    data: &mut dyn SparseIo<IndexIter = Vec<usize>>,
    output: &str,
    opts: &H5adExport,
) -> anyhow::Result<()> {
    info!("Exporting to {}", output);

    let nrow = data
        .num_rows()
        .ok_or_else(|| anyhow::anyhow!("backend has no `nrow`"))?; // features -> var
    let ncol = data
        .num_columns()
        .ok_or_else(|| anyhow::anyhow!("backend has no `ncol`"))?; // cells -> obs
    let nnz = data
        .num_non_zeros()
        .ok_or_else(|| anyhow::anyhow!("backend has no `nnz`"))?;

    let row_names = data.row_names()?; // features
    let col_names = data.column_names()?; // cells
    if row_names.len() != nrow {
        anyhow::bail!("row name count {} != nrow {}", row_names.len(), nrow);
    }
    if col_names.len() != ncol {
        anyhow::bail!("column name count {} != ncol {}", col_names.len(), ncol);
    }

    // Load the CSC (`/by_column`) arrays; they map directly onto X = CSR.
    //
    // This export is BUILT on the preloaded arrays — there is no streaming
    // variant yet — so the preload budget does not get a vote here: a skipped
    // preload would surface as a baffling "failed to access preloaded CSC
    // arrays" three lines down. Refuse up front with the actual reason and the
    // knob to turn instead.
    info!("preloading columns ({} cells, {} non-zeros)...", ncol, nnz);
    data.preload_columns()?;
    let (indptr, indices, values) = data.csc_column_arrays().ok_or_else(|| {
        anyhow::anyhow!(
            "to-h5ad needs the whole CSC resident ({} non-zeros x 12 bytes) and the preload \
             was declined by the memory budget; raise LEGUME_PRELOAD_BUDGET_BYTES to export \
             this matrix",
            nnz
        )
    })?;
    debug_assert_eq!(indptr.len(), ncol + 1);
    debug_assert_eq!(indices.len(), values.len());

    let file = hdf5::File::create(output)?;

    // Root: mark as an AnnData container.
    write_str_attr(&file, "encoding-type", "anndata")?;
    write_str_attr(&file, "encoding-version", "0.1.0")?;

    // X: csr_matrix of shape [n_obs = ncol, n_var = nrow].
    write_csr_x(&file, ncol, nrow, indptr, indices, values)?;

    // obs (cells) and var (features) dataframes.
    write_dataframe(&file, "obs", &col_names, opts.obs.as_deref(), true)?;
    write_dataframe(&file, "var", &row_names, opts.var.as_deref(), false)?;

    file.flush()?;
    info!(
        "done: {} ({} obs x {} var, {} non-zeros)",
        output, ncol, nrow, nnz
    );
    Ok(())
}

/// Write `/X` as an AnnData `csr_matrix` group from the backend's CSC arrays.
fn write_csr_x(
    file: &hdf5::File,
    n_obs: usize,
    n_var: usize,
    indptr: &[u64],
    indices: &[u64],
    values: &[f32],
) -> anyhow::Result<()> {
    let x = file.create_group("X")?;
    write_str_attr(&x, "encoding-type", "csr_matrix")?;
    write_str_attr(&x, "encoding-version", "0.1.0")?;
    x.new_attr::<i64>()
        .shape([2])
        .create("shape")?
        .write_raw(&[n_obs as i64, n_var as i64])?;

    // Values keep their f32 dtype. Column indices (features) and the row
    // pointer are stored as the narrowest signed int scipy will accept:
    // i32 unless the matrix is large enough to overflow it.
    write_gzip_dataset(&x, "data", values)?;

    if values.len() > i32::MAX as usize || n_var > i32::MAX as usize {
        write_index_arrays(&x, indptr, indices, |v| v as i64)?;
    } else {
        write_index_arrays(&x, indptr, indices, |v| v as i32)?;
    }
    Ok(())
}

/// Write the CSR `indices`/`indptr` datasets, casting the backend's `u64`
/// arrays to the chosen integer width.
fn write_index_arrays<T: H5Type>(
    x: &hdf5::Group,
    indptr: &[u64],
    indices: &[u64],
    cast: impl Fn(u64) -> T,
) -> anyhow::Result<()> {
    write_gzip_dataset(
        x,
        "indices",
        &indices.iter().map(|&v| cast(v)).collect::<Vec<T>>(),
    )?;
    write_gzip_dataset(
        x,
        "indptr",
        &indptr.iter().map(|&v| cast(v)).collect::<Vec<T>>(),
    )?;
    Ok(())
}

/// Write an AnnData `dataframe` group (`obs` or `var`).
///
/// * `index` — the observation / variable names (`_index` dataset).
/// * `meta_tsv` — optional metadata to align onto `index` and add as columns.
/// * `barcode_fallback` — when matching `meta_tsv` keys, also match a cell
///   named `barcode@donor` (as data-beans names multi-donor cells) by its
///   barcode and donor, or by its barcode alone.
fn write_dataframe(
    file: &hdf5::File,
    name: &str,
    index: &[Box<str>],
    meta_tsv: Option<&str>,
    barcode_fallback: bool,
) -> anyhow::Result<()> {
    let group = file.create_group(name)?;
    write_str_attr(&group, "encoding-type", "dataframe")?;
    write_str_attr(&group, "encoding-version", "0.2.0")?;
    write_str_attr(&group, "_index", "_index")?;

    write_string_array(&group, "_index", index)?;

    let (col_names, columns) = match meta_tsv {
        Some(path) => load_aligned_metadata(path, index, barcode_fallback)?,
        None => (vec![], vec![]),
    };
    for (cname, col) in col_names.iter().zip(columns.iter()) {
        write_string_array(&group, cname, col)?;
    }

    // column-order lists the non-index columns (empty array when none).
    let order: Vec<VarLenUnicode> = col_names
        .iter()
        .map(|c| vlu(c))
        .collect::<anyhow::Result<_>>()?;
    group
        .new_attr::<VarLenUnicode>()
        .shape([order.len()])
        .create("column-order")?
        .write_raw(&order)?;

    Ok(())
}

/// Metadata column names, and each column's entries in the order of the
/// names they were aligned onto.
type AlignedMetadata = (Vec<Box<str>>, Vec<Vec<Box<str>>>);

/// Parse a metadata TSV and align its rows onto `names`.
///
/// The first TSV column is the join key; the rest become columns (a literal
/// `_index` column, if present, is dropped to avoid clashing with the index
/// dataset). Returns `(column_names, columns)` where every column has the same
/// length as `names`; unmatched entries are filled with `"NA"`.
///
/// With `barcode_fallback`, a name `barcode@donor` that is no key itself
/// matches the row keyed `barcode` whose `donor_id` column says `donor`
/// (the name `from-h5ad` gives that row's cell), else the row keyed
/// `barcode`. A key the file lists on several rows stands for them while
/// they say the same in every kept column; once they differ it names no
/// single row, matches nothing, and the entry is left `"NA"`, with a
/// warning.
fn load_aligned_metadata(
    path: &str,
    names: &[Box<str>],
    barcode_fallback: bool,
) -> anyhow::Result<AlignedMetadata> {
    let lines = read_lines(path)?;
    if lines.is_empty() {
        log::warn!("metadata file {} is empty; no columns attached", path);
        return Ok((vec![], vec![]));
    }

    let header: Vec<&str> = lines[0].split('\t').collect();
    // Keep every non-key column except a literal "_index".
    let (keep_idx, col_names): (Vec<usize>, Vec<Box<str>>) = header
        .iter()
        .enumerate()
        .skip(1)
        .filter(|(_, &c)| c != "_index")
        .map(|(i, &c)| (i, Box::<str>::from(c)))
        .unzip();

    let rows: Vec<Vec<&str>> = lines[1..]
        .iter()
        .map(|line| line.split('\t').collect())
        .collect();
    // What a row says in the kept columns, `NA` where it is short.
    let kept = |row: usize| {
        let fields = &rows[row];
        keep_idx
            .iter()
            .map(move |&hidx| fields.get(hidx).copied().unwrap_or("NA"))
    };
    // Key -> the first row listing it, or `None` once its rows disagree.
    let mut key2row: HashMap<Box<str>, Option<usize>> = HashMap::new();
    let mut add_key = |key: Box<str>, row: usize| {
        key2row
            .entry(key)
            .and_modify(|listed| *listed = listed.filter(|&first| kept(first).eq(kept(row))))
            .or_insert(Some(row));
    };
    let donor_col = header.iter().position(|&c| c == DONOR_COLUMN);
    for (i, fields) in rows.iter().enumerate() {
        let Some(&key) = fields.first() else {
            continue;
        };
        add_key(key.into(), i);
        // The name of this row's cell in a multi-donor backend.
        let donor = donor_col.and_then(|d| fields.get(d));
        if let (true, Some(donor)) = (barcode_fallback, donor) {
            if !key.contains(COLUMN_SEP) {
                add_key(format!("{key}{COLUMN_SEP}{donor}").into(), i);
            }
        }
    }

    let mut columns: Vec<Vec<Box<str>>> = vec![Vec::with_capacity(names.len()); col_names.len()];
    let mut n_matched = 0usize;
    let mut n_ambiguous = 0usize;
    for nm in names {
        let listed = key2row.get(nm).or_else(|| {
            if barcode_fallback {
                nm.split(COLUMN_SEP).next().and_then(|bc| key2row.get(bc))
            } else {
                None
            }
        });
        n_ambiguous += usize::from(listed == Some(&None));
        match listed.copied().flatten() {
            Some(row) => {
                n_matched += 1;
                for (column, value) in columns.iter_mut().zip(kept(row)) {
                    column.push(value.into());
                }
            }
            None => {
                for col in columns.iter_mut() {
                    col.push("NA".into());
                }
            }
        }
    }
    if n_ambiguous > 0 {
        log::warn!(
            "{} of {} entries match a key that {} lists on rows that differ; \
             their metadata is left NA",
            n_ambiguous,
            names.len(),
            path
        );
    }

    info!(
        "attached {} metadata column(s) from {}; matched {}/{} entries",
        col_names.len(),
        path,
        n_matched,
        names.len()
    );
    Ok((col_names, columns))
}

//////////////////////
// HDF5 write helpers //
//////////////////////

/// Encode a Rust string as an HDF5 variable-length UTF-8 string.
pub(crate) fn vlu(s: &str) -> anyhow::Result<VarLenUnicode> {
    s.parse::<VarLenUnicode>()
        .map_err(|e| anyhow::anyhow!("cannot encode '{}' as an HDF5 string: {}", s, e))
}

/// Write a scalar UTF-8 string attribute onto a group or dataset.
pub(crate) fn write_str_attr(loc: &hdf5::Location, name: &str, value: &str) -> anyhow::Result<()> {
    loc.new_attr::<VarLenUnicode>()
        .create(name)?
        .write_scalar(&vlu(value)?)?;
    Ok(())
}

/// Create a gzip-compressed 1-D numeric dataset and write `data` into it.
/// Falls back to a plain contiguous dataset when empty (a chunked layout
/// needs a non-zero chunk size).
fn write_gzip_dataset<T: H5Type>(
    group: &hdf5::Group,
    name: &str,
    data: &[T],
) -> anyhow::Result<()> {
    let n = data.len();
    let ds = if n > 0 {
        let chunk = n.min(MAX_CHUNK_ELEMS);
        group
            .new_dataset::<T>()
            .shape(n)
            .deflate(GZIP_LEVEL)
            .chunk([chunk])
            .create(name)?
    } else {
        group.new_dataset::<T>().shape(n).create(name)?
    };
    ds.write_raw(data)?;
    Ok(())
}

/// Write a `string-array`-encoded dataset of UTF-8 strings.
fn write_string_array(group: &hdf5::Group, name: &str, values: &[Box<str>]) -> anyhow::Result<()> {
    let encoded: Vec<VarLenUnicode> = values
        .iter()
        .map(|s| vlu(s))
        .collect::<anyhow::Result<_>>()?;
    let ds = group
        .new_dataset::<VarLenUnicode>()
        .shape(encoded.len())
        .create(name)?;
    ds.write_raw(&encoded)?;
    write_str_attr(&ds, "encoding-type", "string-array")?;
    write_str_attr(&ds, "encoding-version", "0.2.0")?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    fn names(xs: &[&str]) -> Vec<Box<str>> {
        xs.iter().map(|&x| Box::from(x)).collect()
    }

    /// Two donors share a barcode: each cell gets its own donor's row, and
    /// where the file cannot tell them apart, neither gets the other's.
    #[test]
    fn a_barcode_shared_by_two_donors_keeps_each_donors_metadata() {
        let dir = tempfile::tempdir().unwrap();
        let cells = names(&["AAAC-1@donorA", "AAAC-1@donorB", "GGGT-1@donorA"]);

        let with_donor = dir.path().join("with_donor.tsv");
        std::fs::write(
            &with_donor,
            "barcode\tdonor_id\tcell_type\n\
             AAAC-1\tdonorA\tT\n\
             AAAC-1\tdonorB\tB\n\
             GGGT-1\tdonorA\tNK\n",
        )
        .unwrap();
        let (cols, data) =
            load_aligned_metadata(with_donor.to_str().unwrap(), &cells, true).unwrap();
        assert_eq!(cols, names(&["donor_id", "cell_type"]));
        assert_eq!(data[0], names(&["donorA", "donorB", "donorA"]));
        assert_eq!(data[1], names(&["T", "B", "NK"]));

        let bare = dir.path().join("bare.tsv");
        std::fs::write(
            &bare,
            "barcode\tcell_type\nAAAC-1\tT\nAAAC-1\tB\nGGGT-1\tNK\n",
        )
        .unwrap();
        let (_, data) = load_aligned_metadata(bare.to_str().unwrap(), &cells, true).unwrap();
        assert_eq!(data[0], names(&["NA", "NA", "NK"]));
    }

    /// A key on several rows keeps its metadata while the rows agree, in an
    /// obs table and a var table alike, and loses it once they differ.
    #[test]
    fn a_key_listed_twice_is_ambiguous_only_when_its_rows_differ() {
        let dir = tempfile::tempdir().unwrap();
        let table = dir.path().join("table.tsv");
        std::fs::write(
            &table,
            "key\tkind\tnote\n\
             A\tx\tsame\n\
             A\tx\tsame\n\
             B\ty\tone\n\
             B\ty\tother\n\
             C\tz\tonly\n",
        )
        .unwrap();
        let entries = names(&["A", "B", "C", "D"]);
        for barcode_fallback in [true, false] {
            let (cols, data) =
                load_aligned_metadata(table.to_str().unwrap(), &entries, barcode_fallback).unwrap();
            assert_eq!(cols, names(&["kind", "note"]));
            assert_eq!(data[0], names(&["x", "NA", "z", "NA"]));
            assert_eq!(data[1], names(&["same", "NA", "only", "NA"]));
        }

        // Cells named by donor: a repeated row still names its cell, by
        // barcode and donor and by barcode alone.
        let obs = dir.path().join("obs.tsv");
        std::fs::write(
            &obs,
            "barcode\tdonor_id\tcell_type\n\
             AAAC-1\tdonorA\tT\n\
             AAAC-1\tdonorA\tT\n\
             GGGT-1\tdonorA\tNK\n\
             GGGT-1\tdonorA\tB\n",
        )
        .unwrap();
        let cells = names(&["AAAC-1@donorA", "AAAC-1@donorB", "GGGT-1@donorA"]);
        let (_, data) = load_aligned_metadata(obs.to_str().unwrap(), &cells, true).unwrap();
        assert_eq!(data[1], names(&["T", "T", "NA"]));
    }
}