fibertools-rs 0.11.0

use super::subcommands::center::CenterPosition;
use super::utils::input_bam::FiberFilters;
use super::*;
use crate::utils::bamannotations::*;
use crate::utils::basemods::{CPG_TYPE, M6A_TYPE};
use crate::utils::bio_io::*;
use crate::utils::ftexpression::apply_filter_fsd;
use crate::utils::ma_io::{FIRE_TYPE, MSP_TYPE, NUC_TYPE};
use molecular_annotation::MolecularAnnotations;
use rayon::prelude::*;
use rust_htslib::bam::Read;
use rust_htslib::{bam, bam::ext::BamRecordExtensions, bam::record::Aux, bam::HeaderView};
use std::collections::HashMap;
use std::fmt::Write;

#[derive(Debug, Clone, PartialEq)]
pub struct FiberseqData {
    pub record: bam::Record,
    pub annotations: MolecularAnnotations,
    pub ec: f32,
    pub target_name: String,
    pub rg: String,
    pub center_position: Option<CenterPosition>,
}

impl FiberseqData {
    pub fn new(record: bam::Record, target_name: Option<&String>, filters: &FiberFilters) -> Self {
        // read group
        let rg = if let Ok(Aux::String(f)) = record.aux(b"RG") {
            log::trace!("{f}");
            f
        } else {
            "."
        }
        .to_string();
        let mut annotations = crate::utils::ma_io::read_record(&record).unwrap_or_else(|e| {
            log::warn!("Failed to read annotations: {e}");
            MolecularAnnotations::from_record(&record)
        });

        // The library populates m6a/cpg from MM/ML on read; apply the
        // fibertools read-side basemod filters (min ML score, end-strip
        // distance) on top of that here.
        if filters.min_ml_score > 0 || filters.strip_starting_basemods > 0 {
            let seq_len = record.seq_len();
            let strip = filters.strip_starting_basemods.max(0) as usize;
            let upper = seq_len.saturating_sub(strip);
            let min_ml = filters.min_ml_score;
            for t in annotations.annotation_types.iter_mut() {
                if !crate::utils::basemods::is_basemod_type(&t.name) {
                    continue;
                }
                t.annotations.retain(|a| {
                    if a.qualities.first().copied().unwrap_or(0) < min_ml {
                        return false;
                    }
                    if strip > 0 {
                        let p = a.start as usize;
                        if p < strip || p >= upper {
                            return false;
                        }
                    }
                    true
                });
            }
        }

        // get the number of passes
        let ec = if let Ok(Aux::Float(f)) = record.aux(b"ec") {
            log::trace!("{f}");
            f
        } else {
            0.0
        };

        let target_name = match target_name {
            Some(t) => t.clone(),
            None => ".".to_string(),
        };

        let mut fsd = FiberseqData {
            record,
            annotations,
            ec,
            target_name,
            rg,
            center_position: None,
        };

        apply_filter_fsd(&mut fsd, filters).expect("Failed to apply filter to FiberseqData");
        fsd
    }

    pub fn dict_from_head_view(head_view: &HeaderView) -> HashMap<i32, String> {
        if head_view.target_count() == 0 {
            return HashMap::new();
        }
        let target_u8s = head_view.target_names();
        let tids = target_u8s
            .iter()
            .map(|t| head_view.tid(t).expect("Unable to get tid"));
        let target_names = target_u8s
            .iter()
            .map(|&a| String::from_utf8_lossy(a).to_string());

        tids.zip(target_names)
            .map(|(id, t)| (id as i32, t))
            .collect()
    }

    pub fn target_name_from_tid(tid: i32, target_dict: &HashMap<i32, String>) -> Option<&String> {
        target_dict.get(&tid)
    }

    pub fn from_records(
        records: Vec<bam::Record>,
        head_view: &HeaderView,
        filters: &FiberFilters,
    ) -> Vec<Self> {
        let target_dict = Self::dict_from_head_view(head_view);
        records
            .into_par_iter()
            .map(|r| {
                let tid = r.tid();
                (r, Self::target_name_from_tid(tid, &target_dict))
            })
            .map(|(r, target_name)| Self::new(r, target_name, filters))
            .collect::<Vec<_>>()
    }

    //
    // GET FUNCTIONS
    //

    /// View over `msp` annotations derived from `self.annotations`.
    pub fn msp(&self) -> AnnotationTypeView<'_> {
        AnnotationTypeView::new(&self.annotations, MSP_TYPE)
    }

    /// View over `nuc` annotations derived from `self.annotations`.
    pub fn nuc(&self) -> AnnotationTypeView<'_> {
        AnnotationTypeView::new(&self.annotations, NUC_TYPE)
    }

    /// View over `m6a` annotations derived from `self.annotations`.
    pub fn m6a(&self) -> AnnotationTypeView<'_> {
        AnnotationTypeView::new(&self.annotations, M6A_TYPE)
    }

    /// View over `cpg` annotations derived from `self.annotations`.
    pub fn cpg(&self) -> AnnotationTypeView<'_> {
        AnnotationTypeView::new(&self.annotations, CPG_TYPE)
    }

    /// View over `fire` annotations derived from `self.annotations`.
    pub fn fire(&self) -> AnnotationTypeView<'_> {
        AnnotationTypeView::new(&self.annotations, FIRE_TYPE)
    }

    /// Flush `self.annotations` onto the record's MA-family aux tags. The
    /// single write path for subcommands that edit nuc/msp/fire annotations;
    /// call this, then hand the record to the BAM writer.
    ///
    /// MM/ML are deliberately left untouched: this is an edit path, not a
    /// basemod producer, so the record's original MM/ML bytes pass through
    /// byte-identically. Basemod types (m6a/cpg) are `Encoding::MmMl` (set when
    /// they were read), so `write_record` excludes them from the MA tag rather
    /// than writing basemod calls there too. Producers that synthesize or
    /// modify basemods must use `ma_io::write_record_with_basemods` instead.
    pub fn serialize_annotations(&mut self) {
        crate::utils::ma_io::write_record(&mut self.record, &self.annotations);
    }

    pub fn get_qname(&self) -> String {
        String::from_utf8_lossy(self.record.qname()).to_string()
    }

    pub fn get_rq(&self) -> Option<f32> {
        if let Ok(Aux::Float(f)) = self.record.aux(b"rq") {
            Some(f)
        } else {
            None
        }
    }

    /// Detect the sequencing platform (PacBio vs ONT) for this read from its
    /// aux tags, falling back to the read name if those were stripped.
    pub fn platform(&self) -> crate::utils::platform::SeqPlatform {
        crate::utils::platform::platform_from_record(&self.record)
    }

    pub fn get_hp(&self) -> String {
        match self.record.aux(b"HP") {
            Ok(Aux::U8(v)) => format!("H{v}"),
            Ok(Aux::I8(v)) => format!("H{v}"),
            Ok(Aux::U16(v)) => format!("H{v}"),
            Ok(Aux::I16(v)) => format!("H{v}"),
            Ok(Aux::U32(v)) => format!("H{v}"),
            Ok(Aux::I32(v)) => format!("H{v}"),
            _ => "UNK".to_string(),
        }
    }

    //
    //  WRITE BED12 FUNCTIONS
    //
    pub fn write_msp(&self, reference: bool) -> String {
        let msp = self.msp();
        let (starts, _ends, lengths) = if reference {
            (
                msp.reference_starts(),
                msp.reference_ends(),
                msp.reference_lengths(),
            )
        } else {
            (msp.option_starts(), msp.option_ends(), msp.option_lengths())
        };
        self.to_bed12(reference, &starts, &lengths, LINKER_COLOR)
    }

    pub fn write_nuc(&self, reference: bool) -> String {
        let nuc = self.nuc();
        let (starts, _ends, lengths) = if reference {
            (
                nuc.reference_starts(),
                nuc.reference_ends(),
                nuc.reference_lengths(),
            )
        } else {
            (nuc.option_starts(), nuc.option_ends(), nuc.option_lengths())
        };
        self.to_bed12(reference, &starts, &lengths, NUC_COLOR)
    }

    pub fn write_m6a(&self, reference: bool) -> String {
        let m6a = self.m6a();
        let starts = if reference {
            m6a.reference_starts()
        } else {
            m6a.option_starts()
        };
        let lengths = vec![Some(1); starts.len()];
        self.to_bed12(reference, &starts, &lengths, M6A_COLOR)
    }

    pub fn write_cpg(&self, reference: bool) -> String {
        let cpg = self.cpg();
        let starts = if reference {
            cpg.reference_starts()
        } else {
            cpg.option_starts()
        };
        let lengths = vec![Some(1); starts.len()];
        self.to_bed12(reference, &starts, &lengths, CPG_COLOR)
    }

    pub fn to_bed12(
        &self,
        reference: bool,
        starts: &[Option<i64>],
        lengths: &[Option<i64>],
        color: &str,
    ) -> String {
        if starts.is_empty() {
            return "".to_string();
        }
        // skip if no alignments are here
        if self.record.is_unmapped() && reference {
            return "".to_string();
        }

        let ct;
        let start;
        let end;
        let name = String::from_utf8_lossy(self.record.qname()).to_string();
        let mut rtn: String = String::with_capacity(0);
        if reference {
            ct = &self.target_name;
            start = self.record.reference_start();
            end = self.record.reference_end();
        } else {
            ct = &name;
            start = 0;
            end = self.record.seq_len() as i64;
        }
        let score = self.ec.round() as i64;
        let strand = if self.record.is_reverse() { '-' } else { '+' };
        // filter out positions that do not have an exact liftover
        let (filtered_starts, filtered_lengths): (Vec<i64>, Vec<i64>) = starts
            .iter()
            .flatten()
            .zip(lengths.iter().flatten())
            .unzip();
        // skip empty ones
        if filtered_lengths.is_empty() || filtered_starts.is_empty() {
            return "".to_string();
        }
        let b_ct = filtered_starts.len() + 2;
        let b_ln: String = filtered_lengths
            .iter()
            .map(|&ln| ln.to_string() + ",")
            .collect();
        let b_st: String = filtered_starts
            .iter()
            .map(|&st| (st - start).to_string() + ",")
            .collect();
        assert_eq!(filtered_lengths.len(), filtered_starts.len());

        rtn.push_str(ct);
        rtn.push('\t');
        rtn.push_str(&start.to_string());
        rtn.push('\t');
        rtn.push_str(&end.to_string());
        rtn.push('\t');
        rtn.push_str(&name);
        rtn.push('\t');
        rtn.push_str(&score.to_string());
        rtn.push('\t');
        rtn.push(strand);
        rtn.push('\t');
        rtn.push_str(&start.to_string());
        rtn.push('\t');
        rtn.push_str(&end.to_string());
        rtn.push('\t');
        rtn.push_str(color);
        rtn.push('\t');
        rtn.push_str(&b_ct.to_string());
        rtn.push_str("\t0,"); // add a zero length start
        rtn.push_str(&b_ln);
        rtn.push_str("1\t0,"); // add a 1 base length and a 0 start point
        rtn.push_str(&b_st);
        write!(&mut rtn, "{}", format_args!("{}\n", end - start - 1)).unwrap();
        rtn
    }

    //
    // WRITE ALL FUNCTIONS
    //

    pub fn all_header(simplify: bool, quality: bool) -> String {
        let mut x = format!(
            "#{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t",
            "ct", "st", "en", "fiber", "score", "strand", "sam_flag", "HP", "RG", "fiber_length",
        );
        if !simplify {
            x.push_str("fiber_sequence\t")
        }
        if quality {
            x.push_str("fiber_qual\t")
        }
        x.push_str(&format!(
            "{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\n",
            "ec",
            "rq",
            "total_AT_bp",
            "total_m6a_bp",
            "total_nuc_bp",
            "total_msp_bp",
            "total_5mC_bp",
            "nuc_starts",
            "nuc_lengths",
            "ref_nuc_starts",
            "ref_nuc_lengths",
            "msp_starts",
            "msp_lengths",
            "ref_msp_starts",
            "ref_msp_lengths",
            "fire_starts",
            "fire_lengths",
            "fire_qual",
            "ref_fire_starts",
            "ref_fire_lengths",
            "m6a",
            "ref_m6a",
            "m6a_qual",
            "5mC",
            "ref_5mC",
            "5mC_qual"
        ));
        x
    }

    pub fn write_all(&self, simplify: bool, quality: bool) -> String {
        // PB features
        let name = std::str::from_utf8(self.record.qname()).unwrap();
        let score = self.ec.round() as i64;
        let q_len = self.record.seq_len() as i64;
        let rq = match self.get_rq() {
            Some(x) => format!("{x}"),
            None => ".".to_string(),
        };
        // reference features
        let ct;
        let start;
        let end;
        let strand;
        if self.record.is_unmapped() {
            ct = ".";
            start = 0;
            end = 0;
            strand = '.';
        } else {
            ct = &self.target_name;
            start = self.record.reference_start();
            end = self.record.reference_end();
            strand = if self.record.is_reverse() { '-' } else { '+' };
        }
        let sam_flag = self.record.flags();
        let hp = self.get_hp();

        let at_count = self
            .record
            .seq()
            .as_bytes()
            .iter()
            .filter(|&x| *x == b'A' || *x == b'T')
            .count() as i64;

        // get the info
        let m6a = self.m6a();
        let cpg = self.cpg();
        let msp = self.msp();
        let nuc = self.nuc();
        let fire = self.fire();
        let m6a_count = m6a.len();
        let m6a_qual = m6a.qual().iter().map(|a| Some(*a as i64)).collect();
        let cpg_count = cpg.len();
        let cpg_qual = cpg.qual().iter().map(|a| Some(*a as i64)).collect();
        let fire_qual: Vec<Option<i64>> = fire.qual().iter().map(|a| Some(*a as i64)).collect();

        // write the features
        let mut rtn = String::with_capacity(0);
        // add first things 7
        rtn.write_fmt(format_args!(
            "{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t{}\t",
            ct, start, end, name, score, strand, sam_flag, hp, self.rg, q_len
        ))
        .unwrap();
        // add sequence
        if !simplify {
            rtn.write_fmt(format_args!(
                "{}\t",
                String::from_utf8_lossy(&self.record.seq().as_bytes()),
            ))
            .unwrap();
        }
        if quality {
            // TODO add quality offset
            rtn.write_fmt(format_args!(
                "{}\t",
                String::from_utf8_lossy(
                    &self
                        .record
                        .qual()
                        .iter()
                        .map(|x| x + 33)
                        .collect::<Vec<u8>>()
                ),
            ))
            .unwrap();
        }
        // add PB features
        let total_nuc_bp = nuc.lengths().iter().sum::<i64>();
        let total_msp_bp = msp.lengths().iter().sum::<i64>();
        rtn.write_fmt(format_args!(
            "{}\t{}\t{}\t{}\t{}\t{}\t{}\t",
            self.ec, rq, at_count, m6a_count, total_nuc_bp, total_msp_bp, cpg_count
        ))
        .unwrap();
        // add fiber features. FIRE is its own coord group (parallel to
        // nuc/m6a/cpg) — not wedged into the MSP columns.
        let vecs = [
            nuc.option_starts(),
            nuc.option_lengths(),
            nuc.reference_starts(),
            nuc.reference_lengths(),
            msp.option_starts(),
            msp.option_lengths(),
            msp.reference_starts(),
            msp.reference_lengths(),
            fire.option_starts(),
            fire.option_lengths(),
            fire_qual,
            fire.reference_starts(),
            fire.reference_lengths(),
            m6a.option_starts(),
            m6a.reference_starts(),
            m6a_qual,
            cpg.option_starts(),
            cpg.reference_starts(),
            cpg_qual,
        ];
        for vec in &vecs {
            if vec.is_empty() {
                rtn.push('.');
                rtn.push('\t');
            } else {
                let z: String = vec
                    .iter()
                    .map(|x| match x {
                        Some(y) => *y,
                        None => -1,
                    })
                    .map(|x| x.to_string() + ",")
                    .collect();
                rtn.write_fmt(format_args!("{z}\t")).unwrap();
            }
        }
        // replace the last tab with a newline
        let len = rtn.len();
        rtn.replace_range(len - 1..len, "\n");

        rtn
    }
}

pub struct FiberseqRecords<'a, R = bam::Reader>
where
    R: bam::Read,
{
    bam_chunk: BamChunk<'a, R>,
    header: HeaderView,
    filters: FiberFilters,
    cur_chunk: Vec<FiberseqData>,
}

impl<'a> FiberseqRecords<'a, bam::Reader> {
    pub fn new(bam: &'a mut bam::Reader, filters: FiberFilters) -> Self {
        let header = bam.header().clone();
        let bam_recs = bam.records();
        let mut bam_chunk = BamChunk::new(bam_recs, None);
        bam_chunk.set_bit_flag_filter(filters.get_bit_flag());
        let cur_chunk: Vec<FiberseqData> = vec![];
        FiberseqRecords {
            bam_chunk,
            header,
            filters,
            cur_chunk,
        }
    }
}

impl<'a> FiberseqRecords<'a, bam::IndexedReader> {
    pub fn from_rec_iterator(
        bam_recs: bam::Records<'a, bam::IndexedReader>,
        header: HeaderView,
        filters: FiberFilters,
    ) -> Self {
        let mut bam_chunk = BamChunk::new(bam_recs, None);
        bam_chunk.set_bit_flag_filter(filters.get_bit_flag());
        let cur_chunk: Vec<FiberseqData> = vec![];
        FiberseqRecords {
            bam_chunk,
            header,
            filters,
            cur_chunk,
        }
    }
}

impl<R> Iterator for FiberseqRecords<'_, R>
where
    R: bam::Read,
{
    type Item = FiberseqData;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            // if we are out of data check for another chunk in the bam
            if self.cur_chunk.is_empty() {
                match self.bam_chunk.next() {
                    Some(recs) => {
                        self.cur_chunk =
                            FiberseqData::from_records(recs, &self.header, &self.filters);
                        // we will be popping from this list so we want to remove the first element first, not the last
                        self.cur_chunk.reverse();
                    }
                    None => return None,
                }
            }
            let rec = self.cur_chunk.pop()?;
            if self.filters.passes_fire_filter(&rec) {
                return Some(rec);
            }
        }
    }
}