pa-types 1.5.0

/*!
Cigar string types.

Uses SAM cigar conventions. See <https://timd.one/blog/genomics/cigar.php> (table below).

In particular, an "insertion" is a character in the pattern/query (2nd `Pos` coordinate) that is not in the text/reference (first `Pos` coordinate), and opposite for deletions.

We never output `M` characters, but do parse these into `=` and `X`, which need the corresponding text/pattern string to be resolved.

Undetermined bases (`N`), clipping (`S`, `H`), and padding (`P`) are not supported.

```text
+----------------+----------------+----------------------------------------------+---------------+---------------+
| Symbol         | Name           | Brief Description                            | Consumes Query| Consumes Ref  |
+----------------+----------------+----------------------------------------------+---------------+---------------+
| M              | Match          | No insertion or deletions, bases may differ  | ✓             | ✓             |
| I              | Insertion      | Additional base in query (not in reference)  | ✓             | ✗             |
| D              | Deletion       | Query is missing base from reference         | ✗             | ✓             |
| =              | Equal          | No insertions or deletions, bases agree      | ✓             | ✓             |
| X              | Not Equal      | No insertions or deletions, bases differ     | ✓             | ✓             |
| N              | None           | No query bases to align (spliced read)       | ✗             | ✓             |
| S              | Soft-Clipped   | Bases on end of read not aligned but stored  | ✓             | ✗             |
| H              | Hard-Clipped   | Bases on end of read not aligned, not stored | ✗             | ✗             |
| P              | Padding        | Neither read nor reference has a base        | ✗             | ✗             |
+----------------+----------------+----------------------------------------------+---------------+---------------+
```

*/

use itertools::Itertools;
use serde::{Deserialize, Serialize};
use std::fmt::Write;

use crate::*;

/// A single cigar string character.
#[derive(Serialize, Deserialize, Debug, PartialEq, Eq, Clone, Copy)]
pub enum CigarOp {
    /// `=`
    Match,
    /// `X`
    Sub,
    /// `D`
    Del,
    /// `I`
    Ins,
}

/// A cigar string character with the corresponding characters from text and pattern.
#[derive(Serialize, Deserialize, Debug, PartialEq, Eq, Clone, Copy)]
pub enum CigarOpChars {
    /// The text character of the matching pair.
    Match(u8),
    /// The mismatching text and pattern character, in this order.
    Sub(u8, u8),
    /// The text character that is missing from the pattern.
    Del(u8),
    /// The pattern character that is missing from the text.
    Ins(u8),
}

/// A single repeated cigar element, e.g. `5=` or `3I`.
#[derive(Serialize, Deserialize, Debug, PartialEq, Eq, Clone, Copy)]
pub struct CigarElem {
    pub op: CigarOp,
    pub cnt: I,
}

impl CigarElem {
    pub fn new(op: CigarOp, cnt: I) -> Self {
        Self { op, cnt }
    }
}

/// Main type representing a Cigar string.
///
/// Adjacent equal operations are typically merged into [`CigarElem`] with `cnt>1`.
// This is similar to https://docs.rs/bio/1.0.0/bio/alignment/struct.Alignment.html,
// but more specific for our use case.
#[derive(Serialize, Deserialize, Debug, PartialEq, Eq, Clone, Default)]
pub struct Cigar {
    pub ops: Vec<CigarElem>,
}

impl CigarOp {
    /// Convert to one of `=XID`.
    pub fn to_char(&self) -> char {
        match self {
            CigarOp::Match => '=',
            CigarOp::Sub => 'X',
            CigarOp::Ins => 'I',
            CigarOp::Del => 'D',
        }
    }

    /// Convert operation to `(text_pos, pattern_pos)` delta/step of path indices.
    #[inline(always)]
    pub fn delta(&self) -> Pos {
        match self {
            CigarOp::Match | CigarOp::Sub => Pos(1, 1),
            CigarOp::Del => Pos(1, 0), // deletion: consumes ref/text only
            CigarOp::Ins => Pos(0, 1), // insertion: consumes pattern/query only
        }
    }

    /// The edit-distance cost of the operation.
    ///
    /// 0 for `Match`, 1 otherwise.
    #[inline(always)]
    pub fn edit_cost(&self) -> Cost {
        match self {
            CigarOp::Match => 0,
            _ => 1,
        }
    }

    /// Converts path delta `(text_pos, pattern_pos)` to `CigarOp`.
    ///
    /// Ignores [`CigarOp::Sub`]: (1,1) is always [`CigarOp::Match`].
    #[inline(always)]
    pub fn from_delta(delta: Pos) -> Self {
        match delta {
            Pos(0, 1) => CigarOp::Ins,
            Pos(1, 0) => CigarOp::Del,
            Pos(1, 1) => CigarOp::Match,
            _ => panic!("Invalid delta: {:?}", delta),
        }
    }
}

// Mostly to make verify/resolve matches clearer while still using the delta mapping
impl std::ops::Mul<I> for Pos {
    type Output = Pos;
    fn mul(self, rhs: I) -> Pos {
        Pos(self.0 * rhs, self.1 * rhs)
    }
}

impl From<u8> for CigarOp {
    /// Convert from `=MXID` to `CigarOp`.
    ///
    /// `M` is interpreted as `=` ([`CigarOp::Match`]) and should be resolved into `=` or `X` ([`CigarOp::Sub`]) via `resolve_matches`.
    fn from(op: u8) -> Self {
        match op {
            b'=' | b'M' => CigarOp::Match,
            b'X' => CigarOp::Sub,
            b'I' => CigarOp::Ins,
            b'D' => CigarOp::Del,
            _ => panic!("Invalid CigarOp"),
        }
    }
}

impl ToString for Cigar {
    /// Format the cigar to a `String`.
    fn to_string(&self) -> String {
        let mut s = String::new();
        for elem in &self.ops {
            write!(&mut s, "{}{}", elem.cnt, elem.op.to_char()).unwrap();
        }
        s
    }
}

impl Cigar {
    /// Concatenate the single-characters ops.
    pub fn from_ops(ops: impl Iterator<Item = CigarOp>) -> Self {
        Cigar {
            ops: ops
                .chunk_by(|&op| op)
                .into_iter()
                .map(|(op, group)| CigarElem::new(op, group.count() as _))
                .collect(),
        }
    }

    /// Create Cigar from path and corresponding sequences.
    ///
    /// Path must have `(text_pos, pattern_pos)` pairs.
    /// To distinguish between match and sub it uses simple
    /// equality (i.e. c1==c2) via `resolve_matches`, so IUPAC matching is not supported.
    pub fn from_path(text: Seq, pattern: Seq, path: &Path) -> Cigar {
        if path[0] != Pos(0, 0) {
            panic!("Path must start at (0,0)!");
        }
        Self::resolve_matches(
            path.iter()
                .tuple_windows()
                .map(|(&text_pos, &pattern_pos)| {
                    CigarElem::new(CigarOp::from_delta(pattern_pos - text_pos), 1)
                }),
            text,
            pattern,
        )
    }

    /// Return the diff from pattern to text.
    pub fn to_char_pairs<'s>(&'s self, text: &'s [u8], pattern: &'s [u8]) -> Vec<CigarOpChars> {
        let mut pos = Pos(0, 0);
        // let fix_case = !(b'A' ^ b'a');
        let mut out = vec![];
        for el in &self.ops {
            for _ in 0..el.cnt {
                let c = match el.op {
                    CigarOp::Match => {
                        // NOTE: IUPAC characters can be matching even when they're not equal.
                        // assert_eq!(
                        //     (pattern[pos.0 as usize] & fix_case) as char,
                        //     (text[pos.1 as usize] & fix_case) as char,
                        //     "mismatch for {pos:?}"
                        // );
                        CigarOpChars::Match(text[pos.0 as usize])
                    }
                    CigarOp::Sub => {
                        // NOTE: ASCII characters can be mismatching when only differing in case.
                        assert_ne!(
                            text[pos.0 as usize] as char,
                            pattern[pos.1 as usize] as char,
                            "cigar {:?}\npattern {:?}\ntext    {:?}\nmismatch for {pos:?}",
                            self.to_string(),
                            String::from_utf8_lossy(pattern),
                            String::from_utf8_lossy(text)
                        );
                        CigarOpChars::Sub(text[pos.0 as usize], pattern[pos.1 as usize])
                    }
                    CigarOp::Del => {
                        // Note deletion consumes text hence text slice
                        CigarOpChars::Del(text[pos.0 as usize])
                    }
                    CigarOp::Ins => {
                        // Note insertion consumes pattern hence pattern slice
                        CigarOpChars::Ins(pattern[pos.1 as usize])
                    }
                };
                out.push(c);
                pos += el.op.delta();
            }
        }
        out
    }

    /// Get the `Path` corresponding to this [`Cigar`].
    pub fn to_path(&self) -> Path {
        let mut pos = Pos(0, 0);
        let mut path = vec![pos];
        for el in &self.ops {
            for _ in 0..el.cnt {
                pos += el.op.delta();
                path.push(pos);
            }
        }
        path
    }

    /// Get the `Path` and the alignment `Cost` to each position.
    pub fn to_path_with_costs(&self, cm: CostModel) -> Vec<(Pos, Cost)> {
        let mut pos = Pos(0, 0);
        let mut cost = 0;
        let mut path = vec![(pos, cost)];

        for el in &self.ops {
            match el.op {
                CigarOp::Match => {
                    for _ in 0..el.cnt {
                        pos += el.op.delta();
                        path.push((pos, cost));
                    }
                }
                CigarOp::Sub => {
                    for _ in 0..el.cnt {
                        pos += el.op.delta();
                        cost += cm.sub;
                        path.push((pos, cost));
                    }
                }
                CigarOp::Ins => {
                    for len in 1..=(el.cnt as Cost) {
                        pos += el.op.delta();
                        path.push((pos, cost + cm.ins(len)));
                    }
                    cost += cm.ins(el.cnt);
                }
                CigarOp::Del => {
                    for len in 1..=(el.cnt as Cost) {
                        pos += el.op.delta();
                        path.push((pos, cost + cm.del(len)));
                    }
                    cost += cm.del(el.cnt);
                }
            }
        }
        path
    }

    /// Push a [`CigarOp`] to the cigar.
    pub fn push(&mut self, op: CigarOp) {
        if let Some(s) = self.ops.last_mut() {
            if s.op == op {
                s.cnt += 1;
                return;
            }
        }
        self.ops.push(CigarElem { op, cnt: 1 });
    }

    /// Pop a single [`CigarOp`] from the cigar.
    pub fn pop_op(&mut self) -> Option<CigarOp> {
        while let Some(elem) = self.ops.last_mut() {
            let op = elem.op;
            assert!(elem.cnt > 0);
            elem.cnt -= 1;
            if elem.cnt == 0 {
                self.ops.pop();
            }
            return Some(op);
        }
        None
    }

    /// Push a [`CigarElem`] to the cigar.
    pub fn push_elem(&mut self, e: CigarElem) {
        if let Some(s) = self.ops.last_mut() {
            if s.op == e.op {
                s.cnt += e.cnt;
                return;
            }
        }
        self.ops.push(e);
    }

    /// Push `cnt` matches.
    pub fn push_matches(&mut self, cnt: I) {
        if let Some(s) = self.ops.last_mut() {
            if s.op == CigarOp::Match {
                s.cnt += cnt;
                return;
            }
        }
        self.ops.push(CigarElem {
            op: CigarOp::Match,
            cnt: cnt as _,
        });
    }

    /// Check that the cigar is valid between `text` and `pattern` and return the cost.
    pub fn verify(&self, cm: &CostModel, text: Seq, pattern: Seq) -> Result<Cost, &str> {
        let mut pos = Pos(0, 0);
        let mut cost: Cost = 0;

        for &CigarElem { op, cnt } in &self.ops {
            match op {
                CigarOp::Match => {
                    for _ in 0..cnt {
                        if text.get(pos.0 as usize) != pattern.get(pos.1 as usize) {
                            return Err("Expected match but found substitution.");
                        }
                        pos += op.delta();
                    }
                }
                CigarOp::Sub => {
                    for _ in 0..cnt {
                        if text.get(pos.0 as usize) == pattern.get(pos.1 as usize) {
                            return Err("Expected substitution but found match.");
                        }
                        pos += op.delta();
                        cost += cm.sub;
                    }
                }
                CigarOp::Ins => {
                    cost += cm.open + cnt as Cost * cm.extend;
                    pos += op.delta() * cnt;
                }
                CigarOp::Del => {
                    cost += cm.open + cnt as Cost * cm.extend;
                    pos += op.delta() * cnt;
                }
            }
        }
        if pos != Pos(text.len() as I, pattern.len() as I) {
            return Err("Wrong alignment length.");
        }

        Ok(cost)
    }

    /// Splits all 'M'/[`CigarOp::Match`] into matches (`=`) and substitutions (`X`), and joins consecutive equal elements.
    pub fn resolve_matches(ops: impl Iterator<Item = CigarElem>, text: Seq, pattern: Seq) -> Self {
        let mut pos = Pos(0, 0);
        let mut c = Cigar { ops: vec![] };
        for CigarElem { op, cnt } in ops {
            match op {
                CigarOp::Match => {
                    for _ in 0..cnt {
                        c.push(if text[pos.0 as usize] == pattern[pos.1 as usize] {
                            CigarOp::Match
                        } else {
                            CigarOp::Sub
                        });
                        pos += op.delta();
                    }
                    continue;
                }
                _ => {
                    pos += op.delta() * cnt;
                }
            };
            c.push_elem(CigarElem { op, cnt });
        }
        c
    }

    /// A simpler parsing function that only parses strings of characters `M=XID`, without preceding counts.
    /// Consecutive characters are grouped, and `M` and `=` chars are resolved into `=` and `X`.
    pub fn parse_without_counts(s: &str, text: Seq, pattern: Seq) -> Self {
        Self::resolve_matches(
            s.as_bytes().iter().map(|&op| CigarElem {
                op: op.into(),
                cnt: 1,
            }),
            text,
            pattern,
        )
    }

    /// A simpler parsing function that only parses strings of characters `MXID`, without preceding counts.
    /// Consecutive characters are grouped. `M` chars are *not* resolved and assumed to mean `=`.
    pub fn parse_without_resolving(s: &str) -> Self {
        let mut c = Cigar { ops: vec![] };
        for &op in s.as_bytes() {
            c.push(op.into())
        }
        c
    }

    /// Parse a Cigar string with optional counts
    pub fn from_string(s: &str) -> Self {
        let mut c = Cigar { ops: vec![] };
        for slice in s.as_bytes().split_inclusive(|b| !b.is_ascii_digit()) {
            let (&op, cnt_bytes) = slice.split_last().expect("Cigar string cannot be empty");
            let cnt = if cnt_bytes.is_empty() {
                1
            } else {
                unsafe { std::str::from_utf8_unchecked(cnt_bytes) }
                    .parse()
                    .expect("Invalid Cigar count")
            };
            c.push_elem(CigarElem { op: op.into(), cnt });
        }
        c
    }

    /// A more generic (and slower) parsing function that also allows optional counts, e.g. `5M2X3M`.
    /// Consecutive characters are grouped, and `M` and `=` chars are resolved into `=` and `X`.
    pub fn parse(s: &str, text: Seq, pattern: Seq) -> Self {
        Self::resolve_matches(
            s.as_bytes()
                .split_inclusive(|pattern| pattern.is_ascii_alphabetic())
                .map(|pattern_slice| {
                    let (&op, cnt) = pattern_slice.split_last().unwrap();
                    let cnt = if cnt.is_empty() {
                        1
                    } else {
                        unsafe { std::str::from_utf8_unchecked(cnt) }
                            .parse()
                            .unwrap()
                    };
                    CigarElem { op: op.into(), cnt }
                }),
            text,
            pattern,
        )
    }

    /// Clear the internal vector.
    pub fn clear(&mut self) {
        self.ops.clear();
    }

    /// Reverse the cigar string, for `rc(text)` and `rc(pattern)`.
    pub fn reverse(&mut self) {
        self.ops.reverse();
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_delta() {
        for op in [CigarOp::Match, CigarOp::Del, CigarOp::Ins] {
            assert_eq!(CigarOp::from_delta(op.delta()), op);
        }
        assert_eq!(CigarOp::from_delta(Pos(1, 1)), CigarOp::Match);
        //  assert_eq!(CigarOp::from_delta(Pos(1, 1)), CigarOp::Sub);
        assert_eq!(CigarOp::from_delta(Pos(1, 0)), CigarOp::Del); // Consume text
        assert_eq!(CigarOp::from_delta(Pos(0, 1)), CigarOp::Ins); // Consume pattern
    }

    #[test]
    fn test_valid_eq() {
        let c = Cigar::from_path(b"ab", b"aa", &vec![Pos(0, 0), Pos(1, 1), Pos(2, 2)]);
        assert_eq!(c.to_string(), "1=1X");
    }

    #[test]
    fn test_invalid_end_length() {
        // (0,1)is invalid, text is longer than path
        let c = Cigar::from_path(b"ab", b"aa", &vec![Pos(0, 0), Pos(0, 1)]);
        assert!(c.verify(&CostModel::unit(), b"ab", b"aa").is_err());
    }

    #[test]
    fn to_string() {
        let c = Cigar {
            ops: vec![
                CigarElem {
                    op: CigarOp::Ins,
                    cnt: 1,
                },
                CigarElem {
                    op: CigarOp::Match,
                    cnt: 2,
                },
            ],
        };
        assert_eq!(c.to_string(), "1I2=");
    }

    #[test]
    fn from_path() {
        let c = Cigar::from_path(
            b"aaa",
            b"aabc",
            &vec![Pos(0, 0), Pos(1, 1), Pos(2, 2), Pos(3, 3), Pos(3, 4)],
        );
        assert_eq!(c.to_string(), "2=1X1I");
    }

    #[test]
    fn from_string_with_count() {
        let c = Cigar::from_string("24=");
        assert_eq!(c.ops.len(), 1);
        assert_eq!(c.ops[0], CigarElem::new(CigarOp::Match, 24));
    }

    #[test]
    fn from_string_mixed_ops() {
        let c = Cigar::from_string("2=3I1X");
        assert_eq!(
            c.ops,
            vec![
                CigarElem::new(CigarOp::Match, 2),
                CigarElem::new(CigarOp::Ins, 3),
                CigarElem::new(CigarOp::Sub, 1)
            ]
        );
    }

    #[test]
    fn from_string_no_counts() {
        let c = Cigar::from_string("=XIDDD");
        assert_eq!(
            c.ops,
            vec![
                CigarElem::new(CigarOp::Match, 1),
                CigarElem::new(CigarOp::Sub, 1),
                CigarElem::new(CigarOp::Ins, 1),
                CigarElem::new(CigarOp::Del, 3),
            ]
        );
    }

    #[test]
    #[rustfmt::skip]
    fn push_to_path() {
        let mut c = Cigar::default();
                                // 0 0
        c.push(CigarOp::Match); // 1  1
        c.push(CigarOp::Del);   // 2  1  (text +1)
        c.push(CigarOp::Ins);   // 2  2  (pattern +1)
        c.push(CigarOp::Sub);   // 3  3

        assert_eq!(
            c.to_path(),
            [
                Pos(0, 0),
                Pos(1, 1),
                Pos(2, 1),
                Pos(2, 2),
                Pos(3, 3),
            ]
        );
    }

    #[test]
    fn to_char_pairs_all_match() {
        let c = Cigar::from_string("3=");
        let pairs = c.to_char_pairs(b"aaa", b"aaa");
        assert_eq!(
            pairs,
            vec![
                CigarOpChars::Match(b'a'),
                CigarOpChars::Match(b'a'),
                CigarOpChars::Match(b'a'),
            ]
        );
    }

    #[test]
    fn to_char_pairs_sub() {
        let c = Cigar::from_string("1X");
        let pairs = c.to_char_pairs(b"a", b"c");
        assert_eq!(pairs, vec![CigarOpChars::Sub(b'a', b'c')]);
    }

    #[test]
    fn to_char_pairs_ins() {
        let c = Cigar::from_string("1=1I1=");
        let pairs = c.to_char_pairs(b"ac", b"abc");
        assert_eq!(
            pairs,
            vec![
                CigarOpChars::Match(b'a'),
                CigarOpChars::Ins(b'b'),
                CigarOpChars::Match(b'c'),
            ]
        );
    }

    #[test]
    fn to_char_pairs_del() {
        let c = Cigar::from_string("1=1D1=");
        let pairs = c.to_char_pairs(b"abc", b"ac");
        assert_eq!(
            pairs,
            vec![
                CigarOpChars::Match(b'a'),
                CigarOpChars::Del(b'b'),
                CigarOpChars::Match(b'c'),
            ]
        );
    }

    #[test]
    fn to_char_pairs_mixed() {
        let c = Cigar::from_string("2=1X1I");
        let pairs = c.to_char_pairs(b"abZd", b"abYc");
        assert_eq!(
            pairs,
            vec![
                CigarOpChars::Match(b'a'),
                CigarOpChars::Match(b'b'),
                CigarOpChars::Sub(b'Z', b'Y'),
                CigarOpChars::Ins(b'c'),
            ]
        );
    }
}