use super::*;
use obicompactvec::PersistentCompactIntMatrix;
use obikseq::{set_k, Sequence as _, Unitig};
use crate::meta::IndexMode;
use tempfile::tempdir;

fn push_unitigs_and_layer(
    map: &mut LayeredMap<PersistentCompactIntMatrix>,
    seqs: &[&[u8]],
    count: u32,
) {
    let mut w = map.next_layer_writer().unwrap();
    for s in seqs {
        w.write(&Unitig::from_ascii(s)).unwrap();
    }
    w.close().unwrap();
    map.push_layer(|_| count).unwrap();
}

fn canonical(ascii: &[u8]) -> CanonicalKmer {
    obikseq::Kmer::from_ascii(ascii).unwrap().canonical()
}

#[test]
fn create_empty_map() {
    set_k(4);
    let dir = tempdir().unwrap();
    let map = LayeredMap::<()>::create(dir.path(), IndexMode::Exact).unwrap();
    assert_eq!(map.n_layers(), 0);
}

#[test]
fn open_reloads_layer_count() {
    set_k(4);
    let dir = tempdir().unwrap();
    {
        let mut map = LayeredMap::<PersistentCompactIntMatrix>::create(dir.path(), IndexMode::Exact).unwrap();
        push_unitigs_and_layer(&mut map, &[b"AAAACGT"], 1);
    }
    let map = LayeredMap::<PersistentCompactIntMatrix>::open(dir.path()).unwrap();
    assert_eq!(map.n_layers(), 1);
}

#[test]
fn query_finds_kmer_in_layer_zero() {
    set_k(4);
    let dir = tempdir().unwrap();
    let mut map = LayeredMap::<PersistentCompactIntMatrix>::create(dir.path(), IndexMode::Exact).unwrap();
    push_unitigs_and_layer(&mut map, &[b"AAAACGT"], 3);
    let kmer = canonical(b"AAAC");
    let (layer_idx, hit) = map.query(kmer).expect("kmer must be found");
    assert_eq!(layer_idx, 0);
    assert_eq!(hit.data[0], 3);
}

#[test]
fn query_finds_kmer_in_correct_layer() {
    set_k(4);
    let dir = tempdir().unwrap();
    let mut map = LayeredMap::<PersistentCompactIntMatrix>::create(dir.path(), IndexMode::Exact).unwrap();
    push_unitigs_and_layer(&mut map, &[b"AAAACGT"], 1);
    push_unitigs_and_layer(&mut map, &[b"GGGACGT"], 2);
    assert_eq!(map.n_layers(), 2);

    let (li, hit) = map.query(canonical(b"AAAA")).expect("AAAA must be found");
    assert_eq!(li, 0);
    assert_eq!(hit.data[0], 1);

    let (li, hit) = map.query(canonical(b"GGGA")).expect("GGGA must be found");
    assert_eq!(li, 1);
    assert_eq!(hit.data[0], 2);
}

#[test]
fn query_absent_returns_none() {
    set_k(4);
    let dir = tempdir().unwrap();
    let mut map = LayeredMap::<PersistentCompactIntMatrix>::create(dir.path(), IndexMode::Exact).unwrap();
    push_unitigs_and_layer(&mut map, &[b"AAAACGT"], 1);
    let absent = canonical(b"CCCC");
    assert!(map.query(absent).is_none());
}

#[test]
fn push_layer_from_map_convenience() {
    set_k(4);
    let dir = tempdir().unwrap();
    let mut map = LayeredMap::<PersistentCompactIntMatrix>::create(dir.path(), IndexMode::Exact).unwrap();
    let mut w = map.next_layer_writer().unwrap();
    w.write(&Unitig::from_ascii(b"AAAACGT")).unwrap();
    w.close().unwrap();
    let counts: HashMap<CanonicalKmer, u32> = vec![
        (canonical(b"AAAA"), 10u32),
    ].into_iter().collect();
    map.push_layer_from_map(&counts).unwrap();
    let (_, hit) = map.query(canonical(b"AAAA")).unwrap();
    assert_eq!(hit.data[0], 10);
}
