twitcher 0.6.9

use std::io::IsTerminal;

use anyhow::bail;
use clap_verbosity_flag::Verbosity;
use compact_genome::{
    implementation::{
        alphabets::dna_alphabet_or_n::DnaAlphabetOrN, bit_vec_sequence::BitVectorGenome,
    },
    interface::{
        alphabet::{Alphabet, AlphabetError},
        sequence::{GenomeSequence, OwnedGenomeSequence},
    },
};
use lib_tsalign::{
    a_star_aligner::{
        alignment_result::AlignmentResult,
        template_switch_distance::{
            AlignmentType,
            context::DynamicStrategies,
            strategies::{
                AlignmentStrategySelection,
                allow_ts_14_out_of_range::{
                    AdditionalExplicitTSMStartsAndEnds, Ts14OutOfRangeStrategy,
                },
                chaining::NoChainingStrategy,
                descendant::{
                    AnyTemplateSwitchDescendantStrategy, OnlyEqualTemplateSwitchDescendantStrategy,
                    TemplateSwitchDescendantStrategy,
                },
                node_ord::AntiDiagonalNodeOrdStrategy,
                primary_match::AllowPrimaryMatchStrategy,
                primary_range::NoPrunePrimaryRangeStrategy,
                secondary_deletion::AllowSecondaryDeletionStrategy,
                shortcut::NoShortcutStrategy,
                template_switch_count::{
                    MaxTemplateSwitchCountStrategy, NoTemplateSwitchCountStrategy,
                    TemplateSwitchCountStrategy,
                },
                template_switch_min_length::{
                    LookaheadTemplateSwitchMinLengthStrategy, NoTemplateSwitchMinLengthStrategy,
                    PreprocessedLookaheadTemplateSwitchMinLengthStrategy,
                    PreprocessedTemplateSwitchMinLengthStrategy, TemplateSwitchMinLengthStrategy,
                },
                template_switch_total_length::MaxTemplateSwitchTotalLengthStrategy,
            },
        },
        template_switch_distance_a_star_align,
    },
    config::TemplateSwitchConfig,
    costs::U64Cost,
};
use rlimit::Resource;
use serde::{Deserialize, Serialize};
use tracing::{error, info_span, instrument};

use crate::{
    VerbositySelector,
    common::{
        SequencePair,
        aligner::{AlignerSelector, AlignmentQuery, cli::MLSSelector, result},
        coords::GenomeRegion,
    },
    setup_tracing,
};

// Cow but without a `ToOwned` bound
pub enum RefOrOwned<'a, T: 'a> {
    Ref(&'a T),
    Owned(T),
}

impl<T> std::ops::Deref for RefOrOwned<'_, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        match self {
            RefOrOwned::Ref(r) => r,
            RefOrOwned::Owned(o) => o,
        }
    }
}

impl<T: Serialize> Serialize for RefOrOwned<'_, T> {
    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        match self {
            RefOrOwned::Ref(r) => r.serialize(serializer),
            RefOrOwned::Owned(o) => o.serialize(serializer),
        }
    }
}

impl<'de, T: Deserialize<'de>> Deserialize<'de> for RefOrOwned<'_, T> {
    fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        T::deserialize(deserializer).map(RefOrOwned::Owned)
    }
}

impl<T> From<T> for RefOrOwned<'_, T> {
    fn from(value: T) -> Self {
        Self::Owned(value)
    }
}

impl<'a, T> From<&'a T> for RefOrOwned<'a, T> {
    fn from(value: &'a T) -> Self {
        Self::Ref(value)
    }
}

#[derive(Serialize, Deserialize)]
pub struct WorkerQuery<'a> {
    pub aligner: RefOrOwned<'a, AlignerSelector>,
    pub log_level: VerbositySelector,
    pub memory: usize,
    pub query: AlignmentQuery,
    pub metadata: WorkerQueryMetadata,
}

#[derive(Serialize, Deserialize)]
pub struct WorkerQueryMetadata {
    pub cluster_region: GenomeRegion,
}

/// Read the query from stdin and write the result to stdout.
/// Errors are reported either to stderr, through non-zero exit or abort.
#[instrument]
pub fn run_fallible() -> anyhow::Result<()> {
    let wq = read_from_stdin()?;

    setup_tracing(wq.log_level)?;
    let span = info_span!("Aligning sequences", pos = %wq.metadata.cluster_region).entered();

    let bytes = wq.memory as u64;
    rlimit::setrlimit(Resource::AS, bytes, bytes)
        .inspect_err(|e| error!("Can't set resource limit: {e}"))?;
    let WorkerQuery {
        aligner,
        memory,
        query,
        ..
    } = wq;

    let result = match &*aligner {
        AlignerSelector::AStar {
            costs,
            min_length_strategy,
            allow_mixed_descendants,
            no_ts,
        } => {
            let with_ts = align_select_min_length_strategy(AStarQueryData {
                memory,
                query,
                costs,
                min_length_strategy,
                allow_mixed_descendants: *allow_mixed_descendants,
                no_ts: *no_ts,
            })?;
            result::from_tsalign(with_ts)
        }
        AlignerSelector::Fpa(four_point_aligner) => {
            // fpa, and especially the donstream tsalign alignment methods that we use to post-process the alignment expects valid sequences
            verify_sequences(&query.sequences)?;
            four_point_aligner.align(
                query.sequences.reference,
                query.sequences.query,
                query.ranges,
            )
        }
    };
    span.exit();

    let _span = info_span!("Reporting alignment", pos = %wq.metadata.cluster_region).entered();
    let mut out = std::io::stdout().lock();
    rmp_serde::encode::write(&mut out, &result)?;
    drop(out);

    Ok(())
}

fn verify_sequences(seqs: &SequencePair) -> anyhow::Result<()> {
    for (s, bs) in [("reference", &seqs.reference), ("query", &seqs.query)] {
        if let Some(e) = bs
            .iter()
            .find_map(|c| DnaAlphabetOrN::ascii_to_character(*c).err())
        {
            anyhow::bail!("{s} sequence contains an unsupported ASCII character: {e}")
        }
    }
    Ok(())
}

fn read_from_stdin() -> anyhow::Result<WorkerQuery<'static>> {
    let mut stdin = std::io::stdin().lock();
    if stdin.is_terminal() {
        setup_tracing(Verbosity::default())?;
        bail!(
            "The `worker` subcommand is not designed to be called directly but only as a subroutine of twitcher."
        );
    }

    let wq = rmp_serde::from_read(&mut stdin)?;
    Ok(wq)
}

struct AStarQueryData<'a> {
    memory: usize,
    query: AlignmentQuery,
    costs: &'a TemplateSwitchConfig<DnaAlphabetOrN, U64Cost>,
    min_length_strategy: &'a MLSSelector,
    allow_mixed_descendants: bool,
    no_ts: bool,
}

fn align_select_min_length_strategy(
    data: AStarQueryData,
) -> Result<AlignmentResult<AlignmentType, U64Cost>, AlphabetError> {
    match data.min_length_strategy {
        MLSSelector::None => {
            align_select_descendant_strategy::<NoTemplateSwitchMinLengthStrategy<U64Cost>>(data)
        }
        MLSSelector::Lookahead => align_select_descendant_strategy::<
            LookaheadTemplateSwitchMinLengthStrategy<U64Cost>,
        >(data),
        MLSSelector::PreprocessFilter => align_select_descendant_strategy::<
            PreprocessedTemplateSwitchMinLengthStrategy<true, U64Cost>,
        >(data),
        MLSSelector::PreprocessPrice => align_select_descendant_strategy::<
            PreprocessedTemplateSwitchMinLengthStrategy<false, U64Cost>,
        >(data),
        MLSSelector::PreprocessLookahead => align_select_descendant_strategy::<
            PreprocessedLookaheadTemplateSwitchMinLengthStrategy<U64Cost>,
        >(data),
    }
}

fn align_select_descendant_strategy<ML: TemplateSwitchMinLengthStrategy<U64Cost>>(
    data: AStarQueryData,
) -> Result<AlignmentResult<AlignmentType, U64Cost>, AlphabetError> {
    if data.allow_mixed_descendants {
        align_select_no_ts_strategy::<ML, AnyTemplateSwitchDescendantStrategy>(data)
    } else {
        align_select_no_ts_strategy::<ML, OnlyEqualTemplateSwitchDescendantStrategy>(data)
    }
}

fn align_select_no_ts_strategy<
    ML: TemplateSwitchMinLengthStrategy<U64Cost>,
    DS: TemplateSwitchDescendantStrategy,
>(
    data: AStarQueryData,
) -> Result<AlignmentResult<AlignmentType, U64Cost>, AlphabetError> {
    if data.no_ts {
        align_call::<ML, DS, MaxTemplateSwitchCountStrategy>(data, 0)
    } else {
        align_call::<ML, DS, NoTemplateSwitchCountStrategy>(data, ())
    }
}

fn align_call<
    ML: TemplateSwitchMinLengthStrategy<U64Cost>,
    DS: TemplateSwitchDescendantStrategy,
    TS: TemplateSwitchCountStrategy,
>(
    data: AStarQueryData,
    count_mem: TS::Memory,
) -> Result<AlignmentResult<AlignmentType, U64Cost>, AlphabetError> {
    let reference = seq_to_genome(&data.query.sequences.reference)?;
    let query = seq_to_genome(&data.query.sequences.query)?;
    Ok(template_switch_distance_a_star_align::<
        AlignmentStrategySelection<
            DnaAlphabetOrN,
            U64Cost,
            AntiDiagonalNodeOrdStrategy,
            ML,
            NoChainingStrategy<U64Cost>,
            TS,
            AllowSecondaryDeletionStrategy,
            NoShortcutStrategy<U64Cost>,
            AllowPrimaryMatchStrategy,
            // The alignment range only marks the region of interest; the padded context around
            // it is there so that template switches can jump into it. Pruning the primary
            // alignment to the range would strand the alignment whenever a template switch
            // exits outside of it, forcing artificial extra template switches.
            // TODO probably we want to fix this upstream and then remove this line here.
            NoPrunePrimaryRangeStrategy,
            MaxTemplateSwitchTotalLengthStrategy,
            DS,
        >,
        _,
    >(
        reference.as_genome_subsequence(),
        query.as_genome_subsequence(),
        "reference",
        "query",
        data.query.ranges,
        AdditionalExplicitTSMStartsAndEnds::default(),
        data.costs,
        DynamicStrategies {
            ts_14_out_of_range: Ts14OutOfRangeStrategy::default(),
        },
        None,
        Some(data.memory),
        false,
        false,
        count_mem,
    ))
}

/// `verify_query` MUST be called before running this to ensure that the sequences are valid DNA+N
fn seq_to_genome(seq: &[u8]) -> Result<BitVectorGenome<DnaAlphabetOrN>, AlphabetError> {
    let bvg = BitVectorGenome::from_iter_u8(seq.iter().copied())?;
    Ok(bvg)
}