mehari 0.47.0

//! Consequence prediction for structural variants.

use std::{collections::HashMap, sync::Arc};

use crate::{
    annotate::seqvars::provider::Provider as MehariProvider,
    pbs::txs::{Strand, Transcript},
};

/// Enumeration for effect on transcript.
#[derive(
    serde::Serialize, serde::Deserialize, PartialEq, Eq, PartialOrd, Ord, Debug, Clone, Copy,
)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
#[serde(rename_all = "snake_case")]
// #[schema(as=StrucvarsTranscriptEffect)] // TODO: rename back to TranscriptEffect once utoipa's as= is fixed.
pub enum StrucvarsTranscriptEffect {
    /// Affects the full transcript.
    TranscriptVariant,
    /// An exon is affected by the SV.
    ExonVariant,
    /// The splice region is affected by the SV.
    SpliceRegionVariant,
    /// The intron is affected by the SV.
    IntronVariant,
    /// The upstream region of the transcript is affected.
    UpstreamVariant,
    /// The downstream region of the transcript is affected.
    DownstreamVariant,
    /// Only intergenic regions is affected,
    IntergenicVariant,
}

impl StrucvarsTranscriptEffect {
    /// Return vector with all transcript effects.
    pub fn vec_all() -> Vec<StrucvarsTranscriptEffect> {
        vec![
            StrucvarsTranscriptEffect::TranscriptVariant,
            StrucvarsTranscriptEffect::ExonVariant,
            StrucvarsTranscriptEffect::SpliceRegionVariant,
            StrucvarsTranscriptEffect::IntronVariant,
            StrucvarsTranscriptEffect::UpstreamVariant,
            StrucvarsTranscriptEffect::DownstreamVariant,
            StrucvarsTranscriptEffect::IntergenicVariant,
        ]
    }
}

/// Helpful types for using structural variants.
pub mod interface {
    /// Structural Variant type.
    #[derive(
        serde::Serialize,
        serde::Deserialize,
        Debug,
        Clone,
        Copy,
        PartialEq,
        Eq,
        PartialOrd,
        Ord,
        Hash,
    )]
    #[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
    // #[schema(as=StrucvarsSvType)] // TODO: rename back to StrucVarType once utoipa's as= is fixed.
    pub enum StrucvarsSvType {
        #[serde(rename = "DEL")]
        Del,
        #[serde(rename = "DUP")]
        Dup,
        #[serde(rename = "INS")]
        Ins,
        #[serde(rename = "INV")]
        Inv,
        #[serde(rename = "BND")]
        Bnd,
        #[serde(rename = "CNV")]
        Cnv,
    }

    /// Strand orientation of a structural variant.
    #[derive(
        serde::Serialize,
        serde::Deserialize,
        Debug,
        Clone,
        Default,
        PartialEq,
        Eq,
        PartialOrd,
        Ord,
        Hash,
    )]
    pub enum StrandOrientation {
        #[serde(rename = "3to3")]
        ThreeToThree,
        #[serde(rename = "5to5")]
        FiveToFive,
        #[serde(rename = "3to5")]
        ThreeToFive,
        #[serde(rename = "5to3")]
        FiveToThree,
        #[serde(rename = "NtoN")]
        #[default]
        NotApplicable,
    }

    /// Trait for generic description of structural variant.
    ///
    /// The values of ``stop()`` and ``chrom2()`` are unused then they are set to the
    /// corresponding other value: ``start()`` and ``chrom()``.
    pub trait StrucVar {
        /// Chromosome name.
        fn chrom(&self) -> String;
        /// The second involved chromosome, or the same value as `chrom()`.
        fn chrom2(&self) -> String;

        /// 1-based start position of the variant (or position on first chromosome for break-ends)
        fn start(&self) -> i32;
        /// 1-based end position of the variant.
        fn stop(&self) -> i32;

        /// Type of the structural variant
        fn sv_type(&self) -> StrucvarsSvType;
        /// The strand orientation of the structural variant, if applicable.
        fn strand_orientation(&self) -> StrandOrientation;
    }
}

/// Ad-hoc data structure for `tx_regions`.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
struct TxRegion {
    // 0-based begin position
    begin: i32,
    // 0-based end position
    end: i32,
    // "arbitrary" number
    no: usize,
    // effect of the transcript (encodes region type)
    effect: StrucvarsTranscriptEffect,
}

/// Explanation of transcript effect per individual gene.
#[derive(serde::Serialize, serde::Deserialize, Debug, Clone)]
#[cfg_attr(feature = "server", derive(utoipa::ToSchema))]
// #[schema(as=StrucvarsGeneTranscriptEffects)] // TODO: rename back to GeneTranscriptEffects once utoipa's as= is fixed.
pub struct StrucvarsGeneTranscriptEffects {
    /// Gene ID of the annotation without version, e.g. `"672"` for BRCA1 in RefSeq or
    /// `"ENSG00000182378"` for PLCXD1 in GENCODE.
    gene_id: String,
    /// HGNC ID of the gene, e.g. `"HGNC:1100"` for BRCA1. Unset for genes without HGNC ID.
    hgnc_id: Option<String>,
    /// Transcript effects for the gene.
    transcript_effects: Vec<StrucvarsTranscriptEffect>,
}

/// Length of the upstream/downstream region.
static X_STREAM: i32 = 5000;

/// Return list of half-open intervals for a given transcript.
fn tx_regions(tx: &Transcript) -> Vec<TxRegion> {
    assert_eq!(
        tx.genome_alignments.len(),
        1,
        "only one alignment supported"
    );
    if tx.genome_alignments[0].exons.is_empty() {
        // no exons? skip!
        return Vec::new();
    }

    let mut result = Vec::new();
    let mut tx_start = None;
    let mut tx_end = None;
    let genome_alignment = &tx.genome_alignments[0];

    // Loop over all exons to determine leftmost and rightmost genome position.
    for exon_alignment in &genome_alignment.exons {
        if let Some(value) = tx_start {
            if exon_alignment.alt_start_i < value {
                tx_start = Some(exon_alignment.alt_start_i);
            }
        } else {
            tx_start = Some(exon_alignment.alt_start_i);
        }
        if let Some(value) = tx_end {
            if exon_alignment.alt_end_i > value {
                tx_end = Some(exon_alignment.alt_end_i);
            }
        } else {
            tx_end = Some(exon_alignment.alt_end_i);
        }
    }

    let tx_start = tx_start.expect("must have been set");
    let tx_end = tx_end.expect("must have been set");

    // Perform the actual region extraction.
    let mut prev_alt_end_i = 0;
    for (no, exon_alignment) in genome_alignment.exons.iter().enumerate() {
        if exon_alignment.alt_start_i == tx_start {
            // is first, register upstream/downstream
            result.push(TxRegion {
                begin: exon_alignment.alt_start_i - X_STREAM,
                end: exon_alignment.alt_start_i - 1,
                no,
                effect: if genome_alignment.strand == Strand::Plus as i32 {
                    StrucvarsTranscriptEffect::UpstreamVariant
                } else {
                    StrucvarsTranscriptEffect::DownstreamVariant
                },
            });
        } else {
            // is not first, register splice region on the left boundary
            result.push(TxRegion {
                begin: (exon_alignment.alt_start_i - 1) - 8,
                end: (exon_alignment.alt_start_i - 1) + 3,
                no,
                effect: StrucvarsTranscriptEffect::SpliceRegionVariant,
            })
        }

        if exon_alignment.alt_end_i == tx_end {
            // is last, register upstream/downstream
            result.push(TxRegion {
                begin: exon_alignment.alt_end_i,
                end: exon_alignment.alt_end_i + X_STREAM,
                no,
                effect: if genome_alignment.strand == Strand::Plus as i32 {
                    StrucvarsTranscriptEffect::DownstreamVariant
                } else {
                    StrucvarsTranscriptEffect::UpstreamVariant
                },
            });
        } else {
            // is not last, register splice region on the right boundary
            result.push(TxRegion {
                begin: exon_alignment.alt_end_i - 3,
                end: exon_alignment.alt_end_i + 8,
                no,
                effect: StrucvarsTranscriptEffect::SpliceRegionVariant,
            })
        }

        // register the exon
        result.push(TxRegion {
            begin: exon_alignment.alt_start_i - 1,
            end: exon_alignment.alt_end_i,
            no,
            effect: StrucvarsTranscriptEffect::ExonVariant,
        });

        if exon_alignment.alt_start_i != tx_start {
            // is not first exon, register intron "right" of it
            result.push(TxRegion {
                begin: prev_alt_end_i,
                end: exon_alignment.alt_start_i - 1,
                no,
                effect: StrucvarsTranscriptEffect::IntronVariant,
            });
        }

        // store end of prev exon for next intron's start
        prev_alt_end_i = exon_alignment.alt_end_i;
    }

    result
}

/// Return the transcript region / effect for the given breakpoint.
fn gene_tx_effects_for_bp(tx: &Transcript, pos: i32) -> Vec<StrucvarsTranscriptEffect> {
    // Obtain list of regions for transcript.
    let regions = tx_regions(tx);

    // Determine how this relates to the breakpoint.
    let pos = pos - 1; // 1-based to 0-based
    let mut result = regions
        .iter()
        .filter(|r| r.begin <= pos && pos < r.end)
        .map(|r| r.effect)
        .collect::<Vec<_>>();
    if result.is_empty() {
        result.push(StrucvarsTranscriptEffect::IntergenicVariant);
    } else {
        result.sort();
        result.dedup();
    }
    result
}

/// Return the transcript region / effect for the given range.
fn gene_tx_effect_for_range(
    tx: &Transcript,
    start: i32,
    stop: i32,
) -> Vec<StrucvarsTranscriptEffect> {
    // Obtain list of regions for transcript.
    let regions = tx_regions(tx);

    // Determine how this relates to the left and right breakpoints.
    let pos = start - 1; // 1-based to 0-based
    let mut result = regions
        .iter()
        .filter(|region| pos <= region.end && region.begin <= stop)
        .map(|region| region.effect)
        .collect::<Vec<_>>();

    // Remove any duplicates.
    result.sort();
    result.dedup();

    // If we have both upstream and downstream then the full transcript is affected.
    if result.contains(&StrucvarsTranscriptEffect::UpstreamVariant)
        && result.contains(&StrucvarsTranscriptEffect::DownstreamVariant)
    {
        result.push(StrucvarsTranscriptEffect::TranscriptVariant);
    }

    result
}

/// Group the effects on transcripts by gene, with one entry per gene entry. Transcripts
/// without gene entry give one entry per gene ID.
fn effects_by_gene<'a>(
    mehari_provider: &'a MehariProvider,
    tx_effects: impl Iterator<Item = (&'a Transcript, Vec<StrucvarsTranscriptEffect>)>,
) -> Vec<StrucvarsGeneTranscriptEffects> {
    let mut effects_by_gene: HashMap<_, Vec<_>> = HashMap::new();
    for (tx, effects) in tx_effects {
        let gene = if let Some(gene) = mehari_provider.gene_of(&tx.id) {
            (gene.gene_id.as_str(), gene.hgnc_id.as_deref())
        } else {
            (tx.gene_id.as_str(), tx.hgnc_id.as_deref())
        };
        effects_by_gene.entry(gene).or_default().extend(effects);
    }

    effects_by_gene
        .into_iter()
        .map(|((gene_id, hgnc_id), mut transcript_effects)| {
            transcript_effects.sort();
            transcript_effects.dedup();
            StrucvarsGeneTranscriptEffects {
                gene_id: gene_id.into(),
                hgnc_id: hgnc_id.map(String::from),
                transcript_effects,
            }
        })
        .collect()
}

/// Helper that computes effects on transcripts for a single breakend, e.g., one side of BND or INS.
fn compute_tx_effects_for_breakpoint(
    sv: &impl interface::StrucVar,
    mehari_provider: &MehariProvider,
) -> Vec<StrucvarsGeneTranscriptEffects> {
    // Shortcut to the `TranscriptDb`.
    let tx_db = mehari_provider
        .tx_seq_db
        .tx_db
        .as_ref()
        .expect("transcripts must be present");
    // Compute canonical chromosome name and map to accession.
    let chrom_acc = mehari_provider.contig_manager.get_accession(&sv.chrom());
    if chrom_acc.is_none() {
        return Default::default();
    }
    let chrom_acc = chrom_acc.expect("chromosome must be known at this point");
    // Create range to query the interval trees for.
    let query = (sv.start() - X_STREAM)..(sv.start() + X_STREAM);

    if let Some(idx) = mehari_provider.tx_trees.contig_to_idx.get(chrom_acc) {
        // Collect all transcripts that overlap the INS and compute the effect of the INS on
        // the transcript.
        let tree = &mehari_provider.tx_trees.trees[*idx];
        effects_by_gene(
            mehari_provider,
            tree.find(query).into_iter().map(|it| {
                let tx = &tx_db.transcripts[*it.data() as usize];
                (tx, gene_tx_effects_for_bp(tx, sv.start()))
            }),
        )
    } else {
        // We do not have any transcripts for this chromosome.
        Default::default()
    }
}

/// Compute effect for linear SVs.
fn compute_tx_effects_for_linear(
    sv: &impl interface::StrucVar,
    mehari_provider: &MehariProvider,
) -> Vec<StrucvarsGeneTranscriptEffects> {
    // Shortcut to the `TranscriptDb`.
    let tx_db = mehari_provider
        .tx_seq_db
        .tx_db
        .as_ref()
        .expect("transcripts must be present");
    // Compute canonical chromosome name and map to accession.
    let chrom_acc = mehari_provider.contig_manager.get_accession(&sv.chrom());
    if chrom_acc.is_none() {
        return Default::default();
    }
    let chrom_acc = chrom_acc.expect("chromosome must be known at this point");
    // Create range to query the interval trees for.
    let query = (sv.start() - X_STREAM)..(sv.stop() + X_STREAM);

    if let Some(idx) = mehari_provider.tx_trees.contig_to_idx.get(chrom_acc) {
        // Collect all transcripts that overlap the SV and compute the effect of the SV on
        // the transcript.
        let tree = &mehari_provider.tx_trees.trees[*idx];
        effects_by_gene(
            mehari_provider,
            tree.find(query).into_iter().map(|it| {
                let tx = &tx_db.transcripts[*it.data() as usize];
                (tx, gene_tx_effect_for_range(tx, sv.start(), sv.stop()))
            }),
        )
    } else {
        // We do not have any transcripts for this chromosome.
        Default::default()
    }
}

/// Wrap mapper, provider, and map for consequence prediction.
#[derive(derivative::Derivative)]
#[derivative(Debug)]
pub struct ConsequencePredictor {
    /// The internal transcript provider for locating transcripts.
    #[derivative(Debug = "ignore")]
    provider: Arc<MehariProvider>,
}

impl ConsequencePredictor {
    pub fn new(provider: Arc<MehariProvider>) -> Self {
        ConsequencePredictor { provider }
    }

    /// Compute effect(s) of `sv` on transcript of genes.
    pub fn compute_tx_effects(
        &self,
        sv: &impl interface::StrucVar,
    ) -> Vec<StrucvarsGeneTranscriptEffects> {
        match sv.sv_type() {
            interface::StrucvarsSvType::Ins | interface::StrucvarsSvType::Bnd => {
                compute_tx_effects_for_breakpoint(sv, &self.provider)
            }
            interface::StrucvarsSvType::Del
            | interface::StrucvarsSvType::Dup
            | interface::StrucvarsSvType::Inv
            | interface::StrucvarsSvType::Cnv => compute_tx_effects_for_linear(sv, &self.provider),
        }
    }

    /// Return data version string (if set).
    pub fn data_version(&self) -> Option<String> {
        self.provider.as_ref().tx_seq_db.version.clone()
    }
}

#[cfg(test)]
mod test {
    use super::interface::{StrandOrientation, StrucVar, StrucvarsSvType};

    /// A deletion on chr1.
    struct Deletion {
        start: i32,
        stop: i32,
    }

    impl StrucVar for Deletion {
        fn chrom(&self) -> String {
            "1".into()
        }

        fn chrom2(&self) -> String {
            "1".into()
        }

        fn start(&self) -> i32 {
            self.start
        }

        fn stop(&self) -> i32 {
            self.stop
        }

        fn sv_type(&self) -> StrucvarsSvType {
            StrucvarsSvType::Del
        }

        fn strand_orientation(&self) -> StrandOrientation {
            StrandOrientation::NotApplicable
        }
    }

    /// The effects form one entry per gene entry, with its gene ID and HGNC ID. Here, two
    /// transcripts come from two versions of one gene, and only one names the HGNC ID.
    #[test]
    fn one_entry_per_gene() {
        use crate::annotate::seqvars::provider::{Provider, test_db};
        use crate::pbs::txs::{GeneToTxId, Transcript};
        use std::sync::Arc;

        let txs = [
            ("ENST00000000001.1", "ENSG00000000001.9", Some("HGNC:1")),
            ("ENST00000000001.2", "ENSG00000000001.10", None),
        ]
        .map(|(id, gene_id, hgnc_id)| Transcript {
            gene_id: gene_id.into(),
            hgnc_id: hgnc_id.map(String::from),
            ..test_db::transcript(id, 1000)
        });
        let gene = GeneToTxId {
            gene_id: "ENSG00000000001".into(),
            hgnc_id: Some("HGNC:1".into()),
            tx_ids: txs.iter().map(|tx| tx.id.clone()).collect(),
            filtered: Some(false),
            ..Default::default()
        };
        let provider = Provider::new(
            test_db::tx_seq_db(txs.to_vec(), vec![gene]),
            None::<&str>,
            true,
        );

        let effects =
            super::ConsequencePredictor::new(Arc::new(provider)).compute_tx_effects(&Deletion {
                start: 900,
                stop: 1200,
            });
        let mut genes: Vec<_> = effects
            .iter()
            .map(|gene| (gene.gene_id.as_str(), gene.hgnc_id.as_deref()))
            .collect();
        genes.sort();
        assert_eq!(genes, [("ENSG00000000001", Some("HGNC:1"))]);
    }

    /// A transcript without gene entry forms an entry with its own gene ID and HGNC ID.
    #[test]
    fn transcript_without_gene_entry() {
        use crate::annotate::seqvars::provider::{Provider, test_db};
        use crate::pbs::txs::Transcript;
        use std::sync::Arc;

        let tx = Transcript {
            gene_id: "11".into(),
            hgnc_id: Some("HGNC:1".into()),
            ..test_db::transcript("NM_000001.1", 1000)
        };
        let provider = Provider::new(test_db::tx_seq_db(vec![tx], vec![]), None::<&str>, true);

        let effects =
            super::ConsequencePredictor::new(Arc::new(provider)).compute_tx_effects(&Deletion {
                start: 900,
                stop: 1200,
            });
        let genes: Vec<_> = effects
            .iter()
            .map(|gene| (gene.gene_id.as_str(), gene.hgnc_id.as_deref()))
            .collect();
        assert_eq!(genes, [("11", Some("HGNC:1"))]);
    }
}