data_beans/

zarr_io.rs

//! General-purpose zarr I/O for reading numeric arrays and attributes from
//! zarr directories or `.zarr.zip` archives.
//!
//! Also provides a convenience function for reading coordinate-style data
//! (row names + numeric matrix) that returns `MatWithNames<DMatrix<f32>>`,
//! matching the parquet reader interface in legume_numeric::matrix.

use crate::sparse_io::SparseIoBackend;
use legume_numeric::matrix::traits::MatWithNames;
use log::info;
use nalgebra::DMatrix;
use rand_distr::num_traits::FromPrimitive;
use std::sync::Arc;
use zarrs::array::{data_type, Array as ZArray};
use zarrs::config::MetadataRetrieveVersion;
use zarrs::filesystem::FilesystemStore;
use zarrs::storage::ReadableListableStorageTraits as ZReadStorageTraits;
use zarrs_zip::ZipStorageAdapter;

// ── Store management ────────────────────────────────────────────────────

/// A read store paired with an optional write store (present for directories, `None` for zips).
pub type ZarrStoreRw = (Arc<dyn ZReadStorageTraits>, Option<Arc<FilesystemStore>>);

/// The directory a zarr zip holds its store under, read from the stored
/// names, which [`ZipStorageAdapter`] matches too: empty when a metadata file
/// (`zarr.json`, or v2's `.zgroup` / `.zarray`) sits at the archive's root;
/// else the single top-level directory (as
/// [`legume_numeric::matrix::common_io::zip_dir`] writes, whatever the
/// archive is named), ignoring loose files and `__MACOSX`; else, among
/// several, the one named like the archive (`foo` for `foo.zarr.zip`, or the
/// older `foo.zarr`); else empty.
fn zip_zarr_root(zip_path: &std::path::Path) -> anyhow::Result<String> {
    const METADATA: [&str; 3] = ["zarr.json", ".zgroup", ".zarray"];
    let file = std::fs::File::open(zip_path)?;
    let archive = zip::ZipArchive::new(std::io::BufReader::new(file))?;
    let mut dirs = std::collections::BTreeSet::new();
    for name in archive.file_names().flatten() {
        match name.split_once('/') {
            None if METADATA.contains(&&*name) => return Ok(String::new()),
            None | Some(("__MACOSX", _)) => {}
            Some((top, _)) => {
                dirs.insert(top.to_owned());
            }
        }
    }
    if dirs.len() == 1 {
        return Ok(dirs.pop_first().unwrap_or_default());
    }
    let file_name = zip_path.file_name().and_then(|f| f.to_str()).unwrap_or("");
    let stem = file_name.strip_suffix(".zip").unwrap_or(file_name);
    let named = [stem.strip_suffix(".zarr").unwrap_or(stem), stem];
    Ok(named
        .into_iter()
        .find(|d| dirs.contains(*d))
        .unwrap_or_default()
        .to_owned())
}

/// The path prefix inside a zarr zip: its [`zip_zarr_root`] with a trailing
/// `/`, or empty.
fn detect_zip_zarr_prefix(zip_path: &std::path::Path) -> anyhow::Result<Box<str>> {
    let root = zip_zarr_root(zip_path)?;
    Ok(if root.is_empty() {
        root.into()
    } else {
        format!("{root}/").into()
    })
}

/// Open a zarr store, returning both a read store and an optional write store.
///
/// Supports both `.zarr` directories and `.zarr.zip` archives.
/// For zip archives, uses [`ZipStorageAdapter`] for direct random-access
/// reads without extracting to a temp directory; the write store is `None`.
pub fn open_zarr_store_rw(path: &str) -> anyhow::Result<ZarrStoreRw> {
    let p = std::path::Path::new(path);
    let is_zip = p.extension().is_some_and(|e| e == "zip");

    if is_zip {
        let parent = p
            .parent()
            .ok_or_else(|| anyhow::anyhow!("no parent directory for {}", path))?;
        let filename = p
            .file_name()
            .ok_or_else(|| anyhow::anyhow!("no filename for {}", path))?
            .to_str()
            .ok_or_else(|| anyhow::anyhow!("non-UTF8 filename in {}", path))?;
        let zarr_prefix = detect_zip_zarr_prefix(p)?;
        let fs = Arc::new(FilesystemStore::new(parent)?);
        let key = zarrs::storage::StoreKey::new(filename)?;
        info!(
            "Opening zarr zip store: {} (prefix: {:?})",
            path, zarr_prefix
        );
        Ok((
            Arc::new(ZipStorageAdapter::new_with_path(fs, key, &*zarr_prefix)?),
            None,
        ))
    } else {
        let fs = Arc::new(FilesystemStore::new(path)?);
        Ok((fs.clone(), Some(fs)))
    }
}

/// Open a zarr store for reading (convenience wrapper around [`open_zarr_store_rw`]).
pub fn open_zarr_store(path: &str) -> anyhow::Result<Arc<dyn ZReadStorageTraits>> {
    open_zarr_store_rw(path).map(|(r, _)| r)
}

/// The file `-o output` ends as, under `--no-zip` (`zip` false) and
/// `--backend`: `x.zarr.zip`, `x.zarr` or `x.h5`, for a plan to name.
pub fn output_path(output: &str, zip: bool, backend: &SparseIoBackend) -> String {
    let effective = apply_zip_flag(output, zip, backend);
    if effective.ends_with(".zarr.zip") {
        return effective.to_string();
    }
    crate::hdf5_io::resolve_backend_file(&effective, Some(backend.clone()))
        .map_or_else(|_| effective.to_string(), |(_, f)| f.to_string())
}

/// If `zip` is true and the target is zarr, ensure the output path ends with
/// `.zarr.zip`; otherwise pass through unchanged.
///
/// The target is treated as HDF5 (and left alone) when `backend` is
/// [`SparseIoBackend::HDF5`] or the name already carries a `.h5` / `.hdf5`
/// suffix. Without the `backend` signal, a bare `-o foo` under `--backend hdf5`
/// would be zarr-ified here and then flipped back to Zarr by
/// `resolve_backend_file` — silently ignoring the requested backend.
pub fn apply_zip_flag(output: &str, zip: bool, backend: &SparseIoBackend) -> Box<str> {
    // The `.zarr.zip` suffix only makes sense for a zarr target. An HDF5 target
    // — selected by `--backend hdf5` or an explicit `.h5`/`.hdf5` name — is left
    // untouched, so a bare `-o foo` under `--backend hdf5` resolves to `foo.h5`
    // rather than being silently zarr-ified by `resolve_backend_file`.
    let target_is_hdf5 = matches!(backend, SparseIoBackend::HDF5)
        || output.ends_with(".h5")
        || output.ends_with(".hdf5");
    if zip && !target_is_hdf5 && !output.ends_with(".zarr.zip") {
        let base = crate::hdf5_io::strip_backend_suffix(output);
        format!("{}.zarr.zip", base).into()
    } else {
        output.into()
    }
}

/// Where a handler given `-o output` writes, before it writes anything.
pub struct OutputTarget {
    /// The output as the handler takes it: `.zarr.zip` added under zip.
    pub effective: Box<str>,
    pub backend: SparseIoBackend,
    /// The file the backend is written to (for a `.zarr.zip` target, the
    /// un-zipped `.zarr` directory [`staging_zarr_dir`] names).
    pub working: Box<str>,
    /// What the user finds once [`finalize_output`] has run: the archive for
    /// a `.zarr.zip` target, else the working file.
    pub found: Box<str>,
}

/// [`OutputTarget`] for `output`, touching nothing.
pub fn output_target(
    output: &str,
    backend: SparseIoBackend,
    zip: bool,
) -> anyhow::Result<OutputTarget> {
    let effective = apply_zip_flag(output, zip, &backend);
    let (backend, working) = crate::hdf5_io::resolve_backend_file(&effective, Some(backend))?;
    let found = if effective.ends_with(".zarr.zip") {
        effective.clone()
    } else {
        working.clone()
    };
    Ok(OutputTarget {
        effective,
        backend,
        working,
        found,
    })
}

/// Resolve the output target for a handler that writes a fresh backend, and
/// clear any pre-existing file at that path. Returns
/// `(effective_output, backend, working_file)`, where `working_file` is the
/// path to write (the un-zipped `.zarr` directory for a `.zarr.zip` target) —
/// pair it with [`finalize_output`] once the backend is populated.
pub fn prepare_output(
    output: &str,
    backend: SparseIoBackend,
    zip: bool,
) -> anyhow::Result<(Box<str>, SparseIoBackend, Box<str>)> {
    use legume_numeric::matrix::common_io::remove_file;

    let t = output_target(output, backend, zip)?;
    if std::path::Path::new(t.working.as_ref()).exists() {
        remove_file(&t.working)?;
    }
    Ok((t.effective, t.backend, t.working))
}

/// Finalize a handler's output: re-zip the working `.zarr` directory when the
/// target is `.zarr.zip`, and return the path the user will actually find (the
/// archive for a `.zarr.zip` target, otherwise the working file).
pub fn finalize_output<'a>(
    working_file: &'a str,
    effective_output: &'a str,
) -> anyhow::Result<&'a str> {
    finalize_zarr_output(working_file, effective_output)?;
    Ok(if effective_output.ends_with(".zarr.zip") {
        effective_output
    } else {
        working_file
    })
}

/// Extract a `.zarr.zip` archive into `target_dir` as a flat `.zarr`
/// directory: the store under its root directory, with legume-numeric's
/// checks against entries that would leave the target.
pub fn extract_zarr_zip(zip_path: &str, target_dir: &str) -> anyhow::Result<()> {
    let root = zip_zarr_root(std::path::Path::new(zip_path))?;
    legume_numeric::matrix::common_io::unzip_subdir(zip_path, &root, target_dir)
}

/// Copy input backend to a writable output location. Transparently extracts
/// `.zarr.zip` archives into a `.zarr` directory so the output can be opened
/// read/write.
pub fn materialize_writable_backend(src: &str, dst: &str) -> anyhow::Result<()> {
    if src.ends_with(".zarr.zip") {
        info!("extracting {} → {}", src, dst);
        extract_zarr_zip(src, dst)
    } else {
        legume_numeric::matrix::common_io::recursive_copy(src, dst)
    }
}

/// The hidden folder, in the output's own folder, that `.zarr.zip` targets
/// are built in.
const STAGE: &str = ".data-beans-stage";

/// The working `.zarr` directory for a `.zarr.zip` target `x.zarr.zip`:
/// `.data-beans-stage/x.zarr` in the same folder, on the same disk. Built as a
/// plain `x.zarr` beside the archive, every write made a directory there (and
/// subset a `.subset_tmp` one) and removed it seconds later, and a sync client
/// watching the folder, such as Dropbox, put the emptied directories back.
/// [`make_stage`] marks the stage for Dropbox to leave alone, and
/// [`finalize_zarr_output`] removes it once the archive is in place.
pub fn staging_zarr_dir(target: &str) -> String {
    let path = std::path::Path::new(target);
    let name = path
        .file_name()
        .and_then(|f| f.to_str())
        .and_then(|f| f.strip_suffix(".zip"))
        .unwrap_or("data.zarr");
    path.with_file_name(STAGE)
        .join(name)
        .to_string_lossy()
        .into_owned()
}

/// `name` in the stage of the folder `beside` is in (or of the stage `beside`
/// is itself in): where scratch kept with an output goes, such as a subset's
/// `.subset_tmp` or `qc`'s pool. Pair with [`make_stage`] and [`leave_stage`].
pub fn in_stage(beside: &str, name: &str) -> String {
    let stage = match stage_of(beside) {
        Some(stage) => stage.to_path_buf(),
        None => std::path::Path::new(beside).with_file_name(STAGE),
    };
    stage.join(name).to_string_lossy().into_owned()
}

/// The stage `working` is built in, when it is one ([`staging_zarr_dir`]).
fn stage_of(working: &str) -> Option<&std::path::Path> {
    std::path::Path::new(working)
        .parent()
        .filter(|p| p.file_name().is_some_and(|n| n == STAGE))
}

/// Make the stage `working` is built in, if it is in one, and mark it for
/// Dropbox to ignore before anything is written in it.
pub fn make_stage(working: &str) -> anyhow::Result<()> {
    if let Some(stage) = stage_of(working) {
        std::fs::create_dir_all(stage)?;
        mark_sync_ignored(stage);
    }
    Ok(())
}

/// Remove the stage `working` was built in once nothing else is being built
/// there (another run's build keeps it).
pub fn leave_stage(working: &str) {
    if let Some(stage) = stage_of(working) {
        let _ = std::fs::remove_dir(stage);
    }
}

/// Dropbox's "ignored" mark: it neither uploads the directory nor puts it
/// back. Best effort; elsewhere the mark is a harmless attribute.
fn mark_sync_ignored(dir: &std::path::Path) {
    #[cfg(target_os = "macos")]
    let _ = xattr::set(dir, "com.dropbox.ignored", b"1");
    #[cfg(target_os = "linux")]
    let _ = xattr::set(dir, "user.com.dropbox.ignored", b"1");
    #[cfg(not(any(target_os = "macos", target_os = "linux")))]
    let _ = dir;
}

/// If `target_path` ends with `.zarr.zip`, zip the `zarr_dir` into it and
/// remove the directory and its [`staging_zarr_dir`] stage. Otherwise this is
/// a no-op (the directory IS the target).
///
/// To keep in-zip entry names compact, entries are prefixed with the short
/// stem (`foo/...` rather than `foo.zarr/...`).
pub fn finalize_zarr_output(zarr_dir: &str, target_path: &str) -> anyhow::Result<()> {
    if !target_path.ends_with(".zarr.zip") {
        return Ok(());
    }

    let zarr_path = std::path::Path::new(zarr_dir);
    let stem = zarr_path
        .file_name()
        .and_then(|s| s.to_str())
        .map(|s| s.strip_suffix(".zarr").unwrap_or(s))
        .ok_or_else(|| anyhow::anyhow!("invalid zarr_dir: {}", zarr_dir))?
        .to_string();

    // Zip beside the source first (in the stage, so nothing comes and goes in
    // the output folder) so the source .zarr survives any failure (disk full,
    // etc.). Only after the zip succeeds do we atomically rename into place
    // and remove the source — so a crash leaves either the source directory
    // or the final archive, never neither.
    let tmp_target = format!("{zarr_dir}.zip.tmp");
    info!("Zipping zarr output: {} → {}", zarr_dir, target_path);
    if let Err(e) =
        legume_numeric::matrix::common_io::zip_dir_as(zarr_dir, &tmp_target, Some(&stem))
    {
        let _ = std::fs::remove_file(&tmp_target);
        return Err(e);
    }
    if let Err(e) = std::fs::rename(&tmp_target, target_path) {
        let _ = std::fs::remove_file(&tmp_target);
        return Err(e.into());
    }
    std::fs::remove_dir_all(zarr_dir)?;
    leave_stage(zarr_dir);
    Ok(())
}

// ── Attribute reading ───────────────────────────────────────────────────

/// Read an attribute from a zarr array node.
///
/// ```text
/// read_zarr_array_attr::<Vec<String>>(store, "/cell_summary", "column_names")
/// ```
pub fn read_zarr_array_attr<V: serde::de::DeserializeOwned>(
    store: Arc<dyn ZReadStorageTraits>,
    array_path: &str,
    attr_name: &str,
) -> anyhow::Result<V> {
    let arr = ZArray::open_opt(store, array_path, &MetadataRetrieveVersion::Default)?;
    let attr = arr.attributes().get(attr_name).ok_or_else(|| {
        anyhow::anyhow!(
            "attribute '{}' not found on array '{}'",
            attr_name,
            array_path
        )
    })?;
    Ok(serde_json::from_value(attr.clone())?)
}

// ── Flat array readers ──────────────────────────────────────────────────

/// Retrieve a zarr array as a flat `Vec<f32>` (row-major) and its shape.
///
/// Handles f32, f64, u32, u64 source types with automatic conversion.
pub fn read_zarr_flat_f32(
    store: Arc<dyn ZReadStorageTraits>,
    key: &str,
) -> anyhow::Result<(Vec<f32>, Vec<u64>)> {
    let arr = ZArray::open_opt(store, key, &MetadataRetrieveVersion::Default)?;
    let shape = arr.shape().to_vec();
    let subset = arr.subset_all();

    let dt = arr.data_type();
    let data = if *dt == data_type::float32() {
        arr.retrieve_array_subset::<Vec<f32>>(&subset)?
    } else if *dt == data_type::float64() {
        arr.retrieve_array_subset::<Vec<f64>>(&subset)?
            .into_iter()
            .map(|x| x as f32)
            .collect()
    } else if *dt == data_type::uint32() {
        arr.retrieve_array_subset::<Vec<u32>>(&subset)?
            .into_iter()
            .map(|x| x as f32)
            .collect()
    } else if *dt == data_type::uint64() {
        arr.retrieve_array_subset::<Vec<u64>>(&subset)?
            .into_iter()
            .map(|x| x as f32)
            .collect()
    } else {
        anyhow::bail!("unsupported zarr data type: {:?}", dt)
    };

    Ok((data, shape))
}

/// Retrieve a zarr array as a flat `Vec<u32>` and its shape.
pub fn read_zarr_flat_u32(
    store: Arc<dyn ZReadStorageTraits>,
    key: &str,
) -> anyhow::Result<(Vec<u32>, Vec<u64>)> {
    let arr = ZArray::open_opt(store, key, &MetadataRetrieveVersion::Default)?;
    let shape = arr.shape().to_vec();
    let subset = arr.subset_all();

    let dt = arr.data_type();
    let data = if *dt == data_type::uint32() {
        arr.retrieve_array_subset::<Vec<u32>>(&subset)?
    } else if *dt == data_type::uint64() {
        arr.retrieve_array_subset::<Vec<u64>>(&subset)?
            .into_iter()
            .map(|x| x as u32)
            .collect()
    } else {
        anyhow::bail!("unsupported zarr data type for u32: {:?}", dt)
    };

    Ok((data, shape))
}

// ── Generic array/attribute readers (moved from misc.rs) ────────────────

#[cfg(feature = "ndarray")]
/// Read a full ndarray from zarr storage.
pub fn read_zarr_ndarray<T>(
    store: Arc<dyn ZReadStorageTraits>,
    key_name: &str,
) -> anyhow::Result<ndarray::ArrayD<T>>
where
    T: zarrs::array::ElementOwned + FromPrimitive,
{
    let arr = ZArray::open_opt(store, key_name, &MetadataRetrieveVersion::Default)?;

    let dt = arr.data_type();
    if *dt == data_type::float32() {
        let array: ndarray::ArrayD<f32> =
            arr.retrieve_array_subset::<ndarray::ArrayD<f32>>(&arr.subset_all())?;
        Ok(array.mapv(|x| T::from_f32(x).unwrap()))
    } else if *dt == data_type::float64() {
        let array: ndarray::ArrayD<f64> =
            arr.retrieve_array_subset::<ndarray::ArrayD<f64>>(&arr.subset_all())?;
        Ok(array.mapv(|x| T::from_f64(x).unwrap()))
    } else if *dt == data_type::uint32() {
        let array: ndarray::ArrayD<u32> =
            arr.retrieve_array_subset::<ndarray::ArrayD<u32>>(&arr.subset_all())?;
        Ok(array.mapv(|x| T::from_u32(x).unwrap()))
    } else if *dt == data_type::uint64() {
        let array: ndarray::ArrayD<u64> =
            arr.retrieve_array_subset::<ndarray::ArrayD<u64>>(&arr.subset_all())?;
        Ok(array.mapv(|x| T::from_u64(x).unwrap()))
    } else {
        anyhow::bail!("unsupported zarr data type: {:?}", dt)
    }
}

/// Read a numeric vector from zarr storage.
pub fn read_zarr_numerics<T>(
    store: Arc<dyn ZReadStorageTraits>,
    key_name: &str,
) -> anyhow::Result<Vec<T>>
where
    T: zarrs::array::ElementOwned + FromPrimitive,
{
    let arr = ZArray::open_opt(store, key_name, &MetadataRetrieveVersion::Default)?;

    let dt = arr.data_type();
    let ret = if *dt == data_type::float32() {
        arr.retrieve_array_subset::<Vec<f32>>(&arr.subset_all())?
            .into_iter()
            .map(|x| T::from_f32(x).unwrap())
            .collect()
    } else if *dt == data_type::float64() {
        arr.retrieve_array_subset::<Vec<f64>>(&arr.subset_all())?
            .into_iter()
            .map(|x| T::from_f64(x).unwrap())
            .collect()
    } else if *dt == data_type::uint32() {
        arr.retrieve_array_subset::<Vec<u32>>(&arr.subset_all())?
            .into_iter()
            .map(|x| T::from_u32(x).unwrap())
            .collect()
    } else if *dt == data_type::uint64() {
        arr.retrieve_array_subset::<Vec<u64>>(&arr.subset_all())?
            .into_iter()
            .map(|x| T::from_u64(x).unwrap())
            .collect()
    } else {
        anyhow::bail!("unsupported zarr data type: {:?}", dt);
    };

    Ok(ret)
}

/// Extract an attribute from a zarr group node.
///
/// `key_name` is parsed as `"group_path/attr_name"`, e.g.
/// `"/cell_features/features/id"` → group `"/cell_features/features"`, attr `"id"`.
pub fn read_zarr_group_attr<V>(
    store: Arc<dyn ZReadStorageTraits>,
    key_name: &str,
) -> anyhow::Result<V>
where
    V: serde::de::DeserializeOwned,
{
    use anyhow::Context;

    fn parse_key_name(key_name: &str) -> (Box<str>, Box<str>) {
        let trimmed = key_name.strip_prefix('/').unwrap_or(key_name);
        match trimmed.rsplit_once('/') {
            Some((left, right)) => (
                format!("/{}", left).into_boxed_str(),
                right.to_string().into_boxed_str(),
            ),
            None => (
                "/".to_string().into_boxed_str(),
                trimmed.to_string().into_boxed_str(),
            ),
        }
    }

    let (group_name, attr_name) = parse_key_name(key_name);

    let group = zarrs::group::Group::open_opt(
        store,
        group_name.as_ref(),
        &MetadataRetrieveVersion::Default,
    )
    .with_context(|| format!("Failed to open group '{}'", group_name))?;

    let attr_value = group
        .attributes()
        .get(attr_name.as_ref())
        .with_context(|| {
            format!(
                "Attribute '{}' not found in group '{}'",
                attr_name, group_name
            )
        })?;

    Ok(serde_json::from_value(attr_value.clone())?)
}

/// Read a string array from zarr storage.
pub fn read_zarr_strings(
    store: Arc<dyn ZReadStorageTraits>,
    key_name: &str,
) -> anyhow::Result<Vec<Box<str>>> {
    let arr = ZArray::open_opt(store, key_name, &MetadataRetrieveVersion::Default)?;

    Ok(arr
        .retrieve_array_subset::<Vec<String>>(&arr.subset_all())?
        .into_iter()
        .map(|x| x.into_boxed_str())
        .collect())
}

// ── 10x cell ID encoding ────────────────────────────────────────────────

/// Hex-digit → shifted-alpha lookup table for 10x cell ID encoding.
fn hex_to_shifted_lookup() -> [Option<char>; 256] {
    let mut lookup = [None; 256];
    for (i, ch) in "0123456789abcdef".chars().enumerate() {
        lookup[ch as usize] = Some(
            [
                'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p',
            ][i],
        );
    }
    lookup
}

/// Encode a single `(barcode_u32, suffix_u32)` pair into a 10x cell-ID string.
fn encode_10x_cell_id(
    lookup: &[Option<char>; 256],
    barcode: u32,
    suffix: u32,
) -> anyhow::Result<Box<str>> {
    let barcode: String = format!("{:08x}", barcode)
        .chars()
        .map(|ch| lookup[ch as usize].ok_or_else(|| anyhow::anyhow!("invalid hex char: {}", ch)))
        .collect::<anyhow::Result<String>>()?;
    Ok(format!("{}-{}", barcode, suffix).into_boxed_str())
}

#[cfg(feature = "ndarray")]
/// Parse a 10x Xenium `[N, 2]` u32 cell_id ndarray into string barcodes.
///
/// See [10x docs](https://www.10xgenomics.com/support/software/xenium-onboard-analysis/3.4/advanced/xoa-output-zarr#cellID).
pub fn parse_10x_cell_id(
    input: ndarray::ArrayView<u32, ndarray::IxDyn>,
) -> anyhow::Result<Vec<Box<str>>> {
    anyhow::ensure!(
        input.ndim() == 2 && input.shape()[1] == 2,
        "Must be 2D with shape [N, 2]"
    );
    let lookup = hex_to_shifted_lookup();
    input
        .outer_iter()
        .map(|row| encode_10x_cell_id(&lookup, row[0], row[1]))
        .collect()
}

/// Parse a flat row-major `[N, 2]` u32 buffer into 10x Xenium cell-ID strings.
///
/// Same encoding as `parse_10x_cell_id` but works on a flat slice.
pub fn parse_10x_cell_id_flat(data: &[u32], nrows: usize) -> anyhow::Result<Vec<Box<str>>> {
    anyhow::ensure!(data.len() == nrows * 2, "cell_id buffer size mismatch");
    let lookup = hex_to_shifted_lookup();
    (0..nrows)
        .map(|i| encode_10x_cell_id(&lookup, data[i * 2], data[i * 2 + 1]))
        .collect()
}

// ── Matrix reading (generic) ────────────────────────────────────────────

/// Build a column-major `DMatrix<f32>` by selecting columns from a flat
/// row-major `[nrows, ncols_total]` buffer.
fn select_columns_to_dmatrix(
    data: &[f32],
    nrows: usize,
    ncols_total: usize,
    selected: &[usize],
) -> anyhow::Result<DMatrix<f32>> {
    let k = selected.len();
    let mut col_major = vec![0.0f32; nrows * k];
    for (col_out, &col_in) in selected.iter().enumerate() {
        anyhow::ensure!(col_in < ncols_total, "column index {} out of range", col_in);
        for row in 0..nrows {
            col_major[row + col_out * nrows] = data[row * ncols_total + col_in];
        }
    }
    Ok(DMatrix::from_vec(nrows, k, col_major))
}

/// Resolve column selection: explicit indices, name lookup, or default.
fn resolve_columns(
    column_indices: &[usize],
    column_names: &[Box<str>],
    all_col_names: &[Box<str>],
    default_cols: &[usize],
) -> anyhow::Result<Vec<usize>> {
    if !column_indices.is_empty() {
        Ok(column_indices.to_vec())
    } else if !column_names.is_empty() {
        // Keep only names that exist — allows generous defaults
        // covering multiple platforms (e.g. Visium + Xenium).
        let matched: Vec<usize> = column_names
            .iter()
            .filter_map(|name| all_col_names.iter().position(|c| c == name))
            .collect();
        if matched.is_empty() {
            anyhow::bail!(
                "none of the requested columns {:?} found (available: {:?})",
                column_names,
                all_col_names
            );
        }
        Ok(matched)
    } else {
        Ok(default_cols.to_vec())
    }
}

/// Read a numeric matrix from a zarr array, selecting columns by index or
/// name (looked up via an attribute on the array).
///
/// * `file_path` — zarr directory or `.zarr.zip`
/// * `data_array` — path to the numeric `[N, C]` array (e.g. `/cell_summary`)
/// * `col_names_attr` — attribute name on `data_array` that holds column
///   names (e.g. `"column_names"`); pass `None` to skip name lookup
/// * `row_names_array` — path to the row-name array (e.g. `/cell_id`);
///   `None` generates numeric row names `"0", "1", …`
/// * `row_names_10x` — if `true`, parse `row_names_array` as 10x u32 `[N,2]`
///   cell IDs; if `false`, read as string array
/// * `column_indices` / `column_names` — which columns to select
pub fn read_zarr_matrix(
    file_path: &str,
    data_array: &str,
    col_names_attr: Option<&str>,
    row_names_array: Option<&str>,
    row_names_10x: bool,
    column_indices: &[usize],
    column_names: &[Box<str>],
) -> anyhow::Result<MatWithNames<DMatrix<f32>>> {
    let store = open_zarr_store(file_path)?;

    // Read the data array
    let (flat, shape) = read_zarr_flat_f32(store.clone(), data_array)?;
    let nrows = shape[0] as usize;
    let ncols_total = if shape.len() > 1 {
        shape[1] as usize
    } else {
        1
    };

    // Read column names from attribute (if available)
    let all_col_names: Vec<Box<str>> = if let Some(attr) = col_names_attr {
        read_zarr_array_attr(store.clone(), data_array, attr).unwrap_or_else(|_| {
            (0..ncols_total)
                .map(|i| i.to_string().into_boxed_str())
                .collect()
        })
    } else {
        (0..ncols_total)
            .map(|i| i.to_string().into_boxed_str())
            .collect()
    };

    // Resolve column selection
    let selected = resolve_columns(column_indices, column_names, &all_col_names, &[0, 1])?;
    let sel_names: Vec<Box<str>> = selected
        .iter()
        .map(|&i| {
            all_col_names
                .get(i)
                .cloned()
                .unwrap_or_else(|| i.to_string().into_boxed_str())
        })
        .collect();

    // Build the matrix
    let mat = select_columns_to_dmatrix(&flat, nrows, ncols_total, &selected)?;

    // Read row names
    let row_names = if let Some(rn_path) = row_names_array {
        if row_names_10x {
            let (id_data, id_shape) = read_zarr_flat_u32(store.clone(), rn_path)?;
            anyhow::ensure!(
                id_shape[0] as usize == nrows,
                "row names array rows ({}) != data rows ({})",
                id_shape[0],
                nrows
            );
            parse_10x_cell_id_flat(&id_data, nrows)?
        } else {
            read_zarr_strings(store.clone(), rn_path)?
        }
    } else {
        (0..nrows).map(|i| i.to_string().into_boxed_str()).collect()
    };

    info!(
        "Read {} x {} from zarr {}{}: {:?}",
        nrows,
        selected.len(),
        file_path,
        data_array,
        sel_names
    );

    Ok(MatWithNames {
        rows: row_names,
        cols: sel_names,
        mat,
    })
}

// ── Xenium convenience ──────────────────────────────────────────────────

/// Read cell coordinates from a Xenium-style zarr file.
///
/// Shorthand for [`read_zarr_matrix`] with Xenium defaults:
///   - data array: `/cell_summary`
///   - column names attribute: `column_names`
///   - row names: `/cell_id` (10x encoded)
pub fn read_zarr_coordinates(
    file_path: &str,
    column_indices: &[usize],
    column_names: &[Box<str>],
) -> anyhow::Result<MatWithNames<DMatrix<f32>>> {
    read_zarr_matrix(
        file_path,
        "/cell_summary",
        Some("column_names"),
        Some("/cell_id"),
        true,
        column_indices,
        column_names,
    )
}

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

    fn xenium_path() -> Option<std::path::PathBuf> {
        let p = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
            .parent()
            .unwrap()
            .join("docs/temp/cells.zarr.zip");
        p.exists().then_some(p)
    }

    #[test]
    fn test_parse_10x_cell_id_flat() {
        let data = vec![16844u32, 1, 22527, 1];
        let ids = parse_10x_cell_id_flat(&data, 2).unwrap();
        assert_eq!(ids.len(), 2);
        assert!(ids[0].ends_with("-1"), "got: {}", ids[0]);
        assert!(ids[1].ends_with("-1"), "got: {}", ids[1]);
        assert_eq!(ids[0].len(), 10);
    }

    #[test]
    fn test_parse_10x_cell_id_flat_mismatch() {
        assert!(parse_10x_cell_id_flat(&[1, 2, 3], 2).is_err());
    }

    #[test]
    fn test_read_zarr_coordinates_by_name() {
        let Some(p) = xenium_path() else { return };
        let result = read_zarr_coordinates(
            p.to_str().unwrap(),
            &[],
            &["cell_centroid_x".into(), "cell_centroid_y".into()],
        )
        .unwrap();

        assert_eq!(result.mat.ncols(), 2);
        assert!(result.mat.nrows() > 0);
        assert_eq!(result.rows.len(), result.mat.nrows());
        assert_eq!(result.cols[0].as_ref(), "cell_centroid_x");
        assert_eq!(result.cols[1].as_ref(), "cell_centroid_y");
        assert!(result.mat.min() >= 0.0);
    }

    #[test]
    fn test_read_zarr_coordinates_by_index() {
        let Some(p) = xenium_path() else { return };
        let result = read_zarr_coordinates(p.to_str().unwrap(), &[0, 1], &[]).unwrap();
        assert_eq!(result.mat.ncols(), 2);
        assert_eq!(result.cols[0].as_ref(), "cell_centroid_x");
    }

    #[test]
    fn test_read_zarr_coordinates_default() {
        let Some(p) = xenium_path() else { return };
        let result = read_zarr_coordinates(p.to_str().unwrap(), &[], &[]).unwrap();
        assert_eq!(result.mat.ncols(), 2);
    }

    #[test]
    fn test_read_zarr_coordinates_bad_column() {
        let Some(p) = xenium_path() else { return };
        let result = read_zarr_coordinates(
            p.to_str().unwrap(),
            &[],
            &["nonexistent".to_string().into_boxed_str()],
        );
        assert!(result.is_err());
    }

    #[test]
    fn test_read_zarr_coordinates_mixed_defaults() {
        // Simulates pinto's generous defaults: Visium names that won't match
        // plus Xenium names that will — only matching names should be kept.
        let Some(p) = xenium_path() else { return };
        let result = read_zarr_coordinates(
            p.to_str().unwrap(),
            &[],
            &[
                "pxl_row_in_fullres".into(),
                "pxl_col_in_fullres".into(),
                "cell_centroid_x".into(),
                "cell_centroid_y".into(),
            ],
        )
        .unwrap();
        assert_eq!(result.mat.ncols(), 2);
        assert_eq!(result.cols[0].as_ref(), "cell_centroid_x");
        assert_eq!(result.cols[1].as_ref(), "cell_centroid_y");
        assert!(result.mat.nrows() > 0);
    }

    #[test]
    fn test_read_zarr_array_attr() {
        let Some(p) = xenium_path() else { return };
        let store = open_zarr_store(p.to_str().unwrap()).unwrap();
        let names: Vec<Box<str>> =
            read_zarr_array_attr(store, "/cell_summary", "column_names").unwrap();
        assert!(!names.is_empty());
        assert!(names.contains(&"cell_centroid_x".to_string().into_boxed_str()));
    }

    #[test]
    fn test_read_zarr_matrix_generic() {
        let Some(p) = xenium_path() else { return };
        // Read cell_summary selecting columns 2,3 (cell_area, nucleus_centroid_x)
        let result = read_zarr_matrix(
            p.to_str().unwrap(),
            "/cell_summary",
            Some("column_names"),
            Some("/cell_id"),
            true,
            &[2, 3],
            &[],
        )
        .unwrap();
        assert_eq!(result.mat.ncols(), 2);
        assert_eq!(result.cols[0].as_ref(), "cell_area");
        assert_eq!(result.cols[1].as_ref(), "nucleus_centroid_x");
    }

    /// (entry name, contents) pairs for a test zip.
    type Entries<'a> = &'a [(&'a str, &'a [u8])];

    /// A zip `name` in `dir` of the given entries, written as is.
    fn zip_of(dir: &tempfile::TempDir, name: &str, entries: Entries) -> std::path::PathBuf {
        use std::io::Write;
        let path = dir.path().join(name);
        let mut w = zip::ZipWriter::new(std::fs::File::create(&path).unwrap());
        let opts = zip::write::SimpleFileOptions::default();
        for (entry, bytes) in entries {
            w.start_file(*entry, opts).unwrap();
            w.write_all(bytes).unwrap();
        }
        w.finish().unwrap();
        path
    }

    fn extract(zip: &std::path::Path, out: &std::path::Path) -> anyhow::Result<()> {
        extract_zarr_zip(zip.to_str().unwrap(), out.to_str().unwrap())
    }

    /// An entry that climbs out of the target (`run/../../x`, behind the
    /// store's own prefix) refuses the whole archive before anything is
    /// written.
    #[test]
    fn extracting_refuses_an_entry_that_leaves_the_target() {
        let dir = tempfile::tempdir().unwrap();
        let zip = zip_of(
            &dir,
            "run.zarr.zip",
            &[
                ("run/zarr.json", b"{}"),
                ("run/../../escape.txt", b"escaped"),
            ],
        );
        let target = dir.path().join("out").join("run.zarr");
        let err = extract(&zip, &target).unwrap_err().to_string();
        assert!(err.contains("no safe path"), "{err}");
        assert!(
            !dir.path().join("out").join("escape.txt").exists()
                && !dir.path().join("escape.txt").exists(),
            "an entry was written outside the target"
        );
        assert!(
            !target.join("zarr.json").exists(),
            "a refused archive left files behind"
        );
    }

    /// A store zipped by `zip_dir_as` behind its `run/` root, or by `zip_dir`
    /// behind the older `run.zarr/` one, extracts back to the same files.
    #[test]
    fn a_zipped_store_extracts_back_to_its_files() {
        for root in [Some("run"), None] {
            let dir = tempfile::tempdir().unwrap();
            let src = dir.path().join("run.zarr");
            std::fs::create_dir_all(src.join("X/c/0")).unwrap();
            std::fs::write(src.join("zarr.json"), b"{}").unwrap();
            std::fs::write(src.join("X/c/0/1"), b"chunk").unwrap();
            let zip = dir.path().join("run.zarr.zip");
            legume_numeric::matrix::common_io::zip_dir_as(
                src.to_str().unwrap(),
                zip.to_str().unwrap(),
                root,
            )
            .unwrap();
            let out = dir.path().join("back.zarr");
            extract(&zip, &out).unwrap();
            assert_eq!(
                std::fs::read(out.join("zarr.json")).unwrap(),
                b"{}",
                "{root:?}"
            );
            assert_eq!(
                std::fs::read(out.join("X/c/0/1")).unwrap(),
                b"chunk",
                "{root:?}"
            );
        }
    }

    /// The store's prefix is read from the entries: the archive's single
    /// top-level directory, so a renamed archive keeps it and a store without
    /// a root group (a nested array) still gets its root, with loose files
    /// and an OS's metadata beside it ignored and left out of the extracted
    /// store; beside other directories (a backup), the one named like the
    /// archive; and none for a store at the archive's root.
    #[test]
    fn the_store_prefix_comes_from_the_entries() {
        let dir = tempfile::tempdir().unwrap();
        let cases: [(&str, Entries, &str); 7] = [
            (
                "nested.zarr.zip",
                &[("run/m/X/zarr.json", b"{}"), ("run/m/X/c/0", b"chunk")],
                "run/",
            ),
            (
                "run.zarr.zip",
                &[("run/zarr.json", b"{}"), ("run.zarr.bak/zarr.json", b"{}")],
                "run/",
            ),
            (
                "renamed.zarr.zip",
                &[("run/zarr.json", b"{}"), ("run/X/zarr.json", b"{}")],
                "run/",
            ),
            (
                "no_group.zarr.zip",
                &[("run/X/zarr.json", b"{}"), ("run/rows/zarr.json", b"{}")],
                "run/",
            ),
            (
                "stray.zarr.zip",
                &[
                    ("README.txt", b""),
                    ("__MACOSX/run/._zarr.json", b""),
                    ("run/zarr.json", b"{}"),
                ],
                "run/",
            ),
            (
                "v2.zarr.zip",
                &[("old.zarr/.zgroup", b"{}"), ("old.zarr/X/.zarray", b"{}")],
                "old.zarr/",
            ),
            (
                "at_root.zarr.zip",
                &[("zarr.json", b"{}"), ("X/zarr.json", b"{}")],
                "",
            ),
        ];
        for (name, entries, prefix) in cases {
            let zip = zip_of(&dir, name, entries);
            assert_eq!(&*detect_zip_zarr_prefix(&zip).unwrap(), prefix, "{name}");
        }
        for (name, kept, dropped) in [
            ("stray.zarr.zip", "zarr.json", "README.txt"),
            ("no_group.zarr.zip", "rows/zarr.json", "run"),
        ] {
            let out = dir.path().join(name).with_extension("out");
            extract(&dir.path().join(name), &out).unwrap();
            assert!(
                out.join(kept).exists() && !out.join(dropped).exists(),
                "{name}"
            );
        }
    }
}