twitcher 0.7.0

use anyhow::{Context, bail};
use bstr::ByteSlice;
use iterator::{ClusterOrRecords, find_clusters};
use itertools::Itertools;
use lib_tsalign::a_star_aligner::{
    alignment_geometry::{AlignmentCoordinates, AlignmentRange},
    alignment_result::alignment::Alignment,
    template_switch_distance::AlignmentType,
};
use rust_htslib::bam;
use std::{borrow::Cow, path::PathBuf, sync::Arc};
use tokio::{
    fs::File,
    io::{AsyncWriteExt, BufWriter},
    pin,
    sync::mpsc,
};
use tokio_stream::StreamExt;
use tracing::{error, instrument, trace, warn, warn_span};

use crate::{
    common::{
        ImmutableSequence, SequencePair,
        aligner::{AlignmentOrchestrator, AlignmentQuery, InProgress, cli::CliAlignmentArgs},
        alignment::ForwardAlignment,
        contig::ContigName,
        coords::{GenomePosition, GenomeRegion},
        csv::CSVAuxData,
        list_of_regions::Targets,
        reference::{ReferenceQueryResult, ReferenceReader},
    },
    counter,
    vcf::pipeline::{
        clusterizer::{local_phasing::BamPhaseResolverError, phasing::SplitIntoHaplotypesError},
        reader::VCFReader,
        record::InputRecord,
    },
};

use super::Message;

pub mod cluster;
mod iterator;
pub mod local_phasing;
pub mod phasing;

pub struct Clusterizer {
    input: VCFReader,
    targets: Option<Targets>,
    reference: ReferenceReader,
    output: mpsc::Sender<Message>,
    settings: ClusterizerSettings,
    /// Optional indexed BAM for local read-based phasing of `UnphasedHet` records.
    pub bam: Option<bam::IndexedReader>,
    pub phasing: local_phasing::PhasingSettings,
}

#[derive(clap::Args, Debug, Default)]
pub struct ClusterizerSettings {
    /// Select how clusters are pre-selected. Currently, the default option is the only sensible one, and the other one is for debug purposes.
    #[command(flatten)]
    pub cluster_strategy: cluster::ClusteringSettings,

    #[command(flatten)]
    pub aligner: CliAlignmentArgs,

    /// Write a TSV report of clusters that could not be resolved into haplotypes.
    ///
    /// One line per unresolvable cluster, with three tab-separated columns:
    ///
    /// 1. the cluster region as `contig:start-end` (1-based, inclusive)
    ///
    /// 2. the phasing status of the cluster: `trivial` (no phasing needed),
    ///    `no reads` (no `--phasing-bam` given), `error(...)` (phasing failed,
    ///    the parenthesised text names the cause), or `phased[D,D,...]` with one
    ///    decision per record that phasing attempted — `OK_DIRECT`, `OK_FLIP`,
    ///    `NO_EVIDENCE`, `MIN_READS` or `MIN_CONFIDENCE`
    ///
    /// 3. why the cluster is unresolvable: `MoreThanOneUnphasedHet` or
    ///    `MissingAllele`
    ///
    /// The records of these clusters are still passed through to the output VCF unchanged.
    #[arg(long = "output-unresolvable", value_name = "FILE")]
    pub unresolvable_out: Option<PathBuf>,
}

impl Clusterizer {
    pub(super) const fn new(
        input: VCFReader,
        reference: ReferenceReader,
        targets: Option<Targets>,
        output: mpsc::Sender<Message>,
        settings: ClusterizerSettings,
        bam: Option<bam::IndexedReader>,
        phasing: local_phasing::PhasingSettings,
    ) -> Self {
        Self {
            input,
            targets,
            reference,
            output,
            settings,
            bam,
            phasing,
        }
    }

    pub async fn run(self) -> anyhow::Result<()> {
        let Self {
            input,
            targets,
            reference,
            output,
            settings,
            mut bam,
            phasing: phasing_settings,
        } = self;

        let aligner = AlignmentOrchestrator::try_from(&settings.aligner)
            .map_err(|e| anyhow::anyhow!("Cannot initialize aligners: {e}"))?;

        let aligner_padding = settings.aligner.padding;
        let aligner_range_extension = settings.aligner.range_extension;
        let strategy = settings.cluster_strategy;

        let sample_name = input
            .header()
            .samples()
            .first()
            .and_then(|s| Some(s.to_str().ok()?.to_string()));
        let clusters = find_clusters(strategy.clone(), input, targets);
        let ctx = AlignmentContext {
            reference: &reference,
            aligner: &aligner,
            padding: aligner_padding,
            range_extension: aligner_range_extension,
            sample_name: sample_name.as_deref(),
        };

        pin!(clusters);
        let mut cluster_id = 0usize;
        let mut proximity_cluster_id = 0usize;

        let mut unresolvable_out = if let Some(p) = settings.unresolvable_out {
            Some(BufWriter::new(File::create(p).await?))
        } else {
            None
        };

        while let Some(e) = clusters.next().await {
            match e {
                ClusterOrRecords::Cluster(records) => {
                    proximity_cluster_id += 1;
                    counter!("clusters").inc(1);
                    let proximity_id_str = proximity_cluster_id.to_string();

                    let mut classified: Vec<Arc<InputRecord>> =
                        records.iter().map(|r| Arc::new(r.clone())).collect();

                    // `phasing_ran` records whether local phasing actually ran to completion on
                    // this cluster, so the resolution outcome below can be attributed to it.
                    let (phasing_status, phasing_ran) = run_local_phasing(
                        &mut classified,
                        &records,
                        bam.as_mut(),
                        &phasing_settings,
                    );

                    let sub_clusters =
                        match phasing::split_into_haplotype_clusters(&classified, &strategy) {
                            Ok(sc) => sc,
                            Err(e) => {
                                report_unresolvable(
                                    &e,
                                    &records,
                                    &phasing_status,
                                    phasing_ran,
                                    unresolvable_out.as_mut(),
                                )
                                .await?;
                                output
                                    .send(Message::passthrough(records))
                                    .await
                                    .context("the channel closed unexpectedly")?;
                                continue;
                            }
                        };

                    counter!("clusters.resolved").inc(1);
                    if phasing_ran {
                        counter!("clusters.resolved.after_phasing").inc(1);
                    }

                    dispatch_sub_clusters(
                        &ctx,
                        sub_clusters,
                        &proximity_id_str,
                        &mut cluster_id,
                        &output,
                    )
                    .await?;
                }
                ClusterOrRecords::Records(records) => {
                    output
                        .send(Message::passthrough(records))
                        .await
                        .context("Channel closed !?")?;
                }
                // VCF pipeline does not use the in-memory cache (no enable_cache() call),
                // so SequenceDone is a no-op here.
                ClusterOrRecords::SequenceDone => {} // continue
            }
        }
        if let Some(o) = &mut unresolvable_out {
            o.flush().await?;
        }
        Ok(())
    }
}

/// Align each haplotype sub-cluster of one proximity cluster and forward it to the writer.
/// `cluster_id` is the running sub-cluster counter, advanced once per sub-cluster.
async fn dispatch_sub_clusters(
    ctx: &AlignmentContext<'_>,
    sub_clusters: Vec<phasing::HaplotypeSubCluster>,
    cluster_grp: &str,
    cluster_id: &mut usize,
    output: &mpsc::Sender<Message>,
) -> anyhow::Result<()> {
    for mut sub in sub_clusters {
        *cluster_id += 1;
        sub.records.sort_by_key(|(r, _)| r.pos());
        let recs: Vec<(InputRecord, u32)> = sub
            .records
            .iter()
            .map(|(r, a)| ((**r).clone(), *a))
            .collect();
        let sub_phasing = phasing::OutputPhasing::from_subcluster(sub.haplo, sub.phaseset);
        let pending = start_alignment(
            ctx,
            &recs,
            &cluster_id.to_string(),
            cluster_grp,
            sub_phasing,
        )
        .await;
        output
            .send(Message::cluster(pending, recs, sub_phasing))
            .await?;
    }
    Ok(())
}

/// The loop-invariant inputs needed to start the alignment of one sub-cluster.
struct AlignmentContext<'a> {
    reference: &'a ReferenceReader,
    aligner: &'a AlignmentOrchestrator,
    padding: usize,
    range_extension: usize,
    sample_name: Option<&'a str>,
}

/// Try to resolve the phasing of `classified` from read evidence.
///
/// Returns the phasing status as reported by `--output-unresolvable`, and whether phasing
/// actually ran to completion (so the resolution outcome can be attributed to it).
fn run_local_phasing(
    classified: &mut [Arc<InputRecord>],
    records: &[InputRecord],
    bam: Option<&mut bam::IndexedReader>,
    settings: &local_phasing::PhasingSettings,
) -> (Cow<'static, str>, bool) {
    if phasing::is_trivially_resolvable(classified) {
        counter!("clusters.phasing.not_needed").inc(1);
        return ("trivial".into(), false);
    }
    let Some(bam_reader) = bam else {
        counter!("clusters.phasing.skipped.no_reads").inc(1);
        return ("no reads".into(), false);
    };
    let Ok(cluster_region) = extract_region(records) else {
        counter!("clusters.phasing.failed.no_region").inc(1);
        return ("error(region resolution)".into(), false);
    };

    match tokio::task::block_in_place(|| {
        local_phasing::resolve_phasing(classified, &cluster_region, bam_reader, settings)
    }) {
        Err(BamPhaseResolverError::MultiplePhaseSets(_)) => {
            counter!("clusters.phasing.failed.multiple_phasesets").inc(1);
            ("error(multiple phasesets)".into(), false)
        }
        Err(BamPhaseResolverError::BamReadError(_)) => {
            counter!("clusters.phasing.failed.bam_error").inc(1);
            ("error(bam error)".into(), false)
        }
        Err(BamPhaseResolverError::Other(_)) => {
            counter!("clusters.phasing.failed.other").inc(1);
            ("error(other)".into(), false)
        }
        Ok(decisions) => {
            counter!("clusters.phasing.completed").inc(1);
            let results = decisions
                .into_iter()
                .map(|d| match d {
                    local_phasing::Decision::Resolved(local_phasing::Orientation::Direct) => {
                        "OK_DIRECT"
                    }
                    local_phasing::Decision::Resolved(local_phasing::Orientation::Flipped) => {
                        "OK_FLIP"
                    }
                    local_phasing::Decision::NoEvidence => "NO_EVIDENCE",
                    local_phasing::Decision::InsufficientReads => "MIN_READS",
                    local_phasing::Decision::LowConfidence => "MIN_CONFIDENCE",
                })
                .join(",");
            (format!("phased[{results}]").into(), true)
        }
    }
}

/// Count a cluster that could not be split into haplotypes, and report it to the
/// `--output-unresolvable` file if one is configured.
async fn report_unresolvable(
    error: &SplitIntoHaplotypesError,
    records: &[InputRecord],
    phasing_status: &str,
    phasing_ran: bool,
    unresolvable_out: Option<&mut BufWriter<File>>,
) -> anyhow::Result<()> {
    match error {
        SplitIntoHaplotypesError::MoreThanOneUnphasedHet => {
            counter!("clusters.unresolvable.multiple_unphased_het").inc(1);
        }
        SplitIntoHaplotypesError::MissingAllele => {
            counter!("clusters.unresolvable.missing_allele").inc(1);
        }
        SplitIntoHaplotypesError::Other(_) => {
            counter!("clusters.unresolvable.other").inc(1);
        }
    }
    if phasing_ran {
        counter!("clusters.unresolvable.after_phasing").inc(1);
    }
    if let Some(file) = unresolvable_out {
        file.write_all(
            format!(
                "{}\t{}\t{:?}\n",
                GenomeRegion::try_from(records)?,
                phasing_status,
                error
            )
            .as_bytes(),
        )
        .await?;
    }
    Ok(())
}

/// Prepare one sub-cluster and hand it to the aligner.
///
/// `None` on any preparation / alignment-start failure: the sub-cluster still carries its records
/// and phasing, so the writer reconstructs per-haplotype biallelic records instead of dropping
/// the variants.
async fn start_alignment(
    ctx: &AlignmentContext<'_>,
    recs: &[(InputRecord, u32)],
    cluster_id: &str,
    cluster_grp: &str,
    phasing: phasing::OutputPhasing,
) -> Option<Box<(InProgress, CSVAuxData)>> {
    let prepared =
        match prepare_cluster(ctx.reference, recs, ctx.padding, ctx.range_extension).await {
            Ok(Some(prepared)) => prepared,
            Ok(None) => return None,
            Err((reg, err)) => {
                let _s = reg.map(|r| warn_span!("Failed to prepare", pos = %r).entered());
                counter!("alignments.skipped.overlapping_mutations").inc(1);
                warn!("{err}");
                return None;
            }
        };

    let pending = match ctx.aligner.get_or_compute_alignment(
        ctx.reference.get_name(),
        &prepared.reference_region.clone(),
        prepared.region.clone(),
        prepared.query.clone(),
    ) {
        Ok(pending) => pending,
        Err(e) => {
            error!("Could not get or start the computation of an alignment: {e}");
            return None;
        }
    };

    Some(Box::new((
        pending,
        CSVAuxData {
            cluster_id: cluster_id.to_string(),
            cluster_grp: cluster_grp.to_string(),
            sequences: prepared.query.sequences,
            ref_context_region: prepared.reference_region.clone(),
            // VCF alt sequence is synthetic (reference + applied mutations), so there is no
            // distinct alt context window — the region is intentionally the same.
            alt_context_region: prepared.reference_region,
            region: prepared.region,
            vcf_record_region: Some(prepared.vcf_region.to_string()),
            alt_id: ctx.sample_name.map(ToString::to_string),
            forward_alignment: ForwardAlignment(prepared.fw_alignment),
            cost: ctx.aligner.costs.clone(),
            reference_name: ctx.reference.get_name().to_string(),
            output_phasing: Some(phasing),
        },
    )))
}

#[derive(Clone, Debug)]
struct PreparedCluster {
    query: AlignmentQuery,
    reference_region: GenomeRegion,
    /// The region the cluster's mutations actually edit: the window that is aligned, and the
    /// record's identity. See [`extract_edited_region`].
    region: GenomeRegion,
    /// The POS-based extent of the source records, i.e. `region` plus the anchor bases VCF puts
    /// in front of indels. A label for looking the records up again, nothing else depends on it.
    vcf_region: GenomeRegion,
    fw_alignment: Alignment<AlignmentType>,
}

async fn prepare_cluster(
    reference: &ReferenceReader,
    cluster: &[(InputRecord, u32)],
    padding: usize,
    range_extension: usize,
) -> Result<Option<PreparedCluster>, (Option<GenomeRegion>, anyhow::Error)> {
    let vcf_region = extract_region(cluster.iter().map(|(r, _)| r)).map_err(|e| (None, e))?;
    let region = extract_edited_region(cluster).map_err(|e| (Some(vcf_region.clone()), e))?;
    let Some(ReferenceQueryResult {
        region: actual_region,
        sequence: reference_sequence,
        range_in_sequence,
    }) = reference
        .get_seq(region.clone(), padding, padding)
        .await
        .map_err(|e| (Some(vcf_region.clone()), e))?
    else {
        return Ok(None);
    };

    let _span = warn_span!("prepare_cluster", pos = %region, vcf = %vcf_region).entered();

    // `range_in_sequence` is where the cluster sits in the padded sequence, so the paddings that
    // were actually available (both are clamped at the contig boundaries) follow from it.
    let actual_padding_left = range_in_sequence.start;
    let actual_padding_right = reference_sequence.len() - range_in_sequence.end;

    let (query_sequence, fw_alignment) = apply_mutations(
        &reference_sequence,
        actual_region.start().position_0(),
        cluster.iter().map(|(r, a)| (r, *a)),
    )
    .map_err(|e| (Some(vcf_region.clone()), e))?;

    let ranges = AlignmentRange::new_offset_limit(
        AlignmentCoordinates::new(
            actual_padding_left.saturating_sub(range_extension),
            actual_padding_left.saturating_sub(range_extension),
        ),
        AlignmentCoordinates::new(
            (reference_sequence.len() - actual_padding_right + range_extension)
                .min(reference_sequence.len()),
            (query_sequence.len() - actual_padding_right + range_extension)
                .min(query_sequence.len()),
        ),
    );

    Ok(Some(PreparedCluster {
        query: AlignmentQuery {
            sequences: SequencePair {
                reference: reference_sequence,
                query: query_sequence,
            },
            ranges,
        },
        reference_region: actual_region,
        region,
        vcf_region,
        fw_alignment,
    }))
}

fn extract_region<'a>(
    cluster: impl IntoIterator<Item = &'a InputRecord>,
) -> anyhow::Result<GenomeRegion> {
    let (start, end, contig) = {
        let (mut start, mut end, mut contig) = (i64::MAX, 0, None);
        for r in cluster {
            start = start.min(r.pos());
            end = end.max(r.end());
            if contig.is_none()
                && let Some(present) = r.rid()
            {
                let name = r.header().rid2name(present)?;
                contig = Some(ContigName::new(name));
            }
        }
        (
            usize::try_from(start)?,
            usize::try_from(end)?,
            contig.context("No rid present in records??")?,
        )
    };
    let query_region = GenomeRegion::new_bounded(GenomePosition::new_0(contig, start), end - start);
    Ok(query_region)
}

/// The region a cluster's mutations actually edit, as opposed to the raw VCF extent that
/// [`extract_region`] reports.
///
/// VCF anchors an indel one base left of the edit, so a cluster starting with one covers a
/// reference base it does not touch. The `reads` subcommand derives its cluster regions from the
/// CIGAR, where an insertion sits after the preceding match and covers no reference base at all —
/// the same event therefore came out one base wider in `vcf` than in `reads`. Shifting each record
/// past its shared allele prefix removes the anchor and makes the two agree; a cluster of nothing
/// but an insertion becomes a zero-width region at the point it is inserted into.
///
/// The end needs no such correction: dropping the prefix shortens a record's reference allele by
/// exactly as much as it moves its start right, so `end()` already points past the last edited
/// base.
fn extract_edited_region(cluster: &[(InputRecord, u32)]) -> anyhow::Result<GenomeRegion> {
    let raw = extract_region(cluster.iter().map(|(r, _)| r))?;
    let mut start = usize::MAX;
    for (record, allele) in cluster {
        start = start.min(usize::try_from(record.pos())? + cluster::edit_offset(record, *allele));
    }
    let end = raw.end_excl().context("cluster region is unbounded?")?;
    GenomeRegion::from_incl_excl(
        GenomePosition::new_0(raw.contig().clone(), start),
        Some(end),
    )
}

/// Validate and anchor-trim the alleles of one record, returning the position and alleles to
/// splice into the alt sequence. `None` = the record contributes nothing (already logged).
fn normalise_alleles<'a>(
    mut pos: usize,
    mut ref_allele: &'a [u8],
    mut alt_allele: &'a [u8],
) -> Option<(usize, &'a [u8], &'a [u8])> {
    if alt_allele == b"*" {
        // The allele is missing because an upstream deletion spans this position. The
        // sequence change is fully described by that deletion record, so this one must
        // contribute nothing -- and in particular must not advance `last_end`, nor trip
        // the overlap check of the caller, which it does by construction.
        trace!("Skipping record at {pos}: allele is deleted by an upstream deletion");
        return None;
    }

    if let Some(allele) = [ref_allele, alt_allele].into_iter().find(|a| {
        !a.iter()
            .all(|b| matches!(b.to_ascii_uppercase(), b'A' | b'C' | b'G' | b'T' | b'N'))
    }) {
        // Symbolic alleles (`<DEL>`, `<NON_REF>`, ...), breakends and missing alleles
        // carry no literal sequence, so splicing them in would corrupt the alt sequence.
        warn!(
            "Skipping record at {pos}: allele `{}` is not a DNA sequence",
            String::from_utf8_lossy(allele)
        );
        return None;
    }

    // Strip the bases both alleles share at the front: they are not edited, they only shift the
    // record to the right. For a plain VCF indel (`A -> AC`, `AC -> A`) that is the single anchor
    // base, but a repeat-expansion record (`(CAG)x84 -> (CAG)x100`) shares far more, and the
    // shared part must be removed here as well: [`cluster::edit_offset`] removes all of it when it
    // places the cluster, and the two must agree or the mutation lands outside the aligned window.
    let shared_prefix = ref_allele
        .iter()
        .zip(alt_allele)
        .take_while(|(r, a)| r.eq_ignore_ascii_case(a))
        .count();
    ref_allele = ref_allele.get(shared_prefix..).unwrap_or_default();
    alt_allele = alt_allele.get(shared_prefix..).unwrap_or_default();
    pos += shared_prefix;

    if ref_allele.is_empty() && alt_allele.is_empty() {
        warn!("Empty record; Skipping.");
        return None;
    }

    Some((pos, ref_allele, alt_allele))
}

/// Append the alignment of one anchor-trimmed mutation to `cigar`.
fn push_mutation_cigar(
    cigar: &mut Alignment<AlignmentType>,
    ref_allele: &[u8],
    alt_allele: &[u8],
) -> anyhow::Result<()> {
    match (ref_allele.len(), alt_allele.len()) {
        (0, 0) => {
            bail!("Empty record; should have been caught earlier!");
        }
        // TODO all of the following cases can actually be seen as only the n,m case but for clarity we leave these here.
        (1, 1) => {
            // SNV
            cigar.push(AlignmentType::PrimarySubstitution);
        }
        (0, m) => {
            // insertion
            cigar.push_n(m, AlignmentType::PrimaryInsertion);
        }
        (n, 0) => {
            // deletion
            cigar.push_n(n, AlignmentType::PrimaryDeletion);
        }
        (n, m) => {
            // TODO perhaps run some simple local aligner here? For now, we're gonna put match/mismatch and ins/del
            let match_len = n.min(m);
            for (r, a) in ref_allele.iter().zip(alt_allele.iter()).take(match_len) {
                if r.eq_ignore_ascii_case(a) {
                    cigar.push(AlignmentType::PrimaryMatch);
                } else {
                    cigar.push(AlignmentType::PrimarySubstitution);
                }
            }

            let extra = n.max(m) - match_len;
            if n > m {
                cigar.push_n(extra, AlignmentType::PrimaryDeletion);
            } else {
                cigar.push_n(extra, AlignmentType::PrimaryInsertion);
            }
        }
    }
    Ok(())
}

/// Build the alt sequence.
#[instrument(name = "build_query_sequence", skip_all)]
fn apply_mutations<'a>(
    reference: &[u8],
    reference_start: usize,
    mutations: impl Iterator<Item = (&'a InputRecord, u32)>,
) -> anyhow::Result<(ImmutableSequence, Alignment<AlignmentType>)> {
    let mut alt_sequence = Vec::new();
    let mut cigar = Alignment::new();
    let mut last_end = reference_start;
    let reference_end = reference_start + reference.len();

    let pos2off = |pos: usize| pos - reference_start;

    for (m, allele_idx) in mutations {
        let pos = usize::try_from(m.pos())?;
        let end = usize::try_from(m.end())?;
        let alleles = m.alleles();
        let Some((ref_allele, alt_alleles)) = alleles.split_first() else {
            warn!("Record at {pos} has no alt allele");
            continue;
        };
        if allele_idx == 0 {
            warn!(
                "Record at {pos} has invalid allele index {allele_idx} (only {} alleles); skipping",
                alleles.len()
            );
            continue;
        }
        let alt_allele = alt_alleles
            .get(allele_idx as usize - 1)
            .context("alt allele not found")?;
        let Some((pos, ref_allele, alt_allele)) = normalise_alleles(pos, ref_allele, alt_allele)
        else {
            continue;
        };

        if pos < last_end {
            bail!(
                "Overlapping mutation: last mutation ended at {last_end}, \
                 but this one starts at {pos}"
            );
        }

        if end > reference_end {
            warn!(
                "Skipping out-of-bounds mutation at {pos}: \
                 record end {end} exceeds reference end {reference_end}"
            );
            continue;
        }

        if pos > end {
            warn!("Skip negative-sized mutation: the position is {pos} but the end is {end}",);
            continue;
        }

        if pos - last_end > 0 {
            alt_sequence.extend_from_slice(
                reference
                    .get(pos2off(last_end)..pos2off(pos))
                    .context("invalid reference offsets")?,
            );
            cigar.push_n(pos - last_end, AlignmentType::PrimaryMatch);
        }

        // The reference is uppercased on read, so alleles are normalised to match.
        alt_sequence.extend(alt_allele.iter().map(u8::to_ascii_uppercase)); // TODO should we take anything else than the first alt allele?

        push_mutation_cigar(&mut cigar, ref_allele, alt_allele)?;

        last_end = end;
    }

    alt_sequence.extend_from_slice(
        reference
            .get(pos2off(last_end)..)
            .context("invalid ref offsets")?,
    );
    cigar.push_n(
        reference.len() - pos2off(last_end),
        AlignmentType::PrimaryMatch,
    );

    Ok((alt_sequence.into(), cigar))
}

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

    use bstr::ByteSlice as _;
    use rust_htslib::faidx;
    use std::path::Path;

    use super::{extract_edited_region, extract_region, normalise_alleles, prepare_cluster};
    use crate::{
        common::reference::{CliReferenceArg, ReferenceReader},
        vcf::pipeline::{clusterizer::cluster::edit_offset, record::InputRecord},
    };

    fn make_writer() -> Writer {
        let mut header = Header::new();
        header.push_record(b"##contig=<ID=chr1,length=100000000>");
        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)
    }

    /// The reported bug from <https://version.helsinki.fi/kraujasp/twitcher/-/work_items/80>
    /// An insertion followed by an SNV. The insertion's anchor base is not
    /// edited, so the region must start at the same location as the SNV
    #[test]
    fn leading_insertion_does_not_widen_the_edited_region() {
        let w = make_writer();
        let cluster = [
            (make_record(&w, 1_070_825, &[b"C", b"CGG"]), 1),
            (make_record(&w, 1_070_826, &[b"C", b"G"]), 1),
        ];

        let raw = extract_region(cluster.iter().map(|(r, _)| r)).unwrap();
        assert_eq!(raw.to_string(), "chr1:1070826-1070827");

        let edited = extract_edited_region(&cluster).unwrap();
        assert_eq!(edited.start().position_0(), 1_070_826);
        assert_eq!(edited.len(), Some(1));
        assert_eq!(edited.to_string(), "chr1:1070827-1070827");
    }

    /// A leading deletion carries the same anchor base, and must be corrected the same way.
    #[test]
    fn leading_deletion_does_not_widen_the_edited_region() {
        let w = make_writer();
        let cluster = [
            (make_record(&w, 100, &[b"AT", b"A"]), 1),
            (make_record(&w, 105, &[b"C", b"G"]), 1),
        ];

        // The deleted base is at 101, so that is where the `reads` subcommand puts its `D(1)`.
        let edited = extract_edited_region(&cluster).unwrap();
        assert_eq!(edited.start().position_0(), 101);
        assert_eq!(edited.end_excl().unwrap().position_0(), 106);
    }

    /// A cluster of nothing but an insertion covers no reference base at all, exactly as the
    /// CIGAR-derived cluster of an `I` operation does.
    #[test]
    fn lone_insertion_becomes_a_zero_width_region() {
        let w = make_writer();
        let cluster = [(make_record(&w, 100, &[b"A", b"ATTT"]), 1)];

        let edited = extract_edited_region(&cluster).unwrap();
        assert_eq!(edited.start().position_0(), 101);
        assert_eq!(edited.len(), Some(0));
    }

    /// Without an anchor base to strip there is nothing to correct: substitutions keep the raw
    /// VCF extent.
    #[test]
    fn substitutions_keep_the_raw_extent() {
        let w = make_writer();
        let cluster = [
            (make_record(&w, 100, &[b"A", b"T"]), 1),
            (make_record(&w, 104, &[b"GT", b"CA"]), 1),
        ];

        let raw = extract_region(cluster.iter().map(|(r, _)| r)).unwrap();
        assert_eq!(extract_edited_region(&cluster).unwrap(), raw);
    }

    /// The correction uses the allele the cluster was built with, not the first alt.
    #[test]
    fn edited_region_follows_the_selected_allele() {
        let w = make_writer();
        let insertion = [(make_record(&w, 100, &[b"C", b"CGG", b"G"]), 1)];
        let substitution = [(make_record(&w, 100, &[b"C", b"CGG", b"G"]), 2)];

        assert_eq!(
            extract_edited_region(&insertion)
                .unwrap()
                .start()
                .position_0(),
            101
        );
        assert_eq!(
            extract_edited_region(&substitution)
                .unwrap()
                .start()
                .position_0(),
            100
        );
    }

    /// A record whose edit starts after a later record's does not drag the region's start right.
    #[test]
    fn the_edited_start_is_the_minimum_over_all_records() {
        let w = make_writer();
        let cluster = [
            (make_record(&w, 100, &[b"ATTTT", b"A"]), 1),
            (make_record(&w, 102, &[b"T", b"G"]), 1),
        ];

        let edited = extract_edited_region(&cluster).unwrap();
        assert_eq!(edited.start().position_0(), 101);
        assert_eq!(edited.end_excl().unwrap().position_0(), 105);
    }

    /// The position [`normalise_alleles`] splices at must be the one
    /// [`super::cluster::edit_offset`] reports, or the mutation lands outside the window that
    /// [`prepare_cluster`] builds around the edited region.
    #[test]
    fn normalise_alleles_agrees_with_edit_offset() {
        let w = make_writer();
        for alleles in [
            [b"A".as_slice(), b"AT".as_slice()],
            [b"AT", b"A"],
            [b"A", b"T"],
            [b"AT", b"GC"],
            [b"AT", b"AGC"],
            [b"CAGCAGCAG", b"CAGCAGCAGCAG"],
            [b"AAT", b"AAG"],
        ] {
            let record = make_record(&w, 100, &alleles);
            let (pos, ..) = normalise_alleles(100, alleles[0], alleles[1]).unwrap();
            assert_eq!(
                pos - 100,
                edit_offset(&record, 1),
                "alleles {}/{}",
                alleles[0].as_bstr(),
                alleles[1].as_bstr()
            );
        }
    }

    /// A repeat expansion shares far more than the single anchor base.
    #[test]
    fn normalise_alleles_strips_the_entire_shared_prefix() {
        let (pos, reference, alt) = normalise_alleles(100, b"CAGCAGCAG", b"CAGCAGCAGCAG").unwrap();
        assert_eq!(pos, 109);
        assert!(reference.is_empty());
        assert_eq!(alt, b"CAG");
    }

    /// Identical alleles edit nothing and must not be spliced.
    #[test]
    fn normalise_alleles_rejects_an_unedited_record() {
        assert!(normalise_alleles(100, b"ACGT", b"ACGT").is_none());
    }

    /// A 400 base `ACGT` repeat, plus its index, opened as a reference.
    fn acgt_reference(dir: &Path) -> ReferenceReader {
        let path = dir.join("ref.fa");
        let sequence: Vec<u8> = b"ACGT".repeat(100);
        let mut fasta = b">chr1\n".to_vec();
        for line in sequence.chunks(80) {
            fasta.extend_from_slice(line);
            fasta.push(b'\n');
        }
        std::fs::write(&path, fasta).unwrap();
        faidx::build(&path).unwrap();
        ReferenceReader::try_from(&CliReferenceArg::from(path.to_str().unwrap())).unwrap()
    }

    /// The aligned window covers the edited region, not the POS-based extent of the records.
    ///
    /// The deletion below removes the `C` at 101, which is exactly the base the `reads`
    /// subcommand's `1D` covers, so both subcommands align the same window.
    #[tokio::test(flavor = "multi_thread")]
    async fn the_aligned_window_starts_at_the_edited_region() {
        let dir = tempfile::tempdir().unwrap();
        let reference = acgt_reference(dir.path());
        let w = make_writer();
        let cluster = [
            (make_record(&w, 100, &[b"AC", b"A"]), 1),
            (make_record(&w, 105, &[b"C", b"G"]), 1),
        ];

        let prepared = prepare_cluster(&reference, &cluster, 10, 0)
            .await
            .unwrap()
            .unwrap();

        assert_eq!(prepared.region.to_string(), "chr1:102-106");
        assert_eq!(prepared.vcf_region.to_string(), "chr1:101-106");
        // The window is the edited region plus ten bases on either side.
        assert_eq!(prepared.reference_region.to_string(), "chr1:92-116");
        assert_eq!(prepared.query.sequences.reference.len(), 25);
        // One base is deleted, so the alt sequence is one shorter.
        assert_eq!(prepared.query.sequences.query.len(), 24);

        let ranges = &prepared.query.ranges;
        assert_eq!(ranges.reference_offset(), 10);
        assert_eq!(ranges.reference_limit(), 15);
        assert_eq!(ranges.query_offset(), 10);
        assert_eq!(ranges.query_limit(), 14);
    }

    /// A lone insertion edits no reference base at all. The `reads` subcommand builds the same
    /// zero-width region from an `I` operation, so the window must be zero-width here too.
    #[tokio::test(flavor = "multi_thread")]
    async fn a_lone_insertion_yields_a_zero_width_window() {
        let dir = tempfile::tempdir().unwrap();
        let reference = acgt_reference(dir.path());
        let w = make_writer();
        let cluster = [(make_record(&w, 100, &[b"A", b"ATTT"]), 1)];

        let prepared = prepare_cluster(&reference, &cluster, 10, 0)
            .await
            .unwrap()
            .unwrap();

        assert_eq!(prepared.region.len(), Some(0));
        // A zero-length interval, and the only region string samtools cannot be handed back.
        assert_eq!(prepared.region.to_string(), "chr1:102-101");
        assert_eq!(prepared.reference_region.to_string(), "chr1:92-111");
        assert_eq!(prepared.query.sequences.reference.len(), 20);
        assert_eq!(prepared.query.sequences.query.len(), 23);

        let ranges = &prepared.query.ranges;
        assert_eq!(ranges.reference_offset(), 10);
        assert_eq!(ranges.reference_limit(), 10);
        assert_eq!(ranges.query_offset(), 10);
        assert_eq!(ranges.query_limit(), 13);
    }

    /// A zero-width window with no padding has no sequence to align. That must fail this one
    /// cluster with an error, not abort the process inside htslib.
    #[tokio::test(flavor = "multi_thread")]
    async fn a_zero_width_window_without_padding_fails_the_cluster() {
        let dir = tempfile::tempdir().unwrap();
        let reference = acgt_reference(dir.path());
        let w = make_writer();
        let cluster = [(make_record(&w, 100, &[b"A", b"ATTT"]), 1)];

        assert!(prepare_cluster(&reference, &cluster, 0, 0).await.is_err());
    }
}