twitcher 0.6.9

use std::borrow::Borrow;

use crate::vcf::pipeline::record::InputRecord;

pub use crate::common::cluster_settings::ClusteringSettings;

impl ClusteringSettings {
    pub fn belongs<R: Borrow<InputRecord>>(
        &self,
        last: Option<R>,
        candidate: R,
    ) -> anyhow::Result<bool> {
        let Some(last) = last else {
            return Ok(true);
        };
        Ok(candidate.borrow().pos().saturating_sub(last.borrow().end())
            <= self.max_gap.try_into()?)
    }

    pub fn is_cluster<R, F: Fn(&R) -> &InputRecord, G: Fn(&R) -> Option<u32>>(
        &self,
        cluster: &[R],
        get_record: F,
        get_allele: G,
    ) -> bool {
        let Some(first) = cluster.first() else {
            return false;
        };
        let in_first = get_record(first);
        // Reference extent of the cluster, in bases covered. `end()` is the exclusive end of the
        // last record's reference allele and needs no correction: stripping a record's shared
        // prefix moves its start right and shortens its reference allele by the same amount. The
        // first record's start, on the other hand, shifts right by that prefix — the shared bases
        // are not edited and must not count towards the span.
        let first_prefix = self
            .edited_alleles(in_first, get_allele(first))
            .map_or(0, |e| i64::try_from(e.prefix).unwrap_or(0));
        let Some(max_end) = cluster.iter().map(|r| get_record(r).end()).max() else {
            return false;
        };
        #[expect(clippy::expect_used)]
        let l_ref = usize::try_from((in_first.pos() + first_prefix).abs_diff(max_end))
            .expect("coordinates are larger than usize::MAX !?");
        // Accumulate cluster mass, the net length change (Σ alt_len − ref_len) so we can derive
        // the query extent, and the number of events. Records without a usable allele pair are
        // skipped entirely: they contribute no mass, so they must not count towards the minimum
        // number of events either.
        let (mass, net, count) = cluster
            .iter()
            .map(|r| (get_record(r), get_allele(r)))
            .filter_map(|(r, a)| self.edited_alleles(r, a))
            .fold((0.0, 0isize, 0), |(mass, net, count), e| {
                (
                    mass + self.event_mass(e.ref_len, e.alt_len),
                    net.strict_add_unsigned(e.alt_len)
                        .strict_sub_unsigned(e.ref_len),
                    count + 1,
                )
            });
        let span = self.span(l_ref, l_ref.checked_add_signed(net).unwrap_or(0));
        self.is_valid_cluster(mass, span, count)
    }

    /// The edited part of the reference allele and of the relevant alt allele for a record. When no
    /// specific allele is requested, the alt maximizing [`ClusteringSettings::event_mass`] is chosen
    /// so mass and span use the same allele. Returns `None` for records without at least two
    /// alleles, and for alt alleles that edit nothing.
    fn edited_alleles(&self, record: &InputRecord, allele: Option<u32>) -> Option<EditedAlleles> {
        let alleles = record.alleles();
        if alleles.len() < 2 {
            return None;
        }
        let (reference, alts) = alleles.split_first()?;
        if let Some(a) = allele {
            EditedAlleles::new(reference, alleles.get(a as usize)?)
        } else {
            alts.iter()
                .filter_map(|alt| EditedAlleles::new(reference, alt))
                .max_by(|x, y| {
                    self.event_mass(x.ref_len, x.alt_len)
                        .total_cmp(&self.event_mass(y.ref_len, y.alt_len))
                })
        }
    }
}

/// How far right of `POS` the edit of `record` against the given `allele` starts.
pub fn edit_offset(record: &InputRecord, allele: u32) -> usize {
    let alleles = record.alleles();
    let (Some(reference), Some(alt)) = (alleles.first(), alleles.get(allele as usize)) else {
        return 0;
    };
    EditedAlleles::new(reference, alt).map_or(0, |e| e.prefix)
}

/// An allele pair reduced to the bases it actually edits: the leading bases both alleles share are
/// stripped off, since they are not edited and only shift the record to the right.
///
/// This needs no case distinction between variant classes — an MNV shares no prefix, so trimming is
/// a no-op for it. An `A -> AT` insertion becomes `(0, 1)`, one edited base, while the equally
/// sized `A -> CT` becomes `(1, 2)`, two edited bases.
#[derive(Debug, Clone, Copy)]
struct EditedAlleles {
    ref_len: usize,
    alt_len: usize,
    /// Bases shared by both alleles at the front, i.e. how far right of `POS` the edit starts.
    prefix: usize,
}

impl EditedAlleles {
    /// Returns `None` when the two alleles are identical, i.e. the record edits nothing.
    fn new(reference: &[u8], alt: &[u8]) -> Option<Self> {
        let prefix = reference
            .iter()
            .zip(alt)
            .take_while(|(r, a)| r == a)
            .count();
        let (ref_len, alt_len) = (reference.len() - prefix, alt.len() - prefix);
        (ref_len > 0 || alt_len > 0).then_some(Self {
            ref_len,
            alt_len,
            prefix,
        })
    }
}

#[cfg(test)]
mod tests {
    use rust_htslib::bcf::{self, Header, Writer, record::GenotypeAllele};

    use crate::common::cluster_settings::{ClusterStrategy, ClusteringSettings};
    use crate::vcf::pipeline::record::InputRecord;

    use super::edit_offset;

    fn make_writer() -> Writer {
        let mut header = Header::new();
        header.push_record(b"##contig=<ID=chr1,length=1000000>");
        header.push_record(b"##FORMAT=<ID=GT,Number=1,Type=String,Description=\"Genotype\">");
        header.push_sample(b"S1");
        Writer::from_path("/dev/null", &header, true, bcf::Format::Vcf).unwrap()
    }

    /// A minimal heterozygous record at the 0-based `pos` with the given alleles.
    fn make_record(writer: &Writer, pos: i64, alleles: &[&[u8]]) -> InputRecord {
        let mut rec = writer.empty_record();
        let rid = rec.header().name2rid(b"chr1").unwrap();
        rec.set_rid(Some(rid));
        rec.set_pos(pos);
        rec.set_alleles(alleles).unwrap();
        rec.push_genotypes(&[GenotypeAllele::Unphased(0), GenotypeAllele::Unphased(1)])
            .unwrap();
        InputRecord::new(rec)
    }

    fn is_cluster(records: &[InputRecord], settings: &ClusteringSettings) -> bool {
        settings.is_cluster(records, |r| r, |_| None)
    }

    /// The span is the reference extent, not one base more: two SNVs 10 bases apart sit exactly at
    /// the default density threshold (mass 2 / span 10 = 0.2).
    #[test]
    fn legacy_span_is_the_reference_extent() {
        let w = make_writer();
        let settings = ClusteringSettings {
            min_density: 0.2,
            ..Default::default()
        };

        let exact = [
            make_record(&w, 100, &[b"A", b"T"]),
            make_record(&w, 109, &[b"A", b"T"]),
        ];
        assert!(is_cluster(&exact, &settings));

        // One base further apart: mass 2 / span 11 = 0.18, below the threshold.
        let too_sparse = [
            make_record(&w, 100, &[b"A", b"T"]),
            make_record(&w, 110, &[b"A", b"T"]),
        ];
        assert!(!is_cluster(&too_sparse, &settings));
    }

    /// A lone insertion covers no reference base but 20 query bases, so only the symmetric span of
    /// the `edit-mass` strategy makes it dense enough — and only that strategy admits a
    /// single-record cluster at all.
    #[test]
    fn edit_mass_span_follows_the_query_extent() {
        let w = make_writer();
        // 20-base insertion in VCF anchor form: ref "A", alt "A" + 20 bases.
        let alt = vec![b'A'; 21];
        let insertion = [make_record(&w, 100, &[b"A", &alt])];

        assert!(is_cluster(
            &insertion,
            &ClusteringSettings {
                strategy: ClusterStrategy::EditMass,
                ..Default::default()
            }
        ));
        // Legacy needs two records, so the same insertion is not a cluster there.
        assert!(!is_cluster(&insertion, &ClusteringSettings::default()));
    }

    /// An insertion's shared prefix is not edited: `A -> AT` is a one-base edit, while the equally
    /// sized `A -> CT` shares nothing and edits two bases. Only the latter clears the default
    /// `--cluster-min-mass` of 2 on its own.
    #[test]
    fn shared_prefix_is_not_counted_as_edit_mass() {
        let w = make_writer();
        let settings = ClusteringSettings {
            strategy: ClusterStrategy::EditMass,
            ..Default::default()
        };

        let insertion = [make_record(&w, 100, &[b"A", b"AT"])];
        assert!(!is_cluster(&insertion, &settings));

        let delins = [make_record(&w, 100, &[b"A", b"CT"])];
        assert!(is_cluster(&delins, &settings));
    }

    /// The shared prefix also shifts a record's position: the span of the cluster starts where the
    /// first record's edit does, not at its `POS`.
    #[test]
    fn shared_prefix_shifts_the_span_to_the_right() {
        let w = make_writer();
        let settings = ClusteringSettings {
            min_density: 0.2,
            ..Default::default()
        };

        // The insertion edits at 101, the SNV's end is 111 → span 10 for a mass of 2 → exactly at
        // the threshold. Measuring from POS 100 would give span 11 and reject the cluster.
        let records = [
            make_record(&w, 100, &[b"A", b"AT"]),
            make_record(&w, 110, &[b"A", b"T"]),
        ];
        assert!(is_cluster(&records, &settings));
    }

    /// A record whose alt equals its reference edits nothing, so it neither adds mass nor counts as
    /// an event — and above all must not panic.
    #[test]
    fn alt_equal_to_reference_is_not_an_event() {
        let w = make_writer();
        let records = [
            make_record(&w, 100, &[b"A", b"T"]),
            make_record(&w, 101, &[b"A", b"A"]),
        ];
        // One real event only: below the legacy minimum of two.
        assert!(!is_cluster(&records, &ClusteringSettings::default()));
    }

    /// A record without an alt allele contributes no mass, so it must not count towards
    /// `--cluster-min-records` either.
    #[test]
    fn record_without_alt_allele_does_not_count_as_an_event() {
        let w = make_writer();
        let mut ref_only = w.empty_record();
        let rid = ref_only.header().name2rid(b"chr1").unwrap();
        ref_only.set_rid(Some(rid));
        ref_only.set_pos(101);
        ref_only.set_alleles(&[b"A"]).unwrap();

        let records = [
            make_record(&w, 100, &[b"A", b"T"]),
            InputRecord::new(ref_only),
        ];
        // One real event only: below the legacy minimum of two.
        assert!(!is_cluster(&records, &ClusteringSettings::default()));
    }

    /// VCF anchors an indel one base left of the bases it edits, so the shared allele prefix is
    /// exactly how far right of `POS` the edit really starts.
    #[test]
    fn edit_offset_skips_the_vcf_anchor_base() {
        let w = make_writer();

        // Insertion and deletion both carry one anchor base.
        assert_eq!(edit_offset(&make_record(&w, 100, &[b"A", b"AT"]), 1), 1);
        assert_eq!(edit_offset(&make_record(&w, 100, &[b"AT", b"A"]), 1), 1);
        // An SNV and an MNV share nothing, so they edit right at `POS`.
        assert_eq!(edit_offset(&make_record(&w, 100, &[b"A", b"T"]), 1), 0);
        assert_eq!(edit_offset(&make_record(&w, 100, &[b"AT", b"GC"]), 1), 0);
        // A replacement that happens to share a prefix is shifted just the same.
        assert_eq!(edit_offset(&make_record(&w, 100, &[b"AT", b"AGC"]), 1), 1);
    }

    /// The offset follows the allele that was actually selected, not the first alt.
    #[test]
    fn edit_offset_follows_the_selected_allele() {
        let w = make_writer();
        let record = make_record(&w, 100, &[b"C", b"CGG", b"G"]);

        assert_eq!(edit_offset(&record, 1), 1);
        assert_eq!(edit_offset(&record, 2), 0);
        // The reference allele edits nothing, and neither does an allele that does not exist.
        assert_eq!(edit_offset(&record, 0), 0);
        assert_eq!(edit_offset(&record, 7), 0);
    }
}