ahda 0.1.0

// ahda: Pseudoalignment compression and conversion between formats.
//
// Copyright 2025 Tommi Mäklin [tommi@maklin.fi].
//
// Copyrights in this project are retained by contributors. No copyright assignment
// is required to contribute to this project.
//
// Except as otherwise noted (below and/or in individual files), this
// project is licensed under the Apache License, Version 2.0
// <LICENSE-APACHE> or <http://www.apache.org/licenses/LICENSE-2.0> or
// the MIT license, <LICENSE-MIT> or <http://opensource.org/licenses/MIT>,
// at your option.
//

//! C++ API for encoding and decoding .ahda records.
//!
//! The API operates on the set bit indexes of a *flattened pseudoalignment*,
//! which is a `num_queries * num_targets` vector of boolean values, where the
//! value (query_index, target_index) implies an alignment for the query
//! sequence with `query_index` against the target sequence with `target_index`.
//!
//! The flattened pseudoalignment is assumed to be stored in a query-major
//! format, meaning that the pseudoalignments bits of a single query sequence
//! are stored contiguously in memory.
//!
//! If called via [encode_bitmap], the API will automatically determine whether
//! to use a 32-bit or 64-bit address space.
//!
//! If called via [encode_block32], the API will use a 32-bit address space
//! regardless of the input size. This implies a maximum size of `num_queries *
//! num_targets < 2^32` for the input alignment.
//!
//! If called via [encode_block64], the API will use a 64-bit address space.
//! This can fit a `num_queries * num_targets < 2^64` input alignment.
//!
//! ## Usage
//! **TODO** Usage examples for the C++ API.
//!

use crate::decode_from_read_to_roaring;
use crate::headers::file::build_file_header_and_flags;
use crate::headers::file::read_file_header_and_flags;
use crate::headers::block::read_block_header_and_flags;
use crate::encoder::bitmap_encoder::BitmapEncoder;
use crate::compression::BitmapHolder;
use crate::compression::MetadataCompression;
use crate::compression::roaringwrapper::pack_block_roaring;

use std::io::Cursor;

use cxx::CxxString;
use cxx::CxxVector;
use roaring::RoaringBitmap;
use roaring::RoaringTreemap;

#[cxx::bridge(namespace = "ahda")]
mod ffi {

    extern "Rust" {
        fn encode_file_header_and_flags(
            targets: &CxxVector<CxxString>,
            name: &CxxString,
            n_queries: u32,
        ) -> Vec<u8>;

        fn encode_block32(
            queries: &CxxVector<CxxString>,
            query_ids: &CxxVector<u32>,
            set_bits: &CxxVector<u32>,
        ) -> Vec<u8>;

        fn encode_block64(
            queries: &CxxVector<CxxString>,
            query_ids: &CxxVector<u32>,
            set_bits: &CxxVector<u64>,
        ) -> Vec<u8>;

        fn encode_bitmap(
            targets: &CxxVector<CxxString>,
            queries: &CxxVector<CxxString>,
            name: &CxxString,
            set_bits: &CxxVector<u64>,
            n_queries: u32,
        ) -> Vec<u8>;

        fn decode_bitmap(
            bytes: &CxxVector<u8>,
        ) -> Vec<u64>;

        fn decode_target_names(
            bytes: &CxxVector<u8>,
        ) -> Vec<String>;

        pub fn decode_query_ids(
            bytes: &CxxVector<u8>,
        ) -> Vec<u32>;

        pub fn decode_query_names(
            bytes: &CxxVector<u8>,
        ) -> Vec<String>;
}
}

/// Encode the file header and file flags bytes.
///
/// Calls [build_file_header_and_flags] on the input data and then creates the
/// encoded data by calling [encode_file_header_and_flags].
///
/// The output bytes should always be written at the start of the .ahda record.
///
pub fn encode_file_header_and_flags(
    targets: &CxxVector<CxxString>,
    name: &CxxString,
    n_queries: u32,
) -> Vec<u8> {
    let target_names: Vec<Vec<u8>> = targets.iter().map(|x| x.as_bytes().to_vec()).collect();
    let query_name: Vec<u8> = name.as_bytes().to_vec();

    let (mut header, flags) = build_file_header_and_flags(&target_names, n_queries.try_into().unwrap(), &query_name, &MetadataCompression::default()).unwrap();
    let bytes: Vec<u8> = crate::headers::file::encode_file_header_and_flags(&mut header, &flags).unwrap();

    bytes
}

/// Encode a single .ahda block and its block header and flags in 32-bit address space.
///
/// Creates a [RoaringBitmap] from the set bit indexes and calls [pack_block_roaring] to
/// encode the block header, block flags, and block contents.
///
/// The output is a valid block record that can be appended to an .ahda record
/// containing the corresponding file header and flags.
///
pub fn encode_block32(
    queries: &CxxVector<CxxString>,
    query_ids: &CxxVector<u32>,
    set_bits: &CxxVector<u32>,
) -> Vec<u8> {
    let bitmap = BitmapHolder::Roaring32(RoaringBitmap::from_iter(set_bits.iter()));
    let query_names: Vec<Vec<u8>> = queries.iter().map(|x| x.as_bytes().to_vec()).collect();
    let query_names = if query_names.is_empty() { None } else { Some(query_names) };
    let block = pack_block_roaring(query_ids.as_slice(), query_names, bitmap);
    block.unwrap()
}

/// Encode a single .ahda block and its block header and flags in 64-bit address space.
///
/// Creates a [RoaringTreemap] from the set bit indexes and calls [pack_block_roaring] to
/// encode the block header, block flags, and block contents.
///
/// The output is a valid block record that can be appended to an .ahda record
/// containing the corresponding file header and flags.
///
pub fn encode_block64(
    queries: &CxxVector<CxxString>,
    query_ids: &CxxVector<u32>,
    set_bits: &CxxVector<u64>,
) -> Vec<u8> {
    let bitmap = BitmapHolder::Roaring64(RoaringTreemap::from_iter(set_bits.iter()));
    let query_names: Vec<Vec<u8>> = queries.iter().map(|x| x.as_bytes().to_vec()).collect();
    let query_names = if query_names.is_empty() { None } else { Some(query_names) };
    let block = pack_block_roaring(query_ids.as_slice(), query_names, bitmap);
    block.unwrap()
}

/// Encode a complete .ahda record from the set bits in a flattened pseudoalignment.
///
/// Creates an iterator over the set bit indexes and uses a [BitmapEncoder] to
/// encode a valid .ahda record.
///
/// The output is a complete .ahda record that contains the file header, file
/// flags, and all block data required to store the alignment. This can be
/// written to a file without further API calls.
///
pub fn encode_bitmap(
    targets: &CxxVector<CxxString>,
    queries: &CxxVector<CxxString>,
    name: &CxxString,
    set_bits: &CxxVector<u64>,
    n_queries: u32,
) -> Vec<u8> {
    let query_names: Vec<Vec<u8>> = queries.iter().map(|x| x.as_bytes().to_vec()).collect();
    let target_names: Vec<Vec<u8>> = targets.iter().map(|x| x.as_bytes().to_vec()).collect();
    let query_name: Vec<u8> = name.as_bytes().to_vec();

    let mut set_bits_iter = set_bits.as_slice().iter().cloned();
    let mut encoder = BitmapEncoder::new(&mut set_bits_iter, &target_names, &query_name, n_queries as usize).unwrap();
    if !query_names.is_empty() {
        encoder.set_query_names(&query_names);
    }

    let mut bytes: Vec<u8> = encoder.encode_file_header_and_flags().unwrap();
    for block in encoder {
        bytes.extend(block.unwrap());
    }

    bytes
}

/// Decodes the indexes of set bits in a flattened pseudoalignment from an .ahda record.
///
/// Calls [decode_from_read_to_roaring] to extract a [RoaringBitmap]
/// representing the full pseudoalignment record contained in the .ahda record.
///
/// The input should contain the full contents of the .ahda record to be decoded.
///
/// The output is a vector containing the indexes of set bits, ie.
/// pseudoalignments, in a flattened pseudoalignment.
///
pub fn decode_bitmap(
    bytes: &CxxVector<u8>,
) -> Vec<u64> {
    let mut cursor = Cursor::new(bytes.as_slice());
    let (bitmap, _, _, _) = decode_from_read_to_roaring(&mut cursor).unwrap();
    let set_bits: Vec<u64> = bitmap.iter().collect();
    set_bits
}

/// Decodes the target sequence names from the file flags of an .ahda record.
///
/// The input should contain at least the bytes representing the [FileHeader]
/// and [FileFlags] in the .ahda record.
///
/// The output is a vector with the names of the target sequences. The position
/// of each element in the vector corresponds to its index in the
/// pseudoalignment.
///
pub fn decode_target_names(
    bytes: &CxxVector<u8>,
) -> Vec<String> {
    let mut cursor = Cursor::new(bytes.as_slice());
    let (header, flags) = read_file_header_and_flags(&mut cursor).unwrap();
    assert_eq!(header.n_targets as usize, flags.target_names.len());
    flags.target_names.into_iter().map(|x| String::from_utf8(x.to_vec()).unwrap()).collect()
}

/// Decodes the query sequence names from the block flags in an .ahda record.
///
/// The input should contain a complete .ahda record with the file header and
/// flags and all block data.
///
/// The output is a vector with the names of the query sequences. The position
/// of each element in the vector does **not** necessarily correspond to its
/// index in the original data, use [decode_query_ids] to get this information.
///
pub fn decode_query_names(
    bytes: &CxxVector<u8>,
) -> Vec<String> {
    let mut cursor = Cursor::new(bytes.as_slice());
    let (header, _) = read_file_header_and_flags(&mut cursor).unwrap();

    let mut query_names: Vec<String> = Vec::with_capacity(header.n_queries as usize);
    while let Ok((header, flags)) = read_block_header_and_flags(&mut cursor) {
        assert_eq!(header.num_records as usize, query_names.len());
        assert_eq!(flags.query_ids.as_ref().unwrap().len(), query_names.len());
        let mut queries_str = flags.queries.unwrap().into_iter().map(|x| String::from_utf8(x.to_vec()).unwrap());
        query_names.extend(&mut queries_str);
        cursor.set_position(cursor.position() + header.block_len as u64);
    }

    query_names
}

/// Decodes the query sequence indexes from the block flags in an .ahda record.
///
/// The input should contain a complete .ahda record with the file header and
/// flags and all block data.
///
/// The output is a vector with the indexes of the query sequences in the
/// original query sequence file. Use [decode_query_names] to get the vector
/// containing the names that can be indexed by the output from
/// decode_query_ids.
///
pub fn decode_query_ids(
    bytes: &CxxVector<u8>,
) -> Vec<u32> {
    let mut cursor = Cursor::new(bytes.as_slice());
    let (header, _) = read_file_header_and_flags(&mut cursor).unwrap();

    let mut query_ids: Vec<u32> = Vec::with_capacity(header.n_queries as usize);
    while let Ok((header, flags)) = read_block_header_and_flags(&mut cursor) {
        assert_eq!(header.num_records as usize, query_ids.len());
        assert_eq!(flags.query_ids.as_ref().unwrap().len(), query_ids.len());
        query_ids.append(&mut flags.query_ids.unwrap());
        cursor.set_position(cursor.position() + header.block_len as u64);
    }

    query_ids
}