twitcher 0.6.8

use std::sync::Arc;

use anyhow::Context;
use rust_htslib::bam::{HeaderView, Record, record::Cigar};
use tracing::error;

use crate::common::{
    ImmutableSequence,
    cluster_settings::ClusteringSettings,
    contig::ContigName,
    coords::{GenomePosition, GenomeRegion},
    reference::ReferenceReader,
};

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Cluster {
    pub region: GenomeRegion,
    pub read_range: std::ops::Range<usize>,
}

#[derive(Debug)]
struct WorkingCluster {
    ref_start: usize,
    ref_end: usize,
    read_start: usize,
    read_end: usize,
    complexity: f64,
    event_count: usize,
    matches_since_last_mut: usize,
}

impl WorkingCluster {
    const fn new(ref_pos: usize, read_pos: usize) -> Self {
        Self {
            ref_start: ref_pos,
            ref_end: ref_pos,
            read_start: read_pos,
            read_end: read_pos,
            complexity: 0.0,
            event_count: 0,
            matches_since_last_mut: 0,
        }
    }

    fn is_valid(&self, settings: &ClusteringSettings) -> bool {
        // Reference and query (read) extents, in bases covered: both `*_end` fields are the
        // exclusive end of the last event. The strategy decides whether the span is the reference
        // extent alone or the larger of the two.
        let ref_extent = self.ref_end.abs_diff(self.ref_start) as f64;
        let query_extent = self.read_end.abs_diff(self.read_start) as f64;
        let span = settings.span(ref_extent, query_extent);
        settings.is_valid_cluster(self.complexity, span, self.event_count)
    }
}

#[derive(Debug)]
struct ClusterFinder {
    record: Arc<Record>,
    ref_seq: ImmutableSequence,
    chr: ContigName,
    settings: ClusteringSettings,
    ref_start: usize,
    ref_pos: usize,
    read_pos: usize,
    expected_read_len: usize,
    current: Option<WorkingCluster>,
    record_cigar: Vec<Cigar>,
    clusters: Vec<Cluster>,
}

impl ClusterFinder {
    fn new(
        record: Arc<Record>,
        ref_seq: ImmutableSequence,
        chr: ContigName,
        settings: ClusteringSettings,
    ) -> anyhow::Result<Self> {
        let ref_pos = usize::try_from(record.pos())?;
        Ok(Self {
            record,
            ref_seq,
            chr,
            settings,
            ref_start: ref_pos,
            ref_pos,
            read_pos: 0,
            expected_read_len: 0,
            current: None,
            record_cigar: Vec::new(),
            clusters: Vec::new(),
        })
    }

    fn find(mut self) -> (Vec<Cigar>, Vec<Cluster>) {
        for cig in &self.record.cigar() {
            match *cig {
                Cigar::Match(n) => {
                    let mut n = n as usize;
                    while n > 0 {
                        let qry_seq = self.record.seq();
                        let ref_char = self.ref_seq[self.ref_pos - self.ref_start];
                        let qry_char = qry_seq[self.read_pos];
                        let is_equal = qry_char == ref_char;

                        let mut run_len = 1;
                        while run_len < n {
                            let r = self.ref_seq[self.ref_pos + run_len - self.ref_start];
                            let q = qry_seq[self.read_pos + run_len];
                            if (q == r) != is_equal {
                                break;
                            }
                            run_len += 1;
                        }

                        if is_equal {
                            self.consume_equal(run_len);
                            self.record_cigar.push(Cigar::Equal(run_len as u32));
                        } else {
                            self.consume_diff(run_len);
                            self.record_cigar.push(Cigar::Diff(run_len as u32));
                        }
                        n -= run_len;
                    }
                }
                Cigar::Equal(n) => self.consume_equal(n as usize),
                Cigar::Diff(n) => self.consume_diff(n as usize),
                Cigar::Ins(n) => self.consume_ins(n as usize),
                Cigar::Del(n) => self.consume_del(n as usize),
                Cigar::RefSkip(n) => self.consume_refskip(n as usize),
                Cigar::SoftClip(n) => self.consume_softclip(n as usize),
                Cigar::HardClip(_) | Cigar::Pad(_) => self.flush(),
            }

            // We split Match into Equal and Diff earlier, but all other alignments need to be pushed.
            if !matches!(*cig, Cigar::Match(..)) {
                self.record_cigar.push(*cig);
            }
        }
        self.flush();
        (self.record_cigar, self.clusters)
    }

    fn consume_equal(&mut self, n: usize) {
        if let Some(c) = self.current.as_mut() {
            if c.matches_since_last_mut + n > self.settings.max_gap {
                self.flush();
            } else {
                c.matches_since_last_mut += n;
            }
        }
        self.expected_read_len += n;
        self.ref_pos += n;
        self.read_pos += n;
    }

    /// Registers one mutation event of the given `mass` that covers `ref_len` reference and
    /// `read_len` read bases starting at the current position, opening a cluster if none is.
    ///
    /// The cluster boundaries jump to the end of the event rather than growing with every matched
    /// base, so matches between two events are folded into the span while matches that turn out to
    /// be trailing never extend it. This mirrors the VCF path, where the span runs from the first
    /// record's position to the last record's end.
    fn add_event(&mut self, mass: f64, ref_len: usize, read_len: usize) {
        let (ref_pos, read_pos) = (self.ref_pos, self.read_pos);
        let c = self
            .current
            .get_or_insert_with(|| WorkingCluster::new(ref_pos, read_pos));
        c.complexity += mass;
        c.event_count += 1;
        c.matches_since_last_mut = 0;
        c.ref_end = ref_pos + ref_len;
        c.read_end = read_pos + read_len;
    }

    fn consume_diff(&mut self, n: usize) {
        // A contiguous mismatch run is one event: an n-base substitution block, scored exactly
        // like the equivalent n-base MNV record in the VCF path.
        let mass = self.settings.event_mass(n, n);
        self.add_event(mass, n, n);

        self.expected_read_len += n;
        self.ref_pos += n;
        self.read_pos += n;
    }

    fn consume_ins(&mut self, n: usize) {
        // VCF anchor convention: an n-base insertion is `(1, n + 1)`.
        let mass = self.settings.event_mass(1, n + 1);
        self.add_event(mass, 0, n);

        self.expected_read_len += n;
        self.read_pos += n;
    }

    fn consume_del(&mut self, n: usize) {
        // VCF anchor convention: an n-base deletion is `(n + 1, 1)`.
        let mass = self.settings.event_mass(n + 1, 1);
        self.add_event(mass, n, 0);

        self.ref_pos += n;
    }

    fn consume_softclip(&mut self, n: usize) {
        self.flush();
        self.expected_read_len += n;
        self.read_pos += n;
    }

    fn consume_refskip(&mut self, n: usize) {
        self.flush();
        self.ref_pos += n;
    }

    fn flush(&mut self) {
        if let Some(c) = self.current.take()
            && c.is_valid(&self.settings)
        {
            let start = GenomePosition::new_0(self.chr.clone(), c.ref_start);
            let end = GenomePosition::new_0(self.chr.clone(), c.ref_end);
            self.clusters.push(Cluster {
                region: GenomeRegion::from_incl_excl(start, Some(end)).unwrap(),
                read_range: c.read_start..c.read_end,
            });
        }
    }
}

pub async fn find_cigar_clusters(
    header: &HeaderView,
    record: Arc<Record>,
    reference_reader: &ReferenceReader,
    settings: ClusteringSettings,
) -> anyhow::Result<(Vec<Cigar>, Vec<Cluster>)> {
    let chr = u32::try_from(record.tid())
        .ok()
        .map(|tid| ContigName::new(header.tid2name(tid)))
        .context("Cannot find clusters without target id (can't extract reference)")?;
    let rec_reg = GenomeRegion::try_from_bam_record(&record, chr.clone())?;
    let ref_seq = reference_reader.get_seq_exact_unmasked(rec_reg).await?;
    let finder = ClusterFinder::new(record.clone(), ref_seq, chr, settings)
        .inspect_err(|e| error!("{e}"))?;
    Ok(finder.find())
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use crate::common::cluster_settings::{ClusterStrategy, ClusteringSettings};
    use rust_htslib::bam::{
        Record,
        record::{Cigar, CigarString},
    };

    use super::{Cluster, ClusterFinder};

    /// Edit-mass strategy with default thresholds, for the cases that only cluster under it.
    fn edit_mass_settings() -> ClusteringSettings {
        ClusteringSettings {
            strategy: ClusterStrategy::EditMass,
            ..Default::default()
        }
    }

    /// Runs `ClusterFinder` directly on a record built from the given CIGAR and a dummy reference
    /// sequence. Only `Equal`, `Diff`, `Ins`, and `Del` ops are used — these never index into the
    /// reference buffer, so the content of `ref_seq` is irrelevant.
    fn clusters_from_cigar(
        cigar: Vec<Cigar>,
        pos: i64,
        settings: ClusteringSettings,
    ) -> Vec<Cluster> {
        let cigar_string = CigarString(cigar);
        let read_len: usize = cigar_string
            .iter()
            .map(|op| match op {
                Cigar::Equal(n) | Cigar::Diff(n) | Cigar::Ins(n) | Cigar::SoftClip(n) => {
                    *n as usize
                }
                _ => 0,
            })
            .sum();
        let ref_len: usize = cigar_string
            .iter()
            .map(|op| match op {
                Cigar::Equal(n)
                | Cigar::Diff(n)
                | Cigar::Del(n)
                | Cigar::RefSkip(n)
                | Cigar::Match(n) => *n as usize,
                _ => 0,
            })
            .sum();

        let mut record = Record::new();
        let seq = vec![b'A'; read_len];
        let qual = vec![40u8; read_len];
        record.set(b"test", Some(&cigar_string), &seq, &qual);
        record.set_pos(pos);
        let record = Arc::new(record);

        let chr = crate::common::contig::ContigName::new(b"chr1");
        let ref_seq: crate::common::ImmutableSequence = vec![b'A'; ref_len].into();

        let finder = ClusterFinder::new(record, ref_seq, chr, settings).unwrap();
        let (_cigar, clusters) = finder.find();
        clusters
    }

    #[test]
    fn two_diff_runs_close_together_form_cluster() {
        // 5= 2X 3= 2X 5=: two mismatch events separated by 3 matches (within max_gap)
        let cigar = vec![
            Cigar::Equal(5),
            Cigar::Diff(2),
            Cigar::Equal(3),
            Cigar::Diff(2),
            Cigar::Equal(5),
        ];
        let clusters = clusters_from_cigar(cigar, 0, ClusteringSettings::default());
        assert_eq!(clusters.len(), 1);
        let c = &clusters[0];
        // Cluster begins at the first mismatch (ref offset 5 from pos=0)
        assert_eq!(c.region.start().position_0(), 5);
        // ... and ends at the last one: the 3 matches in between are part of the cluster, the
        // trailing 5 are not.
        assert_eq!(c.region.end_excl().unwrap().position_0(), 12);
        assert_eq!(c.read_range, 5..12);
    }

    #[test]
    fn trailing_matches_do_not_extend_the_cluster() {
        // 2X 3= 2X 30=: the 30 matches break the cluster, and none of them may end up inside it.
        let cigar = vec![
            Cigar::Diff(2),
            Cigar::Equal(3),
            Cigar::Diff(2),
            Cigar::Equal(30),
        ];
        let clusters = clusters_from_cigar(cigar, 0, ClusteringSettings::default());
        assert_eq!(clusters.len(), 1);
        let c = &clusters[0];
        assert_eq!(c.region.start().position_0(), 0);
        assert_eq!(c.region.end_excl().unwrap().position_0(), 7);
        assert_eq!(c.read_range, 0..7);
    }

    #[test]
    fn mismatch_run_is_scored_like_an_mnv_under_legacy() {
        // 4X 5= 4X: two 4-base substitution blocks over a span of 13.
        // Legacy scores each block like a 4-base MNV: 2 * (1 + log2(3)) = 5.17 → density 0.398.
        // Scoring per base would give 8.0 → density 0.615, so a 0.5 threshold separates the two.
        let cigar = vec![Cigar::Diff(4), Cigar::Equal(5), Cigar::Diff(4)];
        let settings = ClusteringSettings {
            min_density: 0.5,
            ..Default::default()
        };
        assert_eq!(clusters_from_cigar(cigar.clone(), 0, settings).len(), 0);

        // Edit mass is linear in the edited bases, so the same blocks give 8.0 and pass.
        let settings = ClusteringSettings {
            min_density: 0.5,
            ..edit_mass_settings()
        };
        assert_eq!(clusters_from_cigar(cigar, 0, settings).len(), 1);
    }

    #[test]
    fn diff_runs_separated_by_overlong_gap_are_rejected() {
        // 1X 25= 1X: the 25-match gap exceeds max_gap=20, breaking the cluster each time.
        // Each resulting fragment is a single event — below the legacy default min_records=2 →
        // both discarded.
        let cigar = vec![Cigar::Diff(1), Cigar::Equal(25), Cigar::Diff(1)];
        let clusters = clusters_from_cigar(cigar, 0, ClusteringSettings::default());
        assert_eq!(clusters.len(), 0);
    }

    #[test]
    fn insertion_between_diffs_contributes_complexity() {
        // 2= 1X 1I 1X 2= at pos=10: SNP, 1-base insertion, SNP — all within gap
        let cigar = vec![
            Cigar::Equal(2),
            Cigar::Diff(1),
            Cigar::Ins(1),
            Cigar::Diff(1),
            Cigar::Equal(2),
        ];
        let clusters = clusters_from_cigar(cigar, 10, ClusteringSettings::default());
        assert_eq!(clusters.len(), 1);
        let c = &clusters[0];
        // Cluster starts two bases into the read (after the leading 2=), at absolute ref pos 12
        assert_eq!(c.region.start().position_0(), 12);
        assert_eq!(c.read_range, 2..5);
    }

    #[test]
    fn single_event_rejected_below_legacy_min_records() {
        // 3= 1X 3=: a lone event, below the legacy default min_records=2 → rejected.
        let cigar = vec![Cigar::Equal(3), Cigar::Diff(1), Cigar::Equal(3)];
        let clusters = clusters_from_cigar(cigar, 0, ClusteringSettings::default());
        assert_eq!(clusters.len(), 0);
    }

    #[test]
    fn single_snp_event_rejected_below_min_mass() {
        // edit-mass: 3= 1X 3= is a lone SNP with edit mass 1, below the default min_mass=2.
        let cigar = vec![Cigar::Equal(3), Cigar::Diff(1), Cigar::Equal(3)];
        let clusters = clusters_from_cigar(cigar, 0, edit_mass_settings());
        assert_eq!(clusters.len(), 0);
    }

    #[test]
    fn single_long_deletion_forms_cluster() {
        // edit-mass: 5= 30D 5=, a single 30-base deletion. Under legacy log-complexity +
        // reference-only span this fails the density gate; with linear edit mass (30) over a
        // symmetric span it is a cluster on its own.
        let cigar = vec![Cigar::Equal(5), Cigar::Del(30), Cigar::Equal(5)];
        let clusters = clusters_from_cigar(cigar, 0, edit_mass_settings());
        assert_eq!(clusters.len(), 1);
    }

    #[test]
    fn single_long_insertion_forms_cluster() {
        // edit-mass: 5= 30I 5=, a single 30-base insertion forms a cluster via edit mass alone.
        let cigar = vec![Cigar::Equal(5), Cigar::Ins(30), Cigar::Equal(5)];
        let clusters = clusters_from_cigar(cigar, 0, edit_mass_settings());
        assert_eq!(clusters.len(), 1);
    }

    #[test]
    fn density_filter_rejects_sparse_cluster() {
        // 1X 20= 1X: exactly 20 matches between events — not over max_gap, so gap is NOT broken.
        // But complexity=2.0 over span=22 → density≈0.091, below min_density=0.2.
        let cigar = vec![Cigar::Diff(1), Cigar::Equal(20), Cigar::Diff(1)];
        let clusters = clusters_from_cigar(cigar, 0, ClusteringSettings::default());
        assert_eq!(clusters.len(), 0);
    }
}