dspk-synapse 0.0.10

//! Phy / Kilosort 1–4 sorting folders: one row per spike (time, cluster, template, amplitude,
//! position), dense templates, template similarity, optional PC and template features, the
//! cluster tables (`cluster_group.tsv`, `cluster_KSLabel.tsv`, …, `cluster_info.tsv`) and
//! `params.py`.
//!
//! [`PhyFolder::read`] reads a folder as written; [`PhyFolder::write`] writes a whole folder;
//! [`PhyFolder::write_curation`] writes back only what curation changes (`spike_clusters.npy`,
//! `cluster_group.tsv`, `cluster_info.tsv`), keeping a `.bak` of each the first time. Values are
//! stored as found: nothing is computed (a missing `similar_templates.npy` stays `None`).

use std::collections::BTreeMap;
use std::fs::File;
use std::io::{BufRead, BufReader, BufWriter, Write};
use std::path::{Path, PathBuf};

use dspk_core::{DspError, DspResult, SampleFormat};

use dspk_io::container::npy::{read_npy, write_npy, NpyArray, NpyElement};
use crate::io::templates::DenseTemplates;

/// Cluster (and template) ids as phy stores them.
pub type ClusterId = u32;

/// `dtype` written to `params.py` when none is known.
const DEFAULT_DAT_DTYPE: &str = "int16";

fn io_err(path: &Path) -> impl Fn(std::io::Error) -> DspError + '_ {
    move |e| DspError::Io(format!("{}: {e}", path.display()))
}

/// `params.py` of a phy folder.
#[derive(Debug, Clone, PartialEq, Default)]
pub struct PhyParams {
    pub sample_rate: f64,
    /// As written (may be a path from another machine; see [`resolve_dat_path`]).
    pub dat_path: Option<String>,
    pub n_channels_dat: usize,
    pub dtype: Option<String>,
    pub offset: u64,
    pub hp_filtered: bool,
}

impl PhyParams {
    /// Reads `params.py` (missing file or keys: defaults).
    pub fn read(path: &Path) -> Self {
        let mut p = Self::default();
        let Ok(f) = File::open(path) else { return p };
        for line in BufReader::new(f).lines().map_while(Result::ok) {
            let Some((k, v)) = line.split_once('=') else { continue };
            let v = v.split('#').next().unwrap_or(v).trim();
            // `'a.bin'`, `r"C:\a.bin"` or `['a.bin']`
            let text = || {
                let v = v.strip_prefix('[').and_then(|s| s.strip_suffix(']')).unwrap_or(v).trim();
                let v = v.strip_prefix('r').filter(|s| s.starts_with(['\'', '"'])).unwrap_or(v);
                v.trim_matches(|c| c == '\'' || c == '"').to_string()
            };
            match k.trim() {
                "sample_rate" => p.sample_rate = v.parse().unwrap_or(0.0),
                "n_channels_dat" => p.n_channels_dat = v.parse().unwrap_or(0),
                "offset" => p.offset = v.parse().unwrap_or(0),
                "hp_filtered" => p.hp_filtered = v == "True",
                "dtype" => p.dtype = Some(text()).filter(|s| !s.is_empty()),
                "dat_path" => p.dat_path = Some(text()).filter(|s| !s.is_empty()),
                _ => {}
            }
        }
        p
    }

    /// Writes `params.py` (`dtype` defaults to `int16`).
    pub fn write(&self, path: &Path) -> DspResult<()> {
        let io = io_err(path);
        let mut f = BufWriter::new(File::create(path).map_err(&io)?);
        writeln!(f, "# Generated by dspk").map_err(&io)?;
        if let Some(dat_path) = &self.dat_path {
            writeln!(f, "dat_path = r'{dat_path}'").map_err(&io)?;
        }
        writeln!(f, "n_channels_dat = {}", self.n_channels_dat).map_err(&io)?;
        writeln!(f, "dtype = '{}'", self.dtype.as_deref().unwrap_or(DEFAULT_DAT_DTYPE)).map_err(&io)?;
        writeln!(f, "offset = {}", self.offset).map_err(&io)?;
        writeln!(f, "sample_rate = {}", self.sample_rate).map_err(&io)?;
        writeln!(f, "hp_filtered = {}", if self.hp_filtered { "True" } else { "False" }).map_err(&io)?;
        f.flush().map_err(&io)
    }
}

fn is_recording_target(p: &Path) -> bool {
    p.is_file() || (p.is_dir() && (p.join("acquisition").is_dir() || p.join("zarr.json").is_file()))
}

/// The recording `dat_path` names: as written, relative to the folder, else a file/store of that
/// name in the folder or up to three folders above it (paths written on another machine).
pub fn resolve_dat_path(sorting_dir: &Path, dat_path: &str) -> Option<PathBuf> {
    let direct = PathBuf::from(dat_path);
    if is_recording_target(&direct) {
        return Some(direct);
    }
    let relative = sorting_dir.join(dat_path);
    if is_recording_target(&relative) {
        return Some(relative);
    }
    let name = dat_path.rsplit(['/', '\\']).next().filter(|n| !n.is_empty())?;
    sorting_dir.ancestors().take(4).map(|d| d.join(name)).find(|p| is_recording_target(p))
}

/// The cluster tables of a phy folder, by cluster id (values as written).
#[derive(Debug, Clone, PartialEq, Default)]
pub struct ClusterTables {
    /// `cluster_group.tsv`: good / mua / noise / unsorted (curated).
    pub group: BTreeMap<ClusterId, String>,
    /// `cluster_KSLabel.tsv`: the sorter's own label.
    pub ks_label: BTreeMap<ClusterId, String>,
    /// `cluster_Amplitude.tsv`, `cluster_ContamPct.tsv`.
    pub amplitude: BTreeMap<ClusterId, f32>,
    pub contam_pct: BTreeMap<ClusterId, f32>,
    /// `cluster_info.tsv`: every column of every row, and the column order.
    pub info: BTreeMap<ClusterId, BTreeMap<String, String>>,
    pub info_columns: Vec<String>,
}

impl ClusterTables {
    /// Reads every table present in `dir`.
    pub fn read(dir: &Path) -> Self {
        let number = |m: BTreeMap<ClusterId, String>| m.into_iter().filter_map(|(k, v)| Some((k, v.parse().ok()?))).collect();
        let (info, info_columns) = read_table(&dir.join("cluster_info.tsv"));
        Self {
            group: read_two_columns(&dir.join("cluster_group.tsv")),
            ks_label: read_two_columns(&dir.join("cluster_KSLabel.tsv")),
            amplitude: number(read_two_columns(&dir.join("cluster_Amplitude.tsv"))),
            contam_pct: number(read_two_columns(&dir.join("cluster_ContamPct.tsv"))),
            info,
            info_columns,
        }
    }

    /// Writes every non-empty table into `dir`.
    pub fn write(&self, dir: &Path) -> DspResult<()> {
        let text = |m: &BTreeMap<ClusterId, f32>| m.iter().map(|(&k, v)| (k, v.to_string())).collect::<BTreeMap<_, _>>();
        let two: [(&str, &str, BTreeMap<ClusterId, String>); 4] = [
            ("cluster_group.tsv", "group", self.group.clone()),
            ("cluster_KSLabel.tsv", "KSLabel", self.ks_label.clone()),
            ("cluster_Amplitude.tsv", "Amplitude", text(&self.amplitude)),
            ("cluster_ContamPct.tsv", "ContamPct", text(&self.contam_pct)),
        ];
        for (file, column, rows) in &two {
            if !rows.is_empty() {
                write_two_columns(&dir.join(file), column, rows)?;
            }
        }
        if !self.info_columns.is_empty() {
            write_table(&dir.join("cluster_info.tsv"), &self.info_columns, &self.info)?;
        }
        Ok(())
    }
}

/// `cluster_id \t value` rows (header skipped).
fn read_two_columns(path: &Path) -> BTreeMap<ClusterId, String> {
    let Ok(f) = File::open(path) else { return BTreeMap::new() };
    BufReader::new(f)
        .lines()
        .map_while(Result::ok)
        .filter_map(|line| {
            let (id, value) = line.split_once('\t')?;
            Some((id.trim().parse().ok()?, value.trim().to_string()))
        })
        .collect()
}

fn write_two_columns(path: &Path, column: &str, rows: &BTreeMap<ClusterId, String>) -> DspResult<()> {
    let io = io_err(path);
    let mut f = BufWriter::new(File::create(path).map_err(&io)?);
    writeln!(f, "cluster_id\t{column}").map_err(&io)?;
    for (id, v) in rows {
        writeln!(f, "{id}\t{v}").map_err(&io)?;
    }
    f.flush().map_err(&io)
}

/// A tab-separated table keyed by its `cluster_id` (or `id`) column, and its column names.
fn read_table(path: &Path) -> (BTreeMap<ClusterId, BTreeMap<String, String>>, Vec<String>) {
    let Ok(f) = File::open(path) else { return Default::default() };
    let mut lines = BufReader::new(f).lines().map_while(Result::ok);
    let Some(header) = lines.next() else { return Default::default() };
    let columns: Vec<String> = header.split('\t').map(|s| s.trim().to_string()).collect();
    let Some(id_col) = columns.iter().position(|h| h == "cluster_id" || h == "id") else { return Default::default() };
    let rows = lines
        .filter_map(|line| {
            let cells: Vec<&str> = line.split('\t').collect();
            let id = cells.get(id_col)?.trim().parse().ok()?;
            let row = columns.iter().zip(&cells).filter(|(_, v)| !v.trim().is_empty()).map(|(k, v)| (k.clone(), v.trim().to_string())).collect();
            Some((id, row))
        })
        .collect();
    (rows, columns)
}

/// `columns` as the header, then one row per entry of `rows` (its id in the `cluster_id` / `id`
/// column, missing cells empty).
fn write_table(path: &Path, columns: &[String], rows: &BTreeMap<ClusterId, BTreeMap<String, String>>) -> DspResult<()> {
    let io = io_err(path);
    let mut f = BufWriter::new(File::create(path).map_err(&io)?);
    writeln!(f, "{}", columns.join("\t")).map_err(&io)?;
    for (id, row) in rows {
        let cells: Vec<String> = columns
            .iter()
            .map(|c| if c == "cluster_id" || c == "id" { id.to_string() } else { row.get(c).cloned().unwrap_or_default() })
            .collect();
        writeln!(f, "{}", cells.join("\t")).map_err(&io)?;
    }
    f.flush().map_err(&io)
}

/// An optional array of the folder, read as `T` (absent or unreadable: `None`, with a warning for
/// unreadable ones).
fn optional<T: NpyElement>(dir: &Path, name: &str) -> Option<NpyArray<T>> {
    let path = dir.join(name);
    if !path.exists() {
        return None;
    }
    read_npy::<T>(&path).map_err(|e| tracing::warn!("ignoring {name}: {e}")).ok()
}

/// Rows of an `[n, ≥2]` array as `[x, y]` (first two columns).
fn xy_rows(a: &NpyArray<f32>) -> Option<Vec<[f32; 2]>> {
    let cols = *a.shape.get(1)?;
    (a.shape.len() == 2 && cols >= 2).then(|| a.data.chunks_exact(cols).map(|r| [r[0], r[1]]).collect())
}

/// A phy / Kilosort folder, spike by spike. See the module docs.
#[derive(Debug, Clone, PartialEq, Default)]
pub struct PhyFolder {
    /// Folder it was read from.
    pub folder: Option<PathBuf>,
    pub params: PhyParams,
    pub spike_times: Vec<u64>,
    pub spike_clusters: Vec<ClusterId>,
    /// Template each spike was matched with (its cluster when the sorter saved none).
    pub spike_templates: Vec<ClusterId>,
    pub amplitudes: Vec<f32>,
    /// `[x, y]` µm per spike (empty when not saved).
    pub spike_positions: Vec<[f32; 2]>,
    /// Recording channel of each sorted channel.
    pub channel_map: Vec<usize>,
    /// `[x, y]` µm per sorted channel.
    pub channel_positions: Vec<[f32; 2]>,
    pub channel_shanks: Vec<usize>,
    pub templates: Option<DenseTemplates>,
    pub templates_std: Option<Vec<f32>>,
    pub templates_se: Option<Vec<f32>>,
    /// `[templates, templates]` similarity (`similar_templates.npy`).
    pub similar_templates: Option<Vec<f32>>,
    /// `pc_features.npy` `[spikes, pcs, channels]` and `pc_feature_ind.npy` `[templates, channels]`.
    pub pc_features: Option<(Vec<f32>, [usize; 3])>,
    pub pc_feature_ind: Option<(Vec<usize>, [usize; 2])>,
    /// `template_features.npy` `[spikes, k]` and `template_feature_ind.npy` `[templates, k]`.
    pub template_features: Option<(Vec<f32>, [usize; 2])>,
    pub template_feature_ind: Option<(Vec<usize>, [usize; 2])>,
    pub tables: ClusterTables,
}

impl PhyFolder {
    /// Whether `dir` looks like a phy / Kilosort folder.
    pub fn is_phy_folder(dir: &Path) -> bool {
        dir.join("spike_times.npy").is_file()
    }

    /// Reads a phy / Kilosort 1–4 folder (`spike_times.npy` and `spike_clusters.npy` required).
    pub fn read(dir: &Path) -> DspResult<Self> {
        let need = |name: &str| {
            let p = dir.join(name);
            if p.exists() { Ok(p) } else { Err(DspError::UnsupportedFormat(format!("{} has no {name}: not a phy / Kilosort folder", dir.display()))) }
        };
        let spike_times = read_npy::<u64>(&need("spike_times.npy")?)?.data;
        let clusters = read_npy::<i64>(&need("spike_clusters.npy")?)?.data;
        if let Some(bad) = clusters.iter().find(|&&c| c < 0) {
            return Err(DspError::UnsupportedFormat(format!("{}: negative cluster id {bad} in spike_clusters.npy", dir.display())));
        }
        let n = spike_times.len().min(clusters.len());
        let spike_clusters: Vec<ClusterId> = clusters[..n].iter().map(|&c| c as ClusterId).collect();
        let spike_times = spike_times[..n].to_vec();
        let per_spike = |v: Option<Vec<ClusterId>>| v.filter(|v| v.len() >= n).map(|mut v| { v.truncate(n); v });
        let spike_templates = per_spike(optional::<ClusterId>(dir, "spike_templates.npy").map(|a| a.data)).unwrap_or_else(|| spike_clusters.clone());
        let amplitudes = optional::<f32>(dir, "amplitudes.npy").map(|a| a.data).filter(|a| a.len() >= n).map(|mut a| { a.truncate(n); a }).unwrap_or_default();
        let spike_positions = optional::<f32>(dir, "spike_positions.npy").as_ref().and_then(xy_rows).filter(|p| p.len() >= n).map(|mut p| { p.truncate(n); p }).unwrap_or_default();

        let params = PhyParams::read(&dir.join("params.py"));
        let channel_positions = optional::<f32>(dir, "channel_positions.npy").as_ref().and_then(xy_rows).unwrap_or_default();
        let channel_map = optional::<usize>(dir, "channel_map.npy").map(|a| a.data).unwrap_or_default();
        let channel_shanks = optional::<usize>(dir, "channel_shanks.npy").map(|a| a.data).unwrap_or_default();
        let channels = channel_positions.len().max(channel_map.len());

        let templates = optional::<f32>(dir, "templates.npy").filter(|a| a.shape.len() == 3).map(|a| {
            let (count, samples, stored) = (a.shape[0], a.shape[1], a.shape[2]);
            // Sparse templates (Kilosort 1–3): `templates_ind.npy` names each stored column's channel
            let ind = optional::<usize>(dir, "templates_ind.npy").filter(|i| i.shape == [count, stored]);
            match ind {
                Some(ind) if stored < channels => DenseTemplates::from_sparse(&a.data, count, samples, stored, &ind.data, channels),
                _ => DenseTemplates { data: a.data, count, samples, channels: stored },
            }
        });
        let same_shape = |name: &str| optional::<f32>(dir, name).filter(|a| templates.as_ref().is_some_and(|t| a.data.len() == t.data.len())).map(|a| a.data);
        let (templates_std, templates_se) = (same_shape("templates_std.npy"), same_shape("templates_se.npy"));
        let similar_templates = optional::<f32>(dir, "similar_templates.npy")
            .filter(|a| templates.as_ref().is_some_and(|t| a.shape == [t.count, t.count]))
            .map(|a| a.data);
        let three = |a: NpyArray<f32>| (a.shape.len() == 3).then(|| (a.data, [a.shape[0], a.shape[1], a.shape[2]]));
        let two = |a: NpyArray<f32>| (a.shape.len() == 2).then(|| (a.data, [a.shape[0], a.shape[1]]));
        let two_ind = |a: NpyArray<usize>| (a.shape.len() == 2).then(|| (a.data, [a.shape[0], a.shape[1]]));

        let mut folder = Self {
            folder: Some(dir.to_path_buf()),
            params,
            spike_times,
            spike_clusters,
            spike_templates,
            amplitudes,
            spike_positions,
            channel_map,
            channel_positions,
            channel_shanks,
            templates,
            templates_std,
            templates_se,
            similar_templates,
            pc_features: optional::<f32>(dir, "pc_features.npy").and_then(three),
            pc_feature_ind: optional::<usize>(dir, "pc_feature_ind.npy").and_then(two_ind),
            template_features: optional::<f32>(dir, "template_features.npy").and_then(two),
            template_feature_ind: optional::<usize>(dir, "template_feature_ind.npy").and_then(two_ind),
            tables: ClusterTables::read(dir),
        };
        folder.fill_channel_defaults();
        Ok(folder)
    }

    /// Identity channel map and shank 0 where they are missing.
    pub fn fill_channel_defaults(&mut self) {
        let channels = self.channels();
        if self.channel_map.len() != channels {
            self.channel_map = (0..channels).collect();
        }
        if self.channel_shanks.len() != channels {
            self.channel_shanks = vec![0; channels];
        }
    }

    /// Writes the whole folder into `dir` (created if needed): every array present and every
    /// non-empty table, and `params.py`.
    pub fn write(&self, dir: &Path) -> DspResult<()> {
        std::fs::create_dir_all(dir).map_err(io_err(dir))?;
        let n = self.spike_times.len();
        let ids = |v: &[ClusterId]| v.iter().map(|&c| c as i32).collect::<Vec<i32>>();
        write_npy(&dir.join("spike_times.npy"), &self.spike_times, &[n])?;
        write_npy(&dir.join("spike_clusters.npy"), &ids(&self.spike_clusters), &[n])?;
        if self.spike_templates.len() == n {
            write_npy(&dir.join("spike_templates.npy"), &ids(&self.spike_templates), &[n])?;
        }
        if self.amplitudes.len() == n {
            write_npy(&dir.join("amplitudes.npy"), &self.amplitudes, &[n])?;
        }
        if self.spike_positions.len() == n && n > 0 {
            let flat: Vec<f32> = self.spike_positions.iter().flatten().copied().collect();
            write_npy(&dir.join("spike_positions.npy"), &flat, &[n, 2])?;
        }

        if let Some(t) = &self.templates {
            let shape = [t.count, t.samples, t.channels];
            write_npy(&dir.join("templates.npy"), &t.data, &shape)?;
            for (name, buf) in [("templates_std.npy", &self.templates_std), ("templates_se.npy", &self.templates_se)] {
                if let Some(b) = buf.as_ref().filter(|b| b.len() == t.data.len()) {
                    write_npy(&dir.join(name), b, &shape)?;
                }
            }
            if let Some(s) = self.similar_templates.as_ref().filter(|s| s.len() == t.count * t.count) {
                write_npy(&dir.join("similar_templates.npy"), s, &[t.count, t.count])?;
            }
        }
        if let Some((d, s)) = &self.pc_features {
            write_npy(&dir.join("pc_features.npy"), d, s)?;
        }
        if let Some((d, s)) = &self.template_features {
            write_npy(&dir.join("template_features.npy"), d, s)?;
        }
        let as_i32 = |v: &[usize]| v.iter().map(|&c| c as i32).collect::<Vec<i32>>();
        if let Some((d, s)) = &self.pc_feature_ind {
            write_npy(&dir.join("pc_feature_ind.npy"), &as_i32(d), s)?;
        }
        if let Some((d, s)) = &self.template_feature_ind {
            write_npy(&dir.join("template_feature_ind.npy"), &as_i32(d), s)?;
        }

        if !self.channel_map.is_empty() {
            write_npy(&dir.join("channel_map.npy"), &as_i32(&self.channel_map), &[self.channel_map.len()])?;
        }
        if !self.channel_shanks.is_empty() {
            write_npy(&dir.join("channel_shanks.npy"), &as_i32(&self.channel_shanks), &[self.channel_shanks.len()])?;
        }
        if !self.channel_positions.is_empty() {
            let flat: Vec<f32> = self.channel_positions.iter().flatten().copied().collect();
            write_npy(&dir.join("channel_positions.npy"), &flat, &[self.channel_positions.len(), 2])?;
        }

        self.tables.write(dir)?;
        self.params.write(&dir.join("params.py"))
    }

    /// Writes curation results into `dir`: `spike_clusters.npy` (from `spike_clusters`),
    /// `cluster_group.tsv` (from `groups`) and `cluster_info.tsv` (`columns`, then one row per
    /// entry of `rows`). The first time a file is replaced its original is kept as `<name>.bak`;
    /// nothing else in the folder is touched.
    pub fn write_curation(
        dir: &Path,
        spike_clusters: &[ClusterId],
        groups: &BTreeMap<ClusterId, String>,
        columns: &[String],
        rows: &BTreeMap<ClusterId, BTreeMap<String, String>>,
    ) -> DspResult<()> {
        let io = io_err(dir);
        std::fs::create_dir_all(dir).map_err(&io)?;
        for name in ["spike_clusters.npy", "cluster_group.tsv", "cluster_info.tsv"] {
            let (file, bak) = (dir.join(name), dir.join(format!("{name}.bak")));
            if file.exists() && !bak.exists() {
                std::fs::copy(&file, &bak).map_err(&io)?;
            }
        }
        let clusters: Vec<i32> = spike_clusters.iter().map(|&c| c as i32).collect();
        write_npy(&dir.join("spike_clusters.npy"), &clusters, &[clusters.len()])?;
        write_two_columns(&dir.join("cluster_group.tsv"), "group", groups)?;
        write_table(&dir.join("cluster_info.tsv"), columns, rows)
    }

    /// Sorted channels (templates, positions or channel map, whichever says).
    pub fn channels(&self) -> usize {
        self.templates.as_ref().map_or(0, |t| t.channels).max(self.channel_positions.len()).max(self.channel_map.len())
    }

    pub fn sample_rate(&self) -> f64 {
        self.params.sample_rate
    }

    /// The recording `params.py` names, if it can be found.
    pub fn recording_path(&self) -> Option<PathBuf> {
        resolve_dat_path(self.folder.as_deref()?, self.params.dat_path.as_deref()?)
    }

    /// Samples per channel of the recording `params.py` names, when it is a flat binary file found
    /// on disk: `(file size − offset) / (n_channels_dat · dtype size)`. `None` when the file is not
    /// found or is not a plain file, or `dtype` / `n_channels_dat` are missing.
    pub fn recording_samples(&self) -> Option<u64> {
        let path = self.recording_path().filter(|p| p.is_file())?;
        let format = SampleFormat::parse(self.params.dtype.as_deref()?)?;
        let frame = (self.params.n_channels_dat * format.bytes()) as u64;
        let bytes = std::fs::metadata(path).ok()?.len().checked_sub(self.params.offset)?;
        (frame > 0).then(|| bytes / frame)
    }

    /// Similarity of templates `a` and `b` (0 without templates or similarity).
    pub fn similarity(&self, a: usize, b: usize) -> f32 {
        let n = self.templates.as_ref().map_or(0, |t| t.count);
        match &self.similar_templates {
            Some(s) if a < n && b < n => s[a * n + b],
            _ => 0.0,
        }
    }
}

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

    fn scratch(name: &str) -> PathBuf {
        let dir = std::env::temp_dir().join(format!("dsp_phyfolder_{name}_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();
        dir
    }

    /// A Kilosort4-like folder: int64 times, 2-column positions, float64 channel positions.
    pub(crate) fn write_kilosort4_like(dir: &Path) {
        write_npy(&dir.join("spike_times.npy"), &[10i64, 20, 30, 40, 55], &[5]).unwrap();
        write_npy(&dir.join("spike_clusters.npy"), &[0i32, 1, 0, 1, 1], &[5]).unwrap();
        write_npy(&dir.join("spike_templates.npy"), &[0u32, 1, 0, 1, 1], &[5]).unwrap();
        write_npy(&dir.join("amplitudes.npy"), &[1.0f64, 2.0, 1.5, 2.5, 3.0], &[5]).unwrap();
        write_npy(&dir.join("spike_positions.npy"), &[0.0f32, 10.0, 0.0, 30.0, 0.0, 10.0, 0.0, 30.0, 0.0, 30.0], &[5, 2]).unwrap();
        write_npy(&dir.join("channel_positions.npy"), &[0.0f64, 0.0, 0.0, 20.0, 0.0, 40.0], &[3, 2]).unwrap();
        write_npy(&dir.join("channel_map.npy"), &[0i64, 1, 2], &[3]).unwrap();
        // Template 0 peaks on channel 2, template 1 on channel 0; [template, sample, channel] 2 × 4 × 3
        let at = |t: usize, s: usize, c: usize| (t * 4 + s) * 3 + c;
        let mut t = vec![0.0f32; 2 * 4 * 3];
        t[at(0, 1, 2)] = -50.0;
        t[at(0, 2, 2)] = 10.0;
        t[at(1, 1, 0)] = -80.0;
        write_npy(&dir.join("templates.npy"), &t, &[2, 4, 3]).unwrap();
        std::fs::write(dir.join("params.py"), "dat_path = 'C:/Users/x/rec.bin'\nn_channels_dat = 3\ndtype = 'int16'\noffset = 0\nsample_rate = 30000.0\nhp_filtered = False\n").unwrap();
        std::fs::write(dir.join("cluster_KSLabel.tsv"), "cluster_id\tKSLabel\n0\tgood\n1\tmua\n").unwrap();
        std::fs::write(dir.join("cluster_ContamPct.tsv"), "cluster_id\tContamPct\n0\t1.5\n1\t20.0\n").unwrap();
        std::fs::write(dir.join("cluster_info.tsv"), "cluster_id\tch\tquality\n0\t2\tnice\n1\t0\t\n").unwrap();
        std::fs::write(dir.join("rec.bin"), [0u8; 4]).unwrap();
    }

    #[test]
    fn reads_a_kilosort4_like_folder() {
        let dir = scratch("read");
        write_kilosort4_like(&dir);
        let s = PhyFolder::read(&dir).unwrap();
        assert_eq!(s.spike_times, vec![10, 20, 30, 40, 55]);
        assert_eq!(s.spike_clusters, vec![0, 1, 0, 1, 1]);
        assert_eq!(s.amplitudes, vec![1.0, 2.0, 1.5, 2.5, 3.0]);
        assert_eq!(s.spike_positions[1], [0.0, 30.0], "two-column positions are kept");
        assert_eq!(s.channel_positions, vec![[0.0, 0.0], [0.0, 20.0], [0.0, 40.0]]);
        assert_eq!(s.params.sample_rate, 30_000.0);
        assert_eq!(s.params.dat_path.as_deref(), Some("C:/Users/x/rec.bin"));
        assert_eq!(s.recording_path(), Some(dir.join("rec.bin")), "a Windows path falls back to the file's name");
        assert_eq!(s.recording_samples(), Some(4 / (3 * 2)), "4 bytes of 3-channel int16");
        let t = s.templates.as_ref().unwrap();
        assert_eq!((t.best_channel(0), t.best_channel(1)), (2, 0));
        assert_eq!(t.top_channels(0, 2)[0], 2);
        assert_eq!(t.amplitude(0), 60.0);
        assert!(s.similar_templates.is_none(), "nothing computed on read");
        assert_eq!(s.tables.ks_label[&1], "mua");
        assert_eq!(s.tables.contam_pct[&1], 20.0);
        assert_eq!(s.tables.info[&0]["quality"], "nice");
        assert_eq!(s.tables.info_columns, vec!["cluster_id", "ch", "quality"]);
        std::fs::remove_dir_all(&dir).unwrap();
    }

    #[test]
    fn write_then_read_round_trips() {
        let (src, dst) = (scratch("rt_src"), scratch("rt_dst"));
        write_kilosort4_like(&src);
        let mut a = PhyFolder::read(&src).unwrap();
        a.similar_templates = Some(vec![1.0, 0.2, 0.2, 1.0]);
        a.write(&dst).unwrap();
        let mut b = PhyFolder::read(&dst).unwrap();
        b.folder = a.folder.clone();
        assert_eq!(a, b);
        std::fs::remove_dir_all(&src).unwrap();
        std::fs::remove_dir_all(&dst).unwrap();
    }

    #[test]
    fn write_curation_touches_three_files_and_keeps_backups() {
        let dir = scratch("curation");
        write_kilosort4_like(&dir);
        let groups = BTreeMap::from([(0, "good".to_string()), (2, "noise".to_string())]);
        let columns = vec!["cluster_id".to_string(), "group".to_string()];
        let rows = BTreeMap::from([(0, BTreeMap::from([("group".to_string(), "good".to_string())])), (2, BTreeMap::new())]);
        PhyFolder::write_curation(&dir, &[0, 2, 0, 2, 2], &groups, &columns, &rows).unwrap();
        let s = PhyFolder::read(&dir).unwrap();
        assert_eq!(s.spike_clusters, vec![0, 2, 0, 2, 2]);
        assert_eq!(s.tables.group[&2], "noise");
        assert_eq!(s.tables.info[&0]["group"], "good");
        assert!(dir.join("spike_clusters.npy.bak").exists() && dir.join("cluster_info.tsv.bak").exists());
        assert!(!dir.join("cluster_group.tsv.bak").exists(), "no original to keep");
        // A second save keeps the first backups (the originals)
        PhyFolder::write_curation(&dir, &[1, 1, 1, 1, 1], &groups, &columns, &rows).unwrap();
        let original = read_npy::<i32>(&dir.join("spike_clusters.npy.bak")).unwrap().data;
        assert_eq!(original, vec![0, 1, 0, 1, 1]);
        assert!(dir.join("amplitudes.npy").exists() && !dir.join("amplitudes.npy.bak").exists(), "other files untouched");
        std::fs::remove_dir_all(&dir).unwrap();
    }

    #[test]
    fn rejects_folders_without_spikes() {
        let dir = scratch("empty");
        assert!(PhyFolder::read(&dir).unwrap_err().to_string().contains("spike_times.npy"));
        std::fs::remove_dir_all(&dir).unwrap();
    }

    /// A real Kilosort4 output, when present: `DSP_KS4_FOLDER=<saved_results>`.
    #[test]
    #[ignore]
    fn real_kilosort4_folder() {
        let Some(dir) = std::env::var_os("DSP_KS4_FOLDER") else { return };
        let s = PhyFolder::read(Path::new(&dir)).unwrap();
        let t = s.templates.as_ref().unwrap();
        println!("{} spikes, {} templates × {} samples × {} channels, {} clusters, recording {:?}", s.spike_times.len(), t.count, t.samples, t.channels, s.tables.ks_label.len(), s.recording_path());
        assert_eq!(s.spike_times.len(), s.spike_clusters.len());
    }
}