slow5lib 0.1.0

use std::collections::HashMap;

use slow5lib::aux::{AuxMeta, AuxType, AuxValue};
use slow5lib::header::{Header, ReadGroup, RecordCompression, SignalCompression};
use slow5lib::record::Record;
use slow5lib::{Slow5IndexedReader, Slow5Reader, Slow5Writer};

const FIVE_READS_ZSTD: &str = "tests/data/five_reads.blow5";
const FIVE_READS_NONE: &str = "tests/data/five_reads_nocompress.blow5";
const FIVE_READS_SLOW5: &str = "tests/data/five_reads.slow5";

// Ground-truth values for the first record extracted from the C slow5lib via slow5tools.
const READ1_ID: &str = "50abece6-9476-46be-a86f-d48cc3fddeb6";
const READ1_READ_GROUP: u32 = 0;
const READ1_DIGITISATION: f64 = 2048.0;
const READ1_OFFSET: f64 = -224.0;
const READ1_RANGE: f64 = 383.119049;
const READ1_SAMPLING_RATE: f64 = 5000.0;
const READ1_LEN: usize = 167180;
const READ1_FIRST_10: [i16; 10] = [761, 791, 772, 790, 761, 771, 807, 764, 784, 786];
const READ1_LAST_5: [i16; 5] = [1011, 980, 1026, 987, 1005];

fn check_read1(r: &Record) {
    assert_eq!(r.read_id, READ1_ID);
    assert_eq!(r.read_group, READ1_READ_GROUP);
    assert!(
        (r.digitisation - READ1_DIGITISATION).abs() < 1e-6,
        "digitisation mismatch"
    );
    assert!((r.offset - READ1_OFFSET).abs() < 1e-6, "offset mismatch");
    assert!((r.range - READ1_RANGE).abs() < 1e-6, "range mismatch");
    assert!(
        (r.sampling_rate - READ1_SAMPLING_RATE).abs() < 1e-6,
        "sampling_rate mismatch"
    );
    assert_eq!(r.raw_signal.len(), READ1_LEN, "signal length mismatch");
    assert_eq!(
        &r.raw_signal[..10],
        &READ1_FIRST_10,
        "first 10 samples mismatch"
    );
    assert_eq!(
        &r.raw_signal[r.raw_signal.len() - 5..],
        &READ1_LAST_5,
        "last 5 samples mismatch"
    );
}

// ── Sequential reader tests ───────────────────────────────────────────────────

#[test]
fn sequential_read_zstd_svbzd_count() {
    let mut reader = Slow5Reader::open(FIVE_READS_ZSTD).expect("open five_reads.blow5");
    let count = reader.records().count();
    assert_eq!(count, 5);
}

#[test]
fn sequential_read_zstd_svbzd_first_record() {
    let mut reader = Slow5Reader::open(FIVE_READS_ZSTD).expect("open five_reads.blow5");
    let record = reader
        .records()
        .next()
        .expect("at least one record")
        .expect("record ok");
    check_read1(&record);
}

#[test]
fn sequential_read_nocompress_count() {
    let mut reader = Slow5Reader::open(FIVE_READS_NONE).expect("open five_reads_nocompress.blow5");
    let count = reader.records().count();
    assert_eq!(count, 5);
}

#[test]
fn sequential_read_nocompress_first_record() {
    let mut reader = Slow5Reader::open(FIVE_READS_NONE).expect("open five_reads_nocompress.blow5");
    let record = reader
        .records()
        .next()
        .expect("at least one record")
        .expect("record ok");
    check_read1(&record);
}

// ── SLOW5 text reader tests ───────────────────────────────────────────────────

#[test]
fn sequential_read_slow5_count() {
    let mut reader = Slow5Reader::open(FIVE_READS_SLOW5).expect("open five_reads.slow5");
    let count = reader.records().count();
    assert_eq!(count, 5);
}

#[test]
fn sequential_read_slow5_first_record() {
    let mut reader = Slow5Reader::open(FIVE_READS_SLOW5).expect("open five_reads.slow5");
    let record = reader
        .records()
        .next()
        .expect("at least one record")
        .expect("record ok");
    check_read1(&record);
}

// ── Indexed reader tests ──────────────────────────────────────────────────────

#[test]
fn indexed_read_zstd_svbzd_get() {
    let reader = Slow5IndexedReader::open(FIVE_READS_ZSTD).expect("open indexed");
    let record = reader.get(READ1_ID).expect("get read1");
    check_read1(&record);
}

#[test]
fn indexed_read_zstd_svbzd_not_found() {
    let reader = Slow5IndexedReader::open(FIVE_READS_ZSTD).expect("open indexed");
    let result = reader.get("does-not-exist-id");
    assert!(result.is_err());
}

#[test]
fn indexed_read_zstd_svbzd_id_count() {
    let reader = Slow5IndexedReader::open(FIVE_READS_ZSTD).expect("open indexed");
    assert_eq!(reader.read_ids().count(), 5);
}

// ── Indexed reader tests (SLOW5 text) ─────────────────────────────────────────

#[test]
fn indexed_read_slow5_get() {
    let reader = Slow5IndexedReader::open(FIVE_READS_SLOW5).expect("open indexed slow5");
    let record = reader.get(READ1_ID).expect("get read1 from slow5");
    check_read1(&record);
}

#[test]
fn indexed_read_slow5_not_found() {
    let reader = Slow5IndexedReader::open(FIVE_READS_SLOW5).expect("open indexed slow5");
    let result = reader.get("does-not-exist-id");
    assert!(result.is_err());
}

#[test]
fn indexed_read_slow5_id_count() {
    let reader = Slow5IndexedReader::open(FIVE_READS_SLOW5).expect("open indexed slow5");
    assert_eq!(reader.read_ids().count(), 5);
}

#[test]
fn indexed_read_slow5_aux() {
    let reader = Slow5IndexedReader::open(FIVE_READS_SLOW5).expect("open indexed slow5");
    let record = reader.get(READ1_ID).expect("get read1");
    check_read1_aux(&record);
}

// ── Round-trip tests ──────────────────────────────────────────────────────────

fn make_test_record(read_id: &str, signal: Vec<i16>) -> Record {
    Record {
        read_id: read_id.to_string(),
        read_group: 0,
        digitisation: 2048.0,
        offset: -224.0,
        range: 383.119049,
        sampling_rate: 5000.0,
        raw_signal: signal,
        aux: HashMap::new(),
    }
}

fn roundtrip_header(rec_comp: RecordCompression, sig_comp: SignalCompression) -> Header {
    Header {
        version: (0, 2, 0),
        num_read_groups: 1,
        record_compression: rec_comp,
        signal_compression: sig_comp,
        read_groups: vec![ReadGroup::default()],
        aux_meta: AuxMeta::default(),
    }
}

fn roundtrip_write_read(
    path: &str,
    records: &[Record],
    rec_comp: RecordCompression,
    sig_comp: SignalCompression,
) -> Vec<Record> {
    let header = roundtrip_header(rec_comp, sig_comp);
    let mut writer = Slow5Writer::create(path, header).expect("create writer");
    for rec in records {
        writer.write(rec).expect("write record");
    }
    writer.finish().expect("finish writer");

    let mut reader = Slow5Reader::open(path).expect("open for read");
    reader.records().map(|r| r.expect("record ok")).collect()
}

#[test]
fn roundtrip_zstd_svbzd() {
    let tmp = "/tmp/slow5lib_roundtrip_zstd.blow5";
    let signal: Vec<i16> = READ1_FIRST_10.to_vec();
    let original = make_test_record(READ1_ID, signal.clone());

    let records = roundtrip_write_read(
        tmp,
        std::slice::from_ref(&original),
        RecordCompression::Zstd,
        SignalCompression::SvbZd,
    );

    assert_eq!(records.len(), 1);
    assert_eq!(records[0].read_id, READ1_ID);
    assert_eq!(records[0].raw_signal, signal);
}

#[test]
fn roundtrip_zstd_exzd() {
    let tmp = "/tmp/slow5lib_roundtrip_zstd_exzd.blow5";
    let signal: Vec<i16> = READ1_FIRST_10.to_vec();
    let original = make_test_record(READ1_ID, signal.clone());

    let records = roundtrip_write_read(
        tmp,
        std::slice::from_ref(&original),
        RecordCompression::Zstd,
        SignalCompression::ExZd,
    );

    assert_eq!(records.len(), 1);
    assert_eq!(records[0].read_id, READ1_ID);
    assert_eq!(records[0].raw_signal, signal);
}

#[test]
fn roundtrip_none_exzd() {
    let tmp = "/tmp/slow5lib_roundtrip_none_exzd.blow5";
    let signal: Vec<i16> = READ1_FIRST_10.to_vec();
    let original = make_test_record(READ1_ID, signal.clone());

    let records = roundtrip_write_read(
        tmp,
        std::slice::from_ref(&original),
        RecordCompression::None,
        SignalCompression::ExZd,
    );

    assert_eq!(records.len(), 1);
    assert_eq!(records[0].read_id, READ1_ID);
    assert_eq!(records[0].raw_signal, signal);
}

#[test]
fn roundtrip_none_none() {
    let tmp = "/tmp/slow5lib_roundtrip_none.blow5";
    let signal: Vec<i16> = READ1_FIRST_10.to_vec();
    let original = make_test_record(READ1_ID, signal.clone());

    let records = roundtrip_write_read(
        tmp,
        std::slice::from_ref(&original),
        RecordCompression::None,
        SignalCompression::None,
    );

    assert_eq!(records.len(), 1);
    assert_eq!(records[0].raw_signal, signal);
}

#[test]
fn roundtrip_zlib_svbzd() {
    let tmp = "/tmp/slow5lib_roundtrip_zlib.blow5";
    let signal: Vec<i16> = READ1_FIRST_10.to_vec();
    let original = make_test_record(READ1_ID, signal.clone());

    let records = roundtrip_write_read(
        tmp,
        std::slice::from_ref(&original),
        RecordCompression::Zlib,
        SignalCompression::SvbZd,
    );

    assert_eq!(records.len(), 1);
    assert_eq!(records[0].read_id, READ1_ID);
    assert_eq!(records[0].raw_signal, signal);
}

#[test]
fn roundtrip_zlib_none() {
    let tmp = "/tmp/slow5lib_roundtrip_zlib_none.blow5";
    let signal: Vec<i16> = READ1_FIRST_10.to_vec();
    let original = make_test_record(READ1_ID, signal.clone());

    let records = roundtrip_write_read(
        tmp,
        std::slice::from_ref(&original),
        RecordCompression::Zlib,
        SignalCompression::None,
    );

    assert_eq!(records.len(), 1);
    assert_eq!(records[0].raw_signal, signal);
}

#[test]
fn roundtrip_zstd_svbzd_full_read() {
    let tmp = "/tmp/slow5lib_roundtrip_full.blow5";

    // Read the first record from the C-generated file as our source of truth
    let source_record = {
        let mut reader =
            Slow5Reader::open(FIVE_READS_ZSTD).expect("open five_reads.blow5 for source");
        reader
            .records()
            .next()
            .expect("at least one record")
            .expect("record ok")
    };

    // Write it with our Rust writer and read it back
    let header = roundtrip_header(RecordCompression::Zstd, SignalCompression::SvbZd);
    let mut writer = Slow5Writer::create(tmp, header).expect("create writer");
    writer.write(&source_record).expect("write record");
    writer.finish().expect("finish");

    let mut reader = Slow5Reader::open(tmp).expect("open roundtrip file");
    let roundtripped = reader
        .records()
        .next()
        .expect("one record")
        .expect("record ok");

    assert_eq!(roundtripped.read_id, source_record.read_id);
    assert_eq!(
        roundtripped.raw_signal.len(),
        source_record.raw_signal.len()
    );
    assert_eq!(
        roundtripped.raw_signal, source_record.raw_signal,
        "signal mismatch after roundtrip"
    );
}

#[test]
fn roundtrip_zstd_exzd_full_read() {
    let tmp = "/tmp/slow5lib_roundtrip_full_exzd.blow5";

    // Read the first record from the C-generated file as our source of truth
    let source_record = {
        let mut reader =
            Slow5Reader::open(FIVE_READS_ZSTD).expect("open five_reads.blow5 for source");
        reader
            .records()
            .next()
            .expect("at least one record")
            .expect("record ok")
    };

    // Write it with ExZd signal compression and read it back
    let header = roundtrip_header(RecordCompression::Zstd, SignalCompression::ExZd);
    let mut writer = Slow5Writer::create(tmp, header).expect("create writer");
    writer.write(&source_record).expect("write record");
    writer.finish().expect("finish");

    let mut reader = Slow5Reader::open(tmp).expect("open roundtrip file");
    let roundtripped = reader
        .records()
        .next()
        .expect("one record")
        .expect("record ok");

    assert_eq!(roundtripped.read_id, source_record.read_id);
    assert_eq!(
        roundtripped.raw_signal.len(),
        source_record.raw_signal.len()
    );
    assert_eq!(
        roundtripped.raw_signal, source_record.raw_signal,
        "signal mismatch after roundtrip"
    );
}

#[test]
fn roundtrip_slow5_text() {
    let tmp = "/tmp/slow5lib_roundtrip.slow5";
    let signal: Vec<i16> = READ1_FIRST_10.to_vec();
    let original = make_test_record(READ1_ID, signal.clone());

    let header = Header {
        version: (0, 2, 0),
        num_read_groups: 1,
        record_compression: RecordCompression::None,
        signal_compression: SignalCompression::None,
        read_groups: vec![ReadGroup::default()],
        aux_meta: AuxMeta::default(),
    };
    let mut writer = Slow5Writer::create(tmp, header).expect("create slow5 writer");
    writer.write(&original).expect("write record");
    writer.finish().expect("finish slow5 writer");

    let mut reader = Slow5Reader::open(tmp).expect("open slow5 for read");
    let record = reader
        .records()
        .next()
        .expect("at least one record")
        .expect("record ok");

    assert_eq!(record.read_id, READ1_ID);
    assert_eq!(record.raw_signal, signal);
}

#[test]
fn roundtrip_slow5_text_full_read() {
    let tmp = "/tmp/slow5lib_roundtrip_full.slow5";

    let source_record = {
        let mut reader =
            Slow5Reader::open(FIVE_READS_SLOW5).expect("open five_reads.slow5 for source");
        reader
            .records()
            .next()
            .expect("at least one record")
            .expect("record ok")
    };

    let header = Header {
        version: (0, 2, 0),
        num_read_groups: 1,
        record_compression: RecordCompression::None,
        signal_compression: SignalCompression::None,
        read_groups: vec![ReadGroup::default()],
        aux_meta: AuxMeta::default(),
    };
    let mut writer = Slow5Writer::create(tmp, header).expect("create writer");
    writer.write(&source_record).expect("write record");
    writer.finish().expect("finish");

    let mut reader = Slow5Reader::open(tmp).expect("open roundtrip file");
    let roundtripped = reader
        .records()
        .next()
        .expect("one record")
        .expect("record ok");

    assert_eq!(roundtripped.read_id, source_record.read_id);
    assert_eq!(
        roundtripped.raw_signal.len(),
        source_record.raw_signal.len()
    );
    assert_eq!(
        roundtripped.raw_signal, source_record.raw_signal,
        "signal mismatch after SLOW5 roundtrip"
    );
}

#[test]
fn roundtrip_multiple_records() {
    let tmp = "/tmp/slow5lib_roundtrip_multi.blow5";

    let signals: Vec<Vec<i16>> = vec![
        vec![100, 200, 300, 400, 500],
        vec![-100, -200, -300],
        vec![0, 1, -1, i16::MAX, i16::MIN],
    ];
    let originals: Vec<Record> = signals
        .iter()
        .enumerate()
        .map(|(i, s)| make_test_record(&format!("read-{i:03}"), s.clone()))
        .collect();

    let records = roundtrip_write_read(
        tmp,
        &originals,
        RecordCompression::Zstd,
        SignalCompression::SvbZd,
    );

    assert_eq!(records.len(), 3);
    for (got, want) in records.iter().zip(originals.iter()) {
        assert_eq!(got.read_id, want.read_id);
        assert_eq!(got.raw_signal, want.raw_signal);
    }
}

// ── Auxiliary field tests ─────────────────────────────────────────────────────

// Ground-truth aux values for read1, extracted via slow5tools from five_reads.slow5.
// Aux schema (14 fields): channel_number(char*), median_before(double),
// read_number(int32_t), start_mux(uint8_t), start_time(uint64_t),
// end_reason(enum{...}), tracked_scaling_shift(float), tracked_scaling_scale(float),
// predicted_scaling_shift(float), predicted_scaling_scale(float),
// num_reads_since_mux_change(uint32_t), time_since_mux_change(float),
// num_minknow_events(uint64_t), open_pore_level(float)
fn check_read1_aux(r: &slow5lib::record::Record) {
    assert_eq!(
        r.aux.get("channel_number"),
        Some(&AuxValue::String("5".into()))
    );
    assert_eq!(r.aux.get("median_before"), Some(&AuxValue::Missing));
    assert_eq!(r.aux.get("read_number"), Some(&AuxValue::Int32(4215)));
    assert_eq!(r.aux.get("start_mux"), Some(&AuxValue::UInt8(2)));
    assert_eq!(r.aux.get("start_time"), Some(&AuxValue::UInt64(15728619)));
    assert_eq!(r.aux.get("end_reason"), Some(&AuxValue::Enum(4)));
    assert_eq!(
        r.aux.get("num_reads_since_mux_change"),
        Some(&AuxValue::UInt32(61))
    );
    assert_eq!(
        r.aux.get("num_minknow_events"),
        Some(&AuxValue::UInt64(18915))
    );
    assert_eq!(r.aux.get("open_pore_level"), Some(&AuxValue::Missing));

    // Float comparisons with tolerance
    if let Some(AuxValue::Float(v)) = r.aux.get("tracked_scaling_shift") {
        assert!(
            (*v - 103.24131f32).abs() < 1e-3,
            "tracked_scaling_shift mismatch"
        );
    } else {
        panic!("tracked_scaling_shift missing or wrong type");
    }
    if let Some(AuxValue::Float(v)) = r.aux.get("time_since_mux_change") {
        assert!(
            (*v - 2979.851807f32).abs() < 0.01,
            "time_since_mux_change mismatch"
        );
    } else {
        panic!("time_since_mux_change missing or wrong type");
    }
}

#[test]
fn aux_read_slow5_text() {
    let mut reader = Slow5Reader::open(FIVE_READS_SLOW5).expect("open five_reads.slow5");
    let record = reader
        .records()
        .next()
        .expect("at least one record")
        .expect("record ok");
    check_read1_aux(&record);
}

#[test]
fn aux_read_blow5_sequential() {
    let mut reader = Slow5Reader::open(FIVE_READS_ZSTD).expect("open five_reads.blow5");
    let record = reader
        .records()
        .next()
        .expect("at least one record")
        .expect("record ok");
    check_read1_aux(&record);
}

#[test]
fn aux_read_blow5_indexed() {
    let reader = Slow5IndexedReader::open(FIVE_READS_ZSTD).expect("open indexed");
    let record = reader.get(READ1_ID).expect("get read1");
    check_read1_aux(&record);
}

#[test]
fn aux_roundtrip_blow5() {
    let tmp = "/tmp/slow5lib_aux_roundtrip.blow5";

    let mut aux_meta = AuxMeta::default();
    aux_meta.names = vec![
        "channel_number".into(),
        "read_number".into(),
        "start_mux".into(),
    ];
    aux_meta.types = vec![AuxType::String, AuxType::Int32, AuxType::UInt8];

    let header = Header {
        version: (0, 2, 0),
        num_read_groups: 1,
        record_compression: RecordCompression::Zstd,
        signal_compression: SignalCompression::SvbZd,
        read_groups: vec![slow5lib::header::ReadGroup::default()],
        aux_meta: aux_meta.clone(),
    };

    let mut original = make_test_record(READ1_ID, READ1_FIRST_10.to_vec());
    original
        .aux
        .insert("channel_number".into(), AuxValue::String("42".into()));
    original
        .aux
        .insert("read_number".into(), AuxValue::Int32(999));
    original.aux.insert("start_mux".into(), AuxValue::UInt8(3));

    let mut writer = Slow5Writer::create(tmp, header).expect("create writer");
    writer.write(&original).expect("write record");
    writer.finish().expect("finish");

    let mut reader = Slow5Reader::open(tmp).expect("open roundtrip file");
    let got = reader
        .records()
        .next()
        .expect("one record")
        .expect("record ok");

    assert_eq!(
        got.aux.get("channel_number"),
        Some(&AuxValue::String("42".into()))
    );
    assert_eq!(got.aux.get("read_number"), Some(&AuxValue::Int32(999)));
    assert_eq!(got.aux.get("start_mux"), Some(&AuxValue::UInt8(3)));
}

#[test]
fn aux_roundtrip_slow5_text() {
    let tmp = "/tmp/slow5lib_aux_roundtrip.slow5";

    let mut aux_meta = AuxMeta::default();
    aux_meta.names = vec!["channel_number".into(), "read_number".into()];
    aux_meta.types = vec![AuxType::String, AuxType::Int32];

    let header = Header {
        version: (0, 2, 0),
        num_read_groups: 1,
        record_compression: RecordCompression::None,
        signal_compression: SignalCompression::None,
        read_groups: vec![slow5lib::header::ReadGroup::default()],
        aux_meta: aux_meta.clone(),
    };

    let mut original = make_test_record(READ1_ID, READ1_FIRST_10.to_vec());
    original
        .aux
        .insert("channel_number".into(), AuxValue::String("7".into()));
    original
        .aux
        .insert("read_number".into(), AuxValue::Int32(42));

    let mut writer = Slow5Writer::create(tmp, header).expect("create writer");
    writer.write(&original).expect("write record");
    writer.finish().expect("finish");

    let mut reader = Slow5Reader::open(tmp).expect("open roundtrip slow5");
    let got = reader
        .records()
        .next()
        .expect("one record")
        .expect("record ok");

    assert_eq!(
        got.aux.get("channel_number"),
        Some(&AuxValue::String("7".into()))
    );
    assert_eq!(got.aux.get("read_number"), Some(&AuxValue::Int32(42)));
}

#[test]
fn aux_missing_roundtrip_blow5() {
    let tmp = "/tmp/slow5lib_aux_missing_roundtrip.blow5";

    let mut aux_meta = AuxMeta::default();
    aux_meta.names = vec!["score".into()];
    aux_meta.types = vec![AuxType::Float];

    let header = Header {
        version: (0, 2, 0),
        num_read_groups: 1,
        record_compression: RecordCompression::None,
        signal_compression: SignalCompression::None,
        read_groups: vec![slow5lib::header::ReadGroup::default()],
        aux_meta,
    };

    let original = make_test_record(READ1_ID, vec![100, 200]);

    let mut writer = Slow5Writer::create(tmp, header).expect("create writer");
    writer.write(&original).expect("write record");
    writer.finish().expect("finish");

    let mut reader = Slow5Reader::open(tmp).expect("open roundtrip file");
    let got = reader
        .records()
        .next()
        .expect("one record")
        .expect("record ok");

    assert_eq!(got.aux.get("score"), Some(&AuxValue::Missing));
}