twitcher 0.7.0

use std::{fmt::Debug, hash::Hash, io::Write, sync::Arc};

use anyhow::Context as _;
use bstr::ByteSlice;
use compact_genome::{
    implementation::{alphabets::dna_alphabet_or_n::DnaAlphabetOrN, vec_sequence::VectorGenome},
    interface::sequence::{GenomeSequence, OwnedGenomeSequence},
};
use csv::Writer;
use generic_a_star::cost::AStarCost;
use lib_tsalign::{
    a_star_aligner::{
        alignment_result::alignment::Alignment, template_switch_distance::AlignmentType,
    },
    config::TemplateSwitchConfig,
    costs::U64Cost,
};
use serde::{Deserialize, Serialize};
use tracing::{instrument, trace};
use xxhash_rust::xxh3::Xxh3Builder;

use crate::{
    common::{
        ImmutableSequence, SequencePair,
        aligner::result::{TSData, TwitcherAlignmentWithStatistics},
        alignment::{ForwardAlignment, consumed_query},
        coords::GenomeRegion,
    },
    vcf::pipeline::clusterizer::phasing::OutputPhasing,
};

#[derive(Serialize, Deserialize)]
pub struct CSVRecord {
    pub id: String,
    pub cluster_id: String, // Some sort of running id (maybe even unique across files?), to make it easy to connect vcf and bam output.
    pub cluster_grp: String, // Proximity cluster group: shared across all haplotype sub-clusters from the same original proximity cluster.
    pub ref_ctx_region: String, // from aux data: entire region
    pub alt_ctx_region: String, // from aux data: entire region
    pub cluster_region: String, // from aux data: the region the cluster's mutations edit, i.e. the "focus" region, a subregion of the context region
    pub ref_cluster_offset: usize,
    // pub ref_cluster_limit: usize,
    pub alt_cluster_offset: usize,
    // pub alt_cluster_limit: usize,
    pub ts_1_4_region: String, // TODO add 1-4 region for the template switches, might be useful for dedup
    pub read_id: Option<String>, // null if vcf
    pub allele_count: Option<usize>, // null if reads
    pub fw_cigar: String,      // for cluster region
    pub fw_cigar_ctx: String,  // for context region
    pub fw_mi_ctx: String,     // for context region
    pub fw_cost: u64,          // for cluster region
    pub fw_cost_ctx: u64,      // for context region
    pub ts_cigar: String,      // for cluster region
    pub ts_cigar_ctx: String,  // for context region
    pub ts_cost: u64,          // for cluster region
    pub ts_cost_ctx: u64,      // for context region
    pub ts_num: usize,
    pub ts_1_2: String,
    pub ts_2_3: String,
    pub ts_2_3_min: String,
    pub ts_2_3_max: String,
    pub ts_1_4: String,
    pub ts_start_left_shift: String,
    pub ts_start_right_shift: String,
    pub ts_end_left_shift: String,
    pub ts_end_right_shift: String,
    pub ts_inner_alignment_cigar: String,
    /// `vcf` only: the POS-based extent of the source records, for looking them up again with
    /// `bcftools -r`. This is `cluster_region` plus the anchor bases VCF puts in front of indels;
    /// nothing but the lookup should use it.
    pub vcf_record_region: Option<String>,
}

static SEP: &str = "|";

pub struct TwitcherCSVWriter(Writer<Box<dyn Write + Send>>);

pub struct CSVAuxData {
    pub cluster_id: String,
    pub cluster_grp: String,
    pub sequences: SequencePair,
    pub ref_context_region: GenomeRegion,
    pub alt_context_region: GenomeRegion,
    /// The region the cluster's mutations actually edit: the window that is aligned, the record's
    /// identity through [`compute_record_hash`], and the position every log and report uses.
    ///
    /// In `reads` this is the CIGAR-derived cluster region; in `vcf` it is the extent of the
    /// records without the anchor bases VCF puts in front of indels, so both subcommands describe
    /// the same event with the same coordinates.
    pub region: GenomeRegion,
    /// `vcf` only: the POS-based extent of the source records, a label for finding them again.
    pub vcf_record_region: Option<String>,
    // pub ranges: AlignmentRange,
    pub alt_id: Option<String>,
    pub forward_alignment: ForwardAlignment,
    pub cost: Arc<TemplateSwitchConfig<DnaAlphabetOrN, U64Cost>>,
    pub reference_name: String,
    pub output_phasing: Option<OutputPhasing>,
}

impl Debug for CSVAuxData {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CSVAuxData")
            .field("cluster_id", &self.cluster_id)
            .field("cluster_grp", &self.cluster_grp)
            .field("sequences", &self.sequences)
            .field("ref_context_region", &self.ref_context_region)
            .field("alt_context_region", &self.alt_context_region)
            .field("region", &self.region)
            .field("vcf_record_region", &self.vcf_record_region)
            // .field("ref range", &self.re)
            // .field("alt range", &self.ranges.query_range())
            .field("alt_id", &self.alt_id)
            .field("forward_alignment", &self.forward_alignment.cigar())
            .field("cost", &"<cost def>")
            .field("reference_name", &self.reference_name)
            .field("output_phasing", &self.reference_name)
            .finish()
    }
}

impl TwitcherCSVWriter {
    pub fn new(write: Box<dyn Write + Send>) -> Self {
        Self(Writer::from_writer(write))
    }

    #[instrument(name = "write_csv_record", skip_all, fields(pos = %aux.region, src = aux.alt_id))]
    pub fn write(
        &mut self,
        result: &TwitcherAlignmentWithStatistics,
        aux: CSVAuxData,
    ) -> anyhow::Result<()> {
        trace!("Writing CSV Record");

        let hash = compute_record_hash(result, &aux);
        let hash = format!("{hash:032x}");

        trace!(
            ref_len = aux.sequences.reference.len(),
            query_len = aux.sequences.query.len(),
            ref_offset = result.stats.reference_offset(),
            query_offset = result.stats.query_offset(),
            fw_cigar = aux.forward_alignment.cigar(),
            "Input sequences and forward alignment"
        );

        trace!(
            ref = %aux.sequences.reference.as_bstr(),
            alt = %aux.sequences.query.as_bstr(),
        );

        let alignments = RecordAlignments::compute(result, &aux)?;
        let record = build_record(hash, result, aux, &alignments)?;

        // Write to csv
        tokio::task::block_in_place(|| self.0.serialize(record))?;

        Ok(())
    }

    pub fn flush(&mut self) -> std::io::Result<()> {
        self.0.flush()
    }
}

/// The four alignments (forward and TS, each for the cluster and the context region) and their
/// costs, as reported in one CSV record.
struct RecordAlignments {
    entire_forward_cost: u64,
    entire_ts: Alignment<AlignmentType>,
    entire_ts_cost: u64,
    cluster_forward: Alignment<AlignmentType>,
    cluster_forward_cost: u64,
}

impl RecordAlignments {
    fn compute(result: &TwitcherAlignmentWithStatistics, aux: &CSVAuxData) -> anyhow::Result<Self> {
        let entire_forward_cost = compute_cost(
            &aux.sequences,
            None,
            aux.forward_alignment.0.clone(),
            &aux.cost,
        )?;
        trace!(
            cost = entire_forward_cost,
            "Computed entire forward alignment cost"
        );

        let entire_ts = aux.forward_alignment.insert_ts_alignment(result)?;
        trace!(
            cigar = entire_ts.cigar(),
            "Entire TS alignment after insert"
        );
        let entire_ts_cost = compute_cost(&aux.sequences, None, entire_ts.clone(), &aux.cost)?;
        trace!(cost = entire_ts_cost, "Computed entire TS alignment cost");

        let cluster_forward = aux.forward_alignment.crop_to_ts_region(result)?;
        trace!(
            cigar = cluster_forward.cigar(),
            ref_offset = result.stats.reference_offset(),
            query_offset = result.stats.query_offset(),
            ref_remaining = aux
                .sequences
                .reference
                .len()
                .saturating_sub(result.stats.reference_offset()),
            query_remaining = aux
                .sequences
                .query
                .len()
                .saturating_sub(result.stats.query_offset()),
            "Cluster forward alignment after crop"
        );
        let cluster_forward_cost = compute_cost(
            &aux.sequences,
            Some((result.stats.reference_offset(), result.stats.query_offset())),
            cluster_forward.clone(),
            &aux.cost,
        )?;
        trace!(
            cost = cluster_forward_cost,
            "Computed cluster forward alignment cost"
        );

        Ok(Self {
            entire_forward_cost,
            entire_ts,
            entire_ts_cost,
            cluster_forward,
            cluster_forward_cost,
        })
    }
}

fn build_record(
    hash: String,
    result: &TwitcherAlignmentWithStatistics,
    aux: CSVAuxData,
    alignments: &RecordAlignments,
) -> anyhow::Result<CSVRecord> {
    let cluster_ts_alignment = &result.alignment.alignment;
    let cluster_ts_alignment_cost = result.alignment.cost.as_primitive();
    trace!(
        cigar = cluster_ts_alignment.cigar(),
        cost = cluster_ts_alignment_cost,
        "Cluster TS alignment"
    );

    let ts_datas = TSData::compute(&aux.ref_context_region, result)?;
    let ts_1_4_region = TSData::to_field_fallible(&ts_datas, SEP, |d| {
        if d.pos_1 <= d.pos_4 {
            GenomeRegion::from_incl_incl(d.pos_1.clone(), Some(d.pos_4.clone()))
        } else {
            GenomeRegion::from_incl_incl(d.pos_4.clone(), Some(d.pos_1.clone()))
        }
    })?;

    let allele_count = aux
        .output_phasing
        .map(|p| p.alleles.iter().filter(|a| **a == 1).count());

    Ok(CSVRecord {
        id: hash,
        cluster_id: aux.cluster_id,
        cluster_grp: aux.cluster_grp,
        ref_ctx_region: aux.ref_context_region.to_string(),
        alt_ctx_region: aux.alt_context_region.to_string(),
        cluster_region: aux.region.to_string(),
        ts_1_4_region,
        read_id: aux.alt_id,
        allele_count,
        fw_cigar: alignments.cluster_forward.cigar(),
        fw_cigar_ctx: aux.forward_alignment.cigar(),
        fw_mi_ctx: create_mi_string(&aux.forward_alignment, &aux.sequences.query)?,
        fw_cost: alignments.cluster_forward_cost,
        fw_cost_ctx: alignments.entire_forward_cost,
        ts_cigar: cluster_ts_alignment.cigar(),
        ts_cigar_ctx: alignments.entire_ts.cigar(),
        ts_cost: cluster_ts_alignment_cost,
        ts_cost_ctx: alignments.entire_ts_cost,
        ts_num: ts_datas.len(),
        ts_1_2: TSData::to_field(&ts_datas, SEP, |d| d.jump_1_2),
        ts_2_3: TSData::to_field(&ts_datas, SEP, |d| d.inner_len),
        ts_2_3_min: TSData::to_opt_field(
            &ts_datas,
            SEP,
            |d| d.er.as_ref(),
            |d, er| {
                d.inner_len
                    .saturating_sub(er.max_start.try_into().unwrap_or_default())
                    .saturating_sub((-er.min_end).try_into().unwrap_or_default())
            },
        ),
        ts_2_3_max: TSData::to_opt_field(
            &ts_datas,
            SEP,
            |d| d.er.as_ref(),
            |d, er| {
                d.inner_len
                    + usize::try_from(-er.min_start).unwrap_or_default()
                    + usize::try_from(er.max_end).unwrap_or_default()
            },
        ),
        ts_1_4: TSData::to_field(&ts_datas, SEP, |d| d.apg),
        ts_start_left_shift: TSData::to_opt_field(
            &ts_datas,
            SEP,
            |d| d.er.as_ref(),
            |_, er| -er.min_start,
        ),
        ts_start_right_shift: TSData::to_opt_field(
            &ts_datas,
            SEP,
            |d| d.er.as_ref(),
            |_, er| er.max_start,
        ),
        ts_end_left_shift: TSData::to_opt_field(
            &ts_datas,
            SEP,
            |d| d.er.as_ref(),
            |_, er| -er.min_end,
        ),
        ts_end_right_shift: TSData::to_opt_field(
            &ts_datas,
            SEP,
            |d| d.er.as_ref(),
            |_, er| er.max_end,
        ),
        ts_inner_alignment_cigar: TSData::to_field(&ts_datas, SEP, |d| d.inner_aln.cigar()),
        vcf_record_region: aux.vcf_record_region,
        ref_cluster_offset: result.stats.reference_offset(),
        // ref_cluster_limit: aux.ranges.reference_limit(),
        alt_cluster_offset: result.stats.query_offset(),
        // alt_cluster_limit: aux.ranges.query_limit(),
    })
}

pub fn compute_record_hash(result: &TwitcherAlignmentWithStatistics, aux: &CSVAuxData) -> u128 {
    let mut state = Xxh3Builder::new().build();
    aux.reference_name.hash(&mut state);
    aux.ref_context_region.hash(&mut state);
    aux.sequences.query.hash(&mut state);
    aux.alt_id.hash(&mut state);
    aux.region.hash(&mut state);
    result.alignment.alignment.hash(&mut state);
    state.digest128()
}

fn compute_cost(
    sequences: &SequencePair,
    sequence_offsets: Option<(usize, usize)>,
    mut alignment: Alignment<AlignmentType>,
    costs: &TemplateSwitchConfig<DnaAlphabetOrN, U64Cost>,
) -> anyhow::Result<u64> {
    let (r_off, q_off) = sequence_offsets.unwrap_or_default();
    // TODO this is copying data, which is a bit stupid but at the moment unavoidable due to the interfaces in lib_tsalign
    let r = VectorGenome::<DnaAlphabetOrN>::from_slice_u8(
        sequences
            .reference
            .get(r_off..)
            .with_context(|| format!("r_off {r_off} is out of bounds"))?,
    )?;
    let q = VectorGenome::from_slice_u8(
        sequences
            .query
            .get(q_off..)
            .with_context(|| format!("q_off {q_off} is out of bounds"))?,
    )?;
    let cost = alignment.compute_cost(
        r.as_genome_subsequence(),
        q.as_genome_subsequence(),
        0,
        0,
        costs,
    );
    Ok(cost.as_primitive())
}

fn create_mi_string(
    forward_alignment: &ForwardAlignment,
    query: &ImmutableSequence,
) -> anyhow::Result<String> {
    let mut qi = 0;
    let mut result = String::new();
    let mut counter = 0;
    for ty in forward_alignment.iter_flat() {
        match ty {
            AlignmentType::PrimaryInsertion | AlignmentType::PrimarySubstitution => {
                if counter > 0 {
                    result.push_str(&counter.to_string());
                    counter = 0;
                }
                result.push((*query.get(qi).context("query index out of bounds")?).into());
            }
            AlignmentType::PrimaryMatch => {
                counter += 1;
            }
            _ => {}
        }
        qi = qi
            .checked_add_signed(consumed_query(1, ty, None)?)
            .context("query index out of bounds")?;
    }

    Ok(result)
}