pileuphi_lib 0.9.3

use crate::errors::{Error, ErrorKind};
use crate::{
    alignment::PileupAlignment,
    output::{OrderedPileupOutput, PileupOutputContext},
    refseq::RefSeqHandle,
};
use indexmap::IndexMap;
use rust_htslib::bam::record::Cigar;
use std::ops::AddAssign;

#[derive(Clone, Debug)]
/// Writes itself into a tab-delimited histogram with counts for ACGTN, as well as gaps, and refskips. Indels are also counted, in the format <COUNT>[<INDEL_SEQUENCE>]
///
/// ```text
/// 1  2 3 4 5 6 7 8 9 10 11                                       12
/// |  | | | | | | | | |  |                                         |
/// c1 5 9 0 7 0 0 0 2 0 [2GTTAAC 1*TTAA* 1GTT*** 1***AAC 1**TA**] [1C 1CGG]
///
///  1: reference name
///  2: reference coordinate
///  3: depth (not including gaps)
///  4: #A's
///  5: #G's
///  6: #C's
///  7: #T's
///  8: #N's
///  9: #Gaps
///  10: #Refskips
///  11: space-delimited unique insertions*
///  12: space-delimited unique deletions*
///
///  In column 11, the first insertion has sequence GTTAAC and was observed 2x
/// ```
pub struct BaseDepthString {
    depth: u32,
    a: u32,
    g: u32,
    c: u32,
    t: u32,
    n: u32,
    gap: u32,
    refskip: u32,
    insertions: IndexMap<Vec<u8>, u32>,
    deletions: IndexMap<Vec<u8>, u32>,
}

unsafe impl Sync for BaseDepthString {}
unsafe impl Send for BaseDepthString {}

impl OrderedPileupOutput for BaseDepthString {
    #[inline(always)]
    fn intake(&mut self, ctx: &PileupOutputContext, p: &PileupAlignment) -> Result<(), Error> {
        self.depth += 1;
        self.register_pileup(p, ctx.refseq, ctx.pos)
    }

    #[inline(always)]
    fn write_header<W: std::io::Write>(ctx: &PileupOutputContext, writer: &mut W) -> Result<(), Error> {
        let mut buf = itoa::Buffer::new();

        writer.write_all(ctx.ref_name.as_bytes())?;
        writer.write_all(b"\t")?;

        writer.write_all(buf.format(ctx.pos + 1).as_bytes())?;
        writer.write_all(b"\t")?;
        Ok(())
    }

    #[inline(always)]
    fn write_body<W: std::io::Write>(&self, _ctx: &PileupOutputContext, writer: &mut W) -> Result<(), Error> {
        let mut buf = itoa::Buffer::new();
        writer.write_all(buf.format(self.depth).as_bytes())?;
        writer.write_all(b"\t")?;

        writer.write_all(buf.format(self.a).as_bytes())?;
        writer.write_all(b"\t")?;

        writer.write_all(buf.format(self.g).as_bytes())?;
        writer.write_all(b"\t")?;

        writer.write_all(buf.format(self.c).as_bytes())?;
        writer.write_all(b"\t")?;

        writer.write_all(buf.format(self.t).as_bytes())?;
        writer.write_all(b"\t")?;

        writer.write_all(buf.format(self.n).as_bytes())?;
        writer.write_all(b"\t")?;

        writer.write_all(buf.format(self.gap).as_bytes())?;
        writer.write_all(b"\t")?;

        writer.write_all(buf.format(self.refskip).as_bytes())?;

        writer.write_all(b"\t[")?;

        let n = self.insertions.len() - 1;
        for (i, (ins, cnt)) in self.insertions.iter().enumerate() {
            writer.write_all(buf.format(*cnt).as_bytes())?;
            writer.write_all(ins)?;
            if i < n {
                writer.write_all(b" ")?
            };
        }

        writer.write_all(b"]\t[")?;

        let n = self.deletions.len() - 1;
        for (i, (del, cnt)) in self.deletions.iter().enumerate() {
            writer.write_all(buf.format(*cnt).as_bytes())?;
            writer.write_all(del)?;
            if i < n {
                writer.write_all(b" ")?
            };
        }

        write!(writer, "]")?;
        Ok(())
    }

    #[inline(always)]
    fn write_body_empty<W: std::io::Write>(_ctx: &PileupOutputContext, writer: &mut W) -> Result<(), Error> {
        writer.write_all(b"0\t0\t0\t0\t0\t0\t0\t0\t[]\t[]")?;
        Ok(())
    }

    #[inline(always)]
    fn depth(&self) -> u32 {
        self.depth
    }

    #[inline(always)]
    fn clear(&mut self) {
        self.a = 0;
        self.g = 0;
        self.c = 0;
        self.t = 0;
        self.n = 0;
        self.depth = 0;
        self.gap = 0;
        self.refskip = 0;
        self.insertions.clear();
        self.deletions.clear();
    }

    fn new() -> Self {
        Self::new()
    }
}

#[allow(clippy::new_without_default)]
impl BaseDepthString {
    pub fn new() -> Self {
        Self {
            a: 0,
            g: 0,
            c: 0,
            t: 0,
            n: 0,
            depth: 0,
            gap: 0,
            refskip: 0,
            insertions: IndexMap::new(),
            deletions: IndexMap::new(),
        }
    }

    #[inline(always)]
    pub fn intake(&mut self, ctx: &PileupOutputContext, p: &PileupAlignment) -> Result<(), Error> {
        self.depth += 1;
        self.register_pileup(p, ctx.refseq, ctx.pos)
    }

    #[inline(always)]
    pub fn register_pileup(&mut self, p: &PileupAlignment, refseq: &RefSeqHandle, pos: i64) -> Result<(), Error> {
        match p.del {
            false => {
                let readbase = if p.qpos < p.rec.seq_len() {
                    p.rec.seq()[p.qpos]
                } else {
                    b'n'
                };

                match readbase.to_ascii_uppercase() {
                    b'A' => self.a += 1,
                    b'G' => self.g += 1,
                    b'C' => self.c += 1,
                    b'T' => self.t += 1,
                    b'N' => self.n += 1,
                    other => {
                        return Err(Error::from(ErrorKind::AnomalousData(format!(
                            "Unrecognized nucleotide character: {other}"
                        ))))
                    }
                }
            }

            true => {
                if p.refskip {
                    self.refskip += 1;
                } else {
                    self.gap += 1;
                }
            }
        }

        let mut del_len = -p.indel;

        if p.indel > 0 {
            let mut temp_buf: Vec<u8> = Vec::with_capacity(p.indel as usize);
            expand_insertions(p, &mut temp_buf, &mut del_len)?;
            self.insertions.entry(temp_buf).or_insert(0).add_assign(1);
        }

        if del_len > 0 {
            let mut temp_buf: Vec<u8> = Vec::with_capacity(del_len as usize);
            let mut refbase;

            for i in 1..=del_len as usize {
                refbase = if let Some(refseq) = refseq.as_ref() {
                    refseq[pos as usize + i]
                } else {
                    b'N'
                };
                temp_buf.push(refbase);
            }

            self.deletions.entry(temp_buf).or_insert(0).add_assign(1);
        }
        Ok(())
    }
}

#[inline(always)]
pub fn expand_insertions(p: &PileupAlignment, seq_buf: &mut Vec<u8>, ndel: &mut i32) -> Result<(), Error> {
    let mut read_pos: usize;
    let mut read_base: u8;
    let ncig = p.cstate.cig.len();

    // then produce the sequence representing the insertion
    let mut k = p.cigar_index + 1;
    let mut offset = 1;
    while k < ncig {
        match p.cstate.cig[k] {
            Cigar::Pad(l) => seq_buf.extend(std::iter::repeat_n(b'*', l as usize)),
            Cigar::Ins(l) => {
                for _ in 0..l as usize {
                    read_pos = p.qpos + offset - p.del as usize;
                    read_base = if read_pos < p.rec.seq_len() {
                        p.rec.seq()[read_pos]
                    } else {
                        b'n'
                    };
                    offset += 1;
                    seq_buf.push(read_base.to_ascii_uppercase());
                }
            }
            Cigar::Del(l) => {
                *ndel = l as i32;
                break;
            }
            _ => break,
        }
        k += 1;
    }

    Ok(())
}