termal-alignment 0.1.0

// SPDX-License-Identifier: MIT
// Copyright (c) 2025-2026 Thomas Junier

mod permutation;

use std::{collections::HashMap, fmt};

use itertools::Itertools;

use crate::seq::file::SeqFile;

use crate::alignment::SeqType::{Nucleic, Protein};

// Whether to show the most frequent residue as LC or UC
const UC_CONS_THRESHOLD: f64 = 0.8; // uppercase if at least this
const LC_CONS_THRESHOLD: f64 = 0.2; // lowercase if at least this (else '*')

type ResidueDistribution = HashMap<char, f64>;
type ResidueCounts = HashMap<char, u64>;

#[derive(PartialEq, Clone, Copy, Debug)]
pub enum SeqType {
    Nucleic,
    Protein,
}

#[derive(Clone, Copy, Debug, PartialEq)]
pub enum RefSpec {
    Consensus,
    Rank(usize),
}

pub enum RefSpecError {
    MalformedInt(String),
    ZeroRef,
    RefTooLarge(usize),
}

impl fmt::Display for RefSpecError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let err_msg = match self {
            RefSpecError::MalformedInt(mfi) => format!("Malformed integer {}", mfi),
            RefSpecError::ZeroRef => "Ref # must be > 0".to_string(),
            RefSpecError::RefTooLarge(max) => format!("Ref # too large (max {})", max),
        };
        write!(f, "{}", err_msg)
    }
}

pub struct Alignment {
    pub headers: Vec<String>,
    pub sequences: Vec<String>,
    /* The consensus sequence is now a field of Alignment, and is computed once upon creation. This
     * contrasts with the very first implementation, in which the consensus was recomputed every
     * time the UI was drawn... which was very inefficient but had this funny "twinkling" effect in
     * columns with tied residue frequencies. This was due to the fact that HashMap stores its keys
     * in an unpredictable order, and that different calls to keys() may return them indifferent
     * orders. See best_residue().
     */
    /* These are properties of the whole _alignment_, or at least of whole columns. They cannot be
     * meaningfully attributed to a sequence. */
    pub consensus: String,
    pub entropies: Vec<f64>,
    pub densities: Vec<f64>,

    /* By contrast, the following are properties of sequences (at least in part). Length, for
     * example, does not depend on anything but the sequence itself, and could be a field in a
     * struct that also contains the sequence and its header. */
    pub id_wrt_reference: Vec<f64>, // reference is usually the consensus, but CAN be an aln seq.
    // Recompute if ref changes.
    // Of course the sequence length is an integer, but using an integer type like u32 would make
    // it hard (for me, at least...) to write a function that accepts a Vec of either  lengths or
    // %IDs. Tried Box, and generics, but the extra work doesn't seem warranted.
    pub relative_seq_len: Vec<f64>,
    pub macromolecule_type: SeqType,
    /* Specifies whether the reference sequence should be the consensus (see above) or one of the
     * original sequences (identified by rank).*/
    ref_spec: RefSpec,
}

#[derive(Debug, PartialEq)]
struct BestResidue {
    residue: char,
    frequency: u64,
}

impl Alignment {
    // Makes an Alignment from a SeqFile, which is consumed.
    pub fn from_file(seq_file: SeqFile) -> Alignment {
        let mut headers: Vec<String> = Vec::new();
        let mut sequences: Vec<String> = Vec::new();
        let mut max_len: usize = 0;
        for record in seq_file {
            headers.push(record.header);
            let l = record.sequence.len();
            sequences.push(record.sequence);
            if l > max_len {
                max_len = l;
            }
        }
        // Pad any sequence shorter than max_len, so we are not limited to alignments with exactly
        // identical numbers of positions (reviewer suggestion).
        sequences
            .iter_mut()
            .for_each(|s| *s = format!("{:<width$}", s, width = max_len));
        // NOTE: the 's' can also be written '&*s', which makes the automatic re-borrow explicit.
        let consensus = consensus(&sequences);
        let entropies = entropies(&sequences);
        let densities = densities(&sequences);
        let id_wrt_reference = sequences
            .iter()
            .map(|seq| percent_identity(seq, &consensus))
            .collect();
        let relative_seq_len = sequences.iter().map(|seq| seq_len_nogaps(seq)).collect();
        let first_seq = sequences.first();
        let macromolecule_type = seq_type(first_seq.expect("No sequence found."));

        Alignment {
            headers,
            sequences,
            consensus,
            entropies,
            densities,
            id_wrt_reference,
            relative_seq_len,
            macromolecule_type,
            ref_spec: RefSpec::Consensus,
        }
    }

    // Makes an Alignment from a Vec of headers and a Vec of Strings, which are consumed. Mostly
    // used for testing.
    #[allow(dead_code)]
    pub fn from_vecs(hdrs: Vec<String>, seqs: Vec<String>) -> Alignment {
        assert_eq!(hdrs.len(), seqs.len());
        let headers = hdrs;
        let sequences = seqs;
        let consensus = consensus(&sequences);
        let entropies = entropies(&sequences);
        let densities = densities(&sequences);
        let id_wrt_reference = sequences
            .iter()
            .map(|seq| percent_identity(seq, &consensus))
            .collect();
        let relative_seq_len = sequences.iter().map(|seq| seq_len_nogaps(seq)).collect();
        let first_seq = sequences.first();
        let macromolecule_type = seq_type(first_seq.expect("No sequence found."));

        Alignment {
            headers,
            sequences,
            consensus,
            entropies,
            densities,
            id_wrt_reference,
            relative_seq_len,
            macromolecule_type,
            ref_spec: RefSpec::Consensus,
        }
    }

    pub fn num_seq(&self) -> usize {
        self.sequences.len()
    }

    // TODO: shouldn't this be aln_width?
    pub fn aln_len(&self) -> usize {
        self.sequences[0].len()
    }

    pub fn macromolecule_type(&self) -> SeqType {
        self.macromolecule_type
    }

    pub fn get_ref_spec(&self) -> RefSpec {
        self.ref_spec
    }

    pub fn set_ref_spec(&mut self, spec: RefSpec) -> Result<(), RefSpecError> {
        match spec {
            // Note: the rank in a RefSpec is 0-based. The conversion from user-land is done in
            // app.rs.
            RefSpec::Rank(rk) if rk >= self.num_seq() => {
                return Err(RefSpecError::RefTooLarge(self.num_seq()));
            }
            _ => self.ref_spec = spec,
        }
        // Probable change of ref -> Recompute the identities WRT ref
        let reference = self.reference();
        self.id_wrt_reference = self
            .sequences
            .iter()
            .map(|seq| percent_identity(seq, &reference))
            .collect();
        Ok(())
    }

    pub fn reference(&self) -> String {
        match self.ref_spec {
            RefSpec::Consensus => self.consensus.clone(),
            RefSpec::Rank(rk) => self.sequences[rk].clone(),
        }
    }
}

// TODO should these be methods of Alignment?

fn res_count(sequences: &Vec<String>, col: usize) -> ResidueCounts {
    let mut freqs: ResidueCounts = HashMap::new();
    for seq in sequences {
        let residue = seq.as_bytes()[col] as char;
        *freqs.entry(residue).or_insert(0) += 1;
    }
    freqs
}

pub fn consensus(sequences: &Vec<String>) -> String {
    let mut consensus = String::new();
    for j in 0..sequences[0].len() {
        let dist = res_count(sequences, j); // res -> count map
        let br = best_residue(&dist);
        let rel_freq: f64 = (br.frequency as f64 / sequences.len() as f64) as f64;
        if rel_freq >= UC_CONS_THRESHOLD {
            consensus.push(br.residue.to_ascii_uppercase());
        } else if rel_freq >= LC_CONS_THRESHOLD {
            if br.residue.is_alphabetic() {
                consensus.push(br.residue.to_ascii_lowercase());
            } else {
                //consensus.push('-');
                consensus.push(br.residue);
            }
        } else {
            consensus.push('*');
        }
    }
    consensus
}

pub fn entropies(sequences: &Vec<String>) -> Vec<f64> {
    let mut entropies: Vec<f64> = Vec::new();
    for j in 0..sequences[0].len() {
        let dist = res_count(sequences, j);
        let freq = to_freq_distrib(&dist);
        let e = entropy(&freq);
        entropies.push(e);
    }
    entropies
}

pub fn col_density(sequences: &Vec<String>, col: usize) -> f64 {
    let mut mass = 0;
    for seq in sequences {
        match seq.as_bytes()[col] as char {
            'a'..='z' | 'A'..='Z' => mass += 1,
            '-' | '.' | ' ' => {}
            other => {
                panic!("Character {other} unexpected in an alignment.\nThis might be due to file format, please see option -f.");
            }
        }
    }
    mass as f64 / sequences.len() as f64
}

pub fn densities(sequences: &Vec<String>) -> Vec<f64> {
    (0..sequences[0].len())
        .map(|col| col_density(sequences, col))
        .collect()
}

fn best_residue(dist: &ResidueCounts) -> BestResidue {
    let max_freq = dist.values().max().unwrap();
    let most_frequent_residue = dist
        .keys()
        .find(|&&k| dist.get(&k) == Some(max_freq))
        .unwrap();

    BestResidue {
        residue: *most_frequent_residue,
        frequency: *max_freq,
    }
}

// Convert a residue -> count map into a residue -> frequency map (relative frequency, that is).
// While gaps are allowed (and indeed useful) in the former, they are not included in the latter
// (in particular because they make litle sense when computing entropy).
//
fn to_freq_distrib(counts: &ResidueCounts) -> ResidueDistribution {
    let total_counts: u64 = counts
        .iter()
        .filter(|(res, _count)| **res != '-')
        .map(|(_res, count)| count)
        .sum();
    let mut distrib = ResidueDistribution::new();
    for (residue, count) in counts.iter() {
        if *residue == '-' {
            continue;
        }
        distrib.insert(*residue, *count as f64 / total_counts as f64);
    }
    distrib
}

fn entropy(freqs: &ResidueDistribution) -> f64 {
    // Discard '-'s
    let residues: Vec<&char> = freqs.keys().filter(|&&r| r != '-').collect();
    let sum: f64 = residues
        .into_iter()
        .map(|res| {
            let p = *freqs.get(res).unwrap();
            p * p.ln()
        })
        .sum();

    -sum
}

fn percent_identity(s1: &str, s2: &str) -> f64 {
    let num_identical = s1
        .chars()
        .zip(s2.chars())
        .filter(|(c1, c2)| c1.eq_ignore_ascii_case(c2))
        .count();
    num_identical as f64 / s1.len() as f64
}

fn seq_len_nogaps(s: &str) -> f64 {
    s.chars().filter(|c| c.is_alphabetic()).count() as f64 / s.len() as f64
}

fn seq_type(sequence: &str) -> SeqType {
    let counts = sequence.to_lowercase().chars().counts();
    let counts_u64: HashMap<char, u64> = counts.into_iter().map(|(k, v)| (k, v as u64)).collect();
    let frequencies = to_freq_distrib(&counts_u64);
    let nt_freq: f64 = *frequencies.get(&'a').unwrap_or(&0.0)
        + *frequencies.get(&'c').unwrap_or(&0.0)
        + *frequencies.get(&'g').unwrap_or(&0.0)
        + *frequencies.get(&'t').unwrap_or(&0.0)
        + *frequencies.get(&'u').unwrap_or(&0.0);
    // A quick-and dirty heuristic, I'm afraid
    if nt_freq > 0.75 {
        Nucleic
    } else {
        Protein
    }
}

#[cfg(test)]
mod tests {
    use crate::alignment::{
        best_residue, consensus, densities, entropies, entropy, percent_identity, res_count,
        seq_len_nogaps, seq_type, to_freq_distrib, Alignment, BestResidue, RefSpec, ResidueCounts,
        ResidueDistribution, SeqType,
        SeqType::{Nucleic, Protein},
    };
    use crate::seq::fasta::read_fasta_file;
    use approx::assert_relative_eq;
    use std::collections::HashMap;

    #[test]
    fn test_read_aln() {
        let fasta1 = read_fasta_file("./data/test2.fas").unwrap();
        let aln1 = Alignment::from_file(fasta1);
        assert_eq!("seq1", aln1.headers[0]);
        assert_eq!("seq2", aln1.headers[1]);
        assert_eq!("seq3", aln1.headers[2]);
        assert_eq!("TTGCCG-CGA", aln1.sequences[0]);
        assert_eq!("TTCCCGGCGA", aln1.sequences[1]);
        assert_eq!("TTACCG-CAA", aln1.sequences[2]);
    }

    #[test]
    fn test_consensus() {
        let fasta2 = read_fasta_file("data/test-cons.fas").unwrap();
        let aln2 = Alignment::from_file(fasta2);
        assert_eq!("AQw-n", consensus(&aln2.sequences));
    }

    #[test]
    fn test_res_count() {
        let fasta2 = read_fasta_file("data/test-cons.fas").unwrap();
        let aln2 = Alignment::from_file(fasta2);
        let mut d0: ResidueCounts = HashMap::new();
        d0.insert('A', 6);
        assert_eq!(d0, res_count(&aln2.sequences, 0));

        let mut d1: ResidueCounts = HashMap::new();
        d1.insert('Q', 5);
        d1.insert('T', 1);
        assert_eq!(d1, res_count(&aln2.sequences, 1));

        let mut d2: ResidueCounts = HashMap::new();
        d2.insert('W', 2);
        d2.insert('I', 1);
        d2.insert('S', 1);
        d2.insert('D', 1);
        d2.insert('F', 1);
        assert_eq!(d2, res_count(&aln2.sequences, 2));

        let mut d3: ResidueCounts = HashMap::new();
        d3.insert('-', 3);
        d3.insert('K', 2);
        d3.insert('L', 1);
        assert_eq!(d3, res_count(&aln2.sequences, 3));
    }

    #[test]
    fn test_most_frequent_residue() {
        let d0: ResidueCounts = HashMap::from([('A', 6)]);
        let mut exp: BestResidue = BestResidue {
            residue: 'A',
            frequency: 6,
        };
        assert_eq!(exp, best_residue(&d0));

        let d1: ResidueCounts = HashMap::from([('Q', 5), ('T', 1)]);
        exp = BestResidue {
            residue: 'Q',
            frequency: 5,
        };
        assert_eq!(exp, best_residue(&d1));

        let d2: ResidueCounts = HashMap::from([('W', 2), ('I', 1), ('S', 1), ('D', 1), ('F', 1)]);
        exp = BestResidue {
            residue: 'W',
            frequency: 2,
        };
        assert_eq!(exp, best_residue(&d2));

        // col 3 cannot be tested <- ties

        let d4: ResidueCounts = HashMap::from([('-', 3), ('K', 2), ('L', 1)]);
        exp = BestResidue {
            residue: '-',
            frequency: 3,
        };
        assert_eq!(exp, best_residue(&d4));
    }

    #[test]
    fn test_to_freq_distrib() {
        let eps = 0.001;
        let counts: ResidueCounts = HashMap::from([('K', 3), ('L', 3), ('G', 6), ('-', 6)]);
        let rfreqs = to_freq_distrib(&counts);
        assert_relative_eq!(0.25, *rfreqs.get(&'K').unwrap(), epsilon = eps);
        assert_relative_eq!(0.25, *rfreqs.get(&'L').unwrap(), epsilon = eps);
        assert_relative_eq!(0.5, *rfreqs.get(&'G').unwrap(), epsilon = eps);
    }

    #[test]
    fn test_entropy_1() {
        let eps = 0.00001;
        let distrib: ResidueDistribution = ResidueDistribution::from([('A', 1.0)]);
        assert_relative_eq!(0.0, entropy(&distrib), epsilon = eps);
    }

    #[test]
    fn test_entropy_2() {
        let eps = 0.00001;
        let distrib: ResidueDistribution = ResidueDistribution::from([('A', 0.5), ('F', 0.5)]);
        // This should be ln(2), and as it happens Rust has a constant for this; remarkably, clippy
        // detects the literal constant below and suggests using the (arguably more accurate)
        // built-in definition.
        // assert_relative_eq!(0.6931471805599453, entropy(&distrib), epsilon = eps);
        assert_relative_eq!(std::f64::consts::LN_2, entropy(&distrib), epsilon = eps);
    }

    #[test]
    fn test_entropy_3() {
        let eps = 0.00001;
        let distrib: ResidueDistribution =
            ResidueDistribution::from([('A', 0.5), ('F', 0.25), ('T', 0.25)]);
        assert_relative_eq!(1.0397207708399179, entropy(&distrib), epsilon = eps);
    }

    #[test]
    fn test_entropies() {
        let fasta2 = read_fasta_file("data/test-cons.fas").unwrap();
        let aln2 = Alignment::from_file(fasta2);
        let entrs = entropies(&aln2.sequences);
        let eps = 0.001;
        assert_relative_eq!(0.0, entrs[0], epsilon = eps);
        assert_relative_eq!(0.4505, entrs[1], epsilon = eps);
        assert_relative_eq!(1.5607, entrs[2], epsilon = eps);
        assert_relative_eq!(0.6365, entrs[3], epsilon = eps);
    }

    #[test]
    fn test_density() {
        let fasta = read_fasta_file("data/test-density.msa").unwrap();
        let aln = Alignment::from_file(fasta);
        let dens = densities(&aln.sequences);
        assert_eq!(1.0, dens[0]);
        assert_eq!(0.8, dens[1]);
        assert_eq!(0.6, dens[2]);
        assert_eq!(0.4, dens[3]);
        assert_eq!(0.2, dens[4]);
        assert_eq!(0.0, dens[5]);
    }

    #[test]
    fn test_order_aln() {
        let fasta = read_fasta_file("./data/test4.aln").unwrap();
        let aln1 = Alignment::from_file(fasta);
        // Check original order
        assert_eq!("Zea_001", aln1.headers[0]);
        assert_eq!("Rana_002", aln1.headers[1]);
        assert_eq!("Panthera_050", aln1.headers[49]);
        assert_eq!("tgctgttcgtcaaAgtaggcc", aln1.sequences[0]);
        assert_eq!("tgctgttAgAcaaagtaggcc", aln1.sequences[1]);
        assert_eq!("tgctgttcgtcaaagtaggcc", aln1.sequences[49]);
    }

    #[test]
    fn test_similarity_00() {
        let s1 = "GAATTC";
        assert_eq!(percent_identity(s1, s1), 1.0);
    }

    #[test]
    fn test_similarity_05() {
        let s1 = "GAATTC";
        let s2 = "GAA---";
        assert_eq!(percent_identity(s1, s2), 0.5);
    }

    #[test]
    fn test_similarity_10() {
        let s1 = "GAATTC";
        let s2 = "gaattc";
        assert_eq!(percent_identity(s1, s2), 1.0);
    }

    #[test]
    fn test_seq_len_nogaps_00() {
        assert_eq!(seq_len_nogaps("atgc"), 1.0);
    }

    #[test]
    fn test_seq_len_nogaps_05() {
        assert_eq!(seq_len_nogaps("a-gc"), 0.75);
    }

    #[test]
    fn test_seq_len_nogaps_10() {
        assert_eq!(seq_len_nogaps("--.-"), 0.0);
    }

    #[test]
    fn test_seq_type_00() {
        assert_eq!(Nucleic, seq_type("GAATTC"));
    }

    #[test]
    fn test_seq_type_05() {
        assert_eq!(Protein, seq_type("HGTSDA"));
    }

    #[test]
    fn test_seq_type_10() {
        assert_eq!(Nucleic, seq_type("cgatgcacgatgcncagtgtuucgatcga"));
    }

    #[test]
    fn test_seq_type_15() {
        assert_eq!(Nucleic, seq_type("UUTGAU"));
    }

    // Make sure seq files with unequal lengths get correctly padded
    #[test]
    fn test_unequal_seq_len() {
        let fasta = read_fasta_file("./data/test5.aln").unwrap();
        let _ = Alignment::from_file(fasta);
    }

    // Test the Vec constructor
    #[test]
    fn test_vec_ctor_00() {
        let hdrs = vec![
            String::from("Leo"),
            String::from("Tigris"),
            String::from("Pardus"),
            String::from("Onca"),
        ];
        let seqs = vec![
            String::from("catgcatatg"),
            String::from("aatgcatatg"),
            String::from("tatgcatatg"),
            String::from("gatgcatatg"),
        ];
        let aln = Alignment::from_vecs(hdrs, seqs);
        assert_eq!(4, aln.num_seq());
        assert_eq!(10, aln.aln_len());
        assert_eq!(SeqType::Nucleic, aln.macromolecule_type());
        assert_eq!("Onca", aln.headers[3]);
        assert_eq!("gatgcatatg", aln.sequences[3]);
    }

    // Test the reference specifier
    #[test]
    fn test_reference_specifier() {
        let hdrs = vec![
            String::from("frugilegus"),
            String::from("monedula"),
            String::from("corax"),
            String::from("corone"),
            String::from("cornix"),
        ];
        let seqs = vec![
            String::from("catgcatatg"),
            String::from("aatgcatatg"),
            String::from("tatgcatatg"),
            String::from("tatgcatatg"),
            String::from("gatgcatatg"),
        ];
        let mut aln = Alignment::from_vecs(hdrs, seqs);
        // By default, the reference sequence is the consensus
        assert_eq!(RefSpec::Consensus, aln.get_ref_spec());
        assert_eq!("tATGCATATG", aln.reference());
        // Now set the ref to the first sequence (rank 0)
        let _ = aln.set_ref_spec(RefSpec::Rank(0));
        assert_eq!(RefSpec::Rank(0), aln.get_ref_spec());
        assert_eq!("catgcatatg", aln.reference());
        // Back to consensus
        let _ = aln.set_ref_spec(RefSpec::Consensus);
        assert_eq!(RefSpec::Consensus, aln.get_ref_spec());
        assert_eq!("tATGCATATG", aln.reference());
    }

    // Tests the %id WRT ref (incl. when != consensus)
    #[test]
    fn test_pct_id_wrt_ref() {
        let hdrs = vec![
            String::from("frugilegus"),
            String::from("monedula"),
            String::from("corax"),
            String::from("corone"),
            String::from("cornix"),
        ];
        // consensus: ACg-
        let seqs = vec![
            String::from("A---"),
            String::from("AC--"),
            String::from("ACG-"),
            String::from("ACGT"),
            String::from("ACGT"),
        ];
        let mut aln = Alignment::from_vecs(hdrs, seqs);
        // Check the ref, which by default is the consensus
        assert_eq!("ACg-", aln.reference());
        assert_eq!(vec![0.5, 0.75, 1.0, 0.75, 0.75], aln.id_wrt_reference);
        // Now switch to seq #0 for reference
        let _ = aln.set_ref_spec(RefSpec::Rank(0));
        assert_eq!("A---", aln.reference());
        assert_eq!(vec![1.0, 0.75, 0.5, 0.25, 0.25], aln.id_wrt_reference);
        // Switch back to consensus
        let _ = aln.set_ref_spec(RefSpec::Consensus);
        assert_eq!("ACg-", aln.reference());
        assert_eq!(vec![0.5, 0.75, 1.0, 0.75, 0.75], aln.id_wrt_reference);
    }
}