slow5lib 0.1.0

/// SLOW5 text format parser and writer helpers.
///
/// Wire layout of the file header (all lines terminated by `\n`):
///
/// ```text
/// #slow5_version  {major}.{minor}.{patch}
/// #num_read_groups  {n}
/// @attr_name  val_rg0  [val_rg1 ...]      (one value per read group, tab-separated)
/// ...
/// #type0  type1  ...                       (primary + aux field types)
/// #col0   col1   ...                       (primary + aux field names)
/// ```
///
/// Data lines are tab-separated with the signal field containing comma-separated i16 values.
///
/// References: hasindu2008/slow5lib slow5.c `slow5_hdr_init` (text branch) and
/// `slow5_hdr_to_mem` (text branch).
use std::collections::BTreeSet;
use std::io::{BufRead, Write};

use crate::Result;
use crate::error::SlowError;
use crate::header::{Header, ReadGroup, RecordCompression, SignalCompression};
use crate::record::{RawRecord, Record};

/// Type header line for the eight primary fields.
const TYPE_LINE: &str = "#char*\tuint32_t\tdouble\tdouble\tdouble\tdouble\tuint64_t\tint16_t*\n";

/// Column name header line for the eight primary fields.
const COL_LINE: &str = "#read_id\tread_group\tdigitisation\toffset\trange\tsampling_rate\tlen_raw_signal\traw_signal\n";

/// Column indices for the eight mandatory primary fields within a data record line.
pub(crate) struct PrimaryColIndices {
    pub read_id: usize,
    pub read_group: usize,
    pub digitisation: usize,
    pub offset: usize,
    pub range: usize,
    pub sampling_rate: usize,
    pub len_raw_signal: usize,
    pub raw_signal: usize,
}

/// Parse the SLOW5 text header from `reader`.
///
/// On return, the reader is positioned at the first data record line.
#[allow(unused_assignments)] // col_idx initial None is always overwritten before it is read
pub(crate) fn parse_header(reader: &mut impl BufRead) -> Result<(Header, PrimaryColIndices)> {
    let mut version = (0u8, 0u8, 0u8);
    let mut num_read_groups: u32 = 1;
    // attribute key -> [val_rg0, val_rg1, ...]
    let mut attr_vals: std::collections::HashMap<String, Vec<String>> =
        std::collections::HashMap::new();
    let mut attr_keys_order: Vec<String> = Vec::new();
    let mut col_idx: Option<PrimaryColIndices> = None;
    // The type line (e.g. "#char*\tuint32_t\t...") precedes the column name line.
    let mut type_line: Option<String> = None;
    // The column name line (e.g. "#read_id\t...") -- saved for aux_meta parsing.
    let mut name_line: Option<String> = None;

    let mut line = String::new();

    loop {
        line.clear();
        let n = reader.read_line(&mut line).map_err(SlowError::Io)?;
        if n == 0 {
            return Err(SlowError::InvalidFormat(
                "unexpected EOF while parsing SLOW5 header".into(),
            ));
        }

        let t = line.trim_end();

        if t.is_empty() {
            continue;
        }

        if let Some(rest) = t.strip_prefix('@') {
            // @attr_name\tval_rg0[\tval_rg1...]
            let mut parts = rest.splitn(2, '\t');
            let key = match parts.next() {
                Some(k) if !k.is_empty() => k.to_string(),
                _ => continue,
            };
            let vals: Vec<String> = parts
                .next()
                .map(|s| s.split('\t').map(str::to_string).collect())
                .unwrap_or_default();
            if !attr_vals.contains_key(&key) {
                attr_keys_order.push(key.clone());
            }
            attr_vals.insert(key, vals);
            continue;
        }

        if let Some(rest) = t.strip_prefix('#') {
            // Split off the first field
            let (first, remainder) = match rest.find('\t') {
                Some(i) => (&rest[..i], Some(&rest[i + 1..])),
                None => (rest, None),
            };

            match first {
                "slow5_version" => {
                    let ver = remainder.unwrap_or("0.0.0");
                    let mut it = ver.splitn(3, '.');
                    version.0 = it.next().and_then(|s| s.parse().ok()).unwrap_or(0);
                    version.1 = it.next().and_then(|s| s.parse().ok()).unwrap_or(0);
                    version.2 = it.next().and_then(|s| s.parse().ok()).unwrap_or(0);
                }
                "num_read_groups" => {
                    let parsed = remainder.and_then(|s| s.parse::<u32>().ok()).unwrap_or(1);
                    const MAX_READ_GROUPS: u32 = 65536;
                    if parsed > MAX_READ_GROUPS {
                        return Err(SlowError::InvalidFormat(format!(
                            "num_read_groups {parsed} exceeds {MAX_READ_GROUPS} limit"
                        )));
                    }
                    num_read_groups = parsed;
                }
                "read_id" => {
                    // Column header line; first col is "read_id" (already stripped '#')
                    let mut col_names = vec!["read_id".to_string()];
                    if let Some(r) = remainder {
                        col_names.extend(r.split('\t').map(str::to_string));
                    }
                    col_idx = Some(find_primary_col_indices(&col_names)?);
                    // Save the full line for aux_meta construction after the loop.
                    name_line = Some(t.to_string());
                    break;
                }
                _ => {
                    // This is the type line that precedes the column name line.
                    type_line = Some(t.to_string());
                }
            }
            continue;
        }
    }

    let col_idx = col_idx.ok_or_else(|| {
        SlowError::InvalidFormat("SLOW5 column header line (#read_id ...) not found".into())
    })?;

    let rg_count = num_read_groups.max(1) as usize;
    let mut read_groups: Vec<ReadGroup> = (0..rg_count).map(|_| ReadGroup::default()).collect();

    for key in &attr_keys_order {
        let vals = &attr_vals[key];
        for (i, rg) in read_groups.iter_mut().enumerate() {
            let val = vals.get(i).cloned().unwrap_or_else(|| ".".to_string());
            rg.attributes.insert(key.clone(), val);
        }
    }

    let aux_meta = match (type_line, name_line) {
        (Some(tl), Some(nl)) => crate::aux::parse_aux_meta_from_lines(&tl, &nl)?,
        _ => crate::aux::AuxMeta::default(),
    };

    let header = Header {
        version,
        num_read_groups,
        record_compression: RecordCompression::None,
        signal_compression: SignalCompression::None,
        read_groups,
        aux_meta,
    };

    Ok((header, col_idx))
}

fn find_primary_col_indices(cols: &[String]) -> Result<PrimaryColIndices> {
    let find = |name: &str| -> Result<usize> {
        cols.iter().position(|c| c == name).ok_or_else(|| {
            SlowError::InvalidFormat(format!("SLOW5 column '{name}' not found in header"))
        })
    };

    Ok(PrimaryColIndices {
        read_id: find("read_id")?,
        read_group: find("read_group")?,
        digitisation: find("digitisation")?,
        offset: find("offset")?,
        range: find("range")?,
        sampling_rate: find("sampling_rate")?,
        len_raw_signal: find("len_raw_signal")?,
        raw_signal: find("raw_signal")?,
    })
}

/// Read one record from a SLOW5 text file. Returns `Ok(None)` at EOF.
pub(crate) fn read_next_record(
    reader: &mut impl BufRead,
    col_idx: &PrimaryColIndices,
    aux_meta: &crate::aux::AuxMeta,
) -> Result<Option<RawRecord>> {
    let mut line = String::new();

    loop {
        line.clear();
        let n = reader.read_line(&mut line).map_err(SlowError::Io)?;
        if n == 0 {
            return Ok(None);
        }

        let t = line.trim_end();
        if t.is_empty() || t.starts_with('#') || t.starts_with('@') {
            continue;
        }

        return parse_record_line(t, col_idx, aux_meta).map(Some);
    }
}

pub(crate) fn parse_record_line(
    line: &str,
    col_idx: &PrimaryColIndices,
    aux_meta: &crate::aux::AuxMeta,
) -> Result<RawRecord> {
    let fields: Vec<&str> = line.split('\t').collect();

    let get = |idx: usize, name: &str| -> Result<&str> {
        fields.get(idx).copied().ok_or_else(|| {
            SlowError::InvalidFormat(format!(
                "SLOW5 record too short: missing field '{name}' at column {idx}"
            ))
        })
    };

    let parse_f64 = |s: &str, name: &str| -> Result<f64> {
        s.parse::<f64>()
            .map_err(|e| SlowError::InvalidFormat(format!("invalid f64 for '{name}': {e}")))
    };

    let parse_u32 = |s: &str, name: &str| -> Result<u32> {
        s.parse::<u32>()
            .map_err(|e| SlowError::InvalidFormat(format!("invalid u32 for '{name}': {e}")))
    };

    let read_id = get(col_idx.read_id, "read_id")?.to_string();
    let read_group = parse_u32(get(col_idx.read_group, "read_group")?, "read_group")?;
    let digitisation = parse_f64(get(col_idx.digitisation, "digitisation")?, "digitisation")?;
    let offset = parse_f64(get(col_idx.offset, "offset")?, "offset")?;
    let range = parse_f64(get(col_idx.range, "range")?, "range")?;
    let sampling_rate = parse_f64(
        get(col_idx.sampling_rate, "sampling_rate")?,
        "sampling_rate",
    )?;

    // len_raw_signal present in file but we re-derive from the actual parsed signal
    let _ = get(col_idx.len_raw_signal, "len_raw_signal")?;

    let signal_str = get(col_idx.raw_signal, "raw_signal")?;
    let signal_bytes: Vec<u8> = signal_str
        .split(',')
        .map(|s| {
            s.parse::<i16>()
                .map(|v| v.to_le_bytes())
                .map_err(|e| SlowError::InvalidFormat(format!("invalid i16 in raw_signal: {e}")))
        })
        .collect::<Result<Vec<[u8; 2]>>>()?
        .into_iter()
        .flatten()
        .collect();

    // Parse auxiliary fields from the columns following raw_signal.
    // Aux field column indices start at max(primary_col_indices) + 1, but since
    // the schema (aux_meta) gives us name and type in order, we look them up by
    // the column names stored in aux_meta.names (in order), finding each by
    // searching for the name in the parsed fields vector starting after raw_signal.
    let aux = if aux_meta.is_empty() {
        std::collections::HashMap::new()
    } else {
        // Build a column-name -> index map from the field positions.
        // We already know the primary indices; aux fields appear after raw_signal.
        // We find each aux column by searching fields for its name -- but we don't
        // have the full column name list here, only the data. Instead we rely on
        // the aux fields appearing in schema order immediately after raw_signal.
        // The raw_signal column index tells us where the primary fields end.
        let aux_start_col = col_idx.raw_signal + 1;
        let mut map = std::collections::HashMap::with_capacity(aux_meta.len());
        for (i, (name, typ)) in aux_meta.names.iter().zip(aux_meta.types.iter()).enumerate() {
            let col = aux_start_col + i;
            let s = fields.get(col).copied().unwrap_or(".");
            let val = crate::aux::decode_text(s, typ)?;
            map.insert(name.clone(), val);
        }
        map
    };

    Ok(RawRecord {
        read_id,
        read_group,
        digitisation,
        offset,
        range,
        sampling_rate,
        signal_bytes,
        signal_compression: SignalCompression::None,
        aux,
    })
}

/// Write the SLOW5 text header to `writer`.
pub(crate) fn write_header(writer: &mut impl Write, header: &Header) -> Result<()> {
    let (maj, min, pat) = header.version;
    writeln!(writer, "#slow5_version\t{maj}.{min}.{pat}")?;
    writeln!(writer, "#num_read_groups\t{}", header.num_read_groups)?;

    let all_keys: BTreeSet<&str> = header
        .read_groups
        .iter()
        .flat_map(|rg| rg.attributes.keys().map(String::as_str))
        .collect();

    for key in &all_keys {
        write!(writer, "@{key}")?;
        for rg in &header.read_groups {
            let val = rg.attributes.get(*key).map(String::as_str).unwrap_or(".");
            if val.is_empty() {
                write!(writer, "\t.")?;
            } else {
                write!(writer, "\t{val}")?;
            }
        }
        writeln!(writer)?;
    }

    // Type and column name lines extended with aux fields
    write!(writer, "{}", TYPE_LINE.trim_end_matches('\n'))?;
    for typ in &header.aux_meta.types {
        write!(writer, "\t{}", crate::aux::aux_type_str(typ))?;
    }
    writeln!(writer)?;

    write!(writer, "{}", COL_LINE.trim_end_matches('\n'))?;
    for name in &header.aux_meta.names {
        write!(writer, "\t{name}")?;
    }
    writeln!(writer)?;

    Ok(())
}

/// Write one record to a SLOW5 text file.
pub(crate) fn write_record(
    writer: &mut impl Write,
    record: &Record,
    aux_meta: &crate::aux::AuxMeta,
) -> Result<()> {
    let n = record.raw_signal.len();
    write!(
        writer,
        "{}\t{}\t{}\t{}\t{}\t{}\t{}\t",
        record.read_id,
        record.read_group,
        record.digitisation,
        record.offset,
        record.range,
        record.sampling_rate,
        n,
    )?;

    // Comma-separated signal
    for (i, &s) in record.raw_signal.iter().enumerate() {
        if i > 0 {
            writer.write_all(b",")?;
        }
        write!(writer, "{s}")?;
    }

    // Auxiliary fields in schema order
    for (name, typ) in aux_meta.names.iter().zip(aux_meta.types.iter()) {
        let val = record
            .aux
            .get(name)
            .unwrap_or(&crate::aux::AuxValue::Missing);
        write!(writer, "\t{}", crate::aux::encode_text(val))?;
        let _ = typ; // type is needed for binary encoding; text uses the value's own Display
    }

    writeln!(writer)?;
    Ok(())
}