// The explicit `core::simd` distance kernel needs the nightly `portable_simd`
// feature. Gated on the `simd` Cargo feature so default (stable) builds are
// unaffected — the kernel falls back to the autovectorizable scalar path.
#![cfg_attr(feature = "simd", feature(portable_simd))]
//! Neighbor-Joining phylogenetic tree inference library.
//!
//! # Data flow
//!
//! ```text
//! [FASTA / Python dict / JS object]
//!         │
//!         ▼
//!      NJConfig  (config.rs)
//!         │
//!         ▼
//!   detect_alphabet()  ──►  Alphabet::DNA | Alphabet::Protein
//!         │
//!         ▼
//!     MSA<DNA|Protein>  (msa.rs)
//!      ├── bootstrap() ──► bootstrap_clade_counts()
//!      └── into_dist::<Model>()
//!               │
//!               ▼
//!           DistMat  (dist.rs)
//!               │
//!               ▼
//!         neighbor_joining()  ──►  NJState::run()  (nj.rs)
//!               │
//!               ▼
//!           TreeNode  (tree.rs)
//!               │
//!               ▼
//!           to_newick()  ──►  Newick String
//! ```
//!
//! # Public API
//!
//! The single public entry point is [`nj`], which accepts an [`NJConfig`] and
//! returns a Newick string. Everything else is internal implementation detail
//! exposed only to the Python and WASM wrapper crates.
//!
//! # Model–alphabet compatibility
//!
//! | Model | DNA | Protein |
//! |-------|-----|---------|
//! | `PDiff` | ✓ | ✓ |
//! | `JukesCantor` | ✓ | — |
//! | `Kimura2P` | ✓ | — |
//! | `TajimaNei` | ✓ | — |
//! | `Tamura` | ✓ | — |
//! | `Poisson` | — | ✓ |
//! | `KimuraProtein` | — | ✓ |
//!
//! Providing an incompatible model returns an `Err` from [`nj`].
pub mod alphabet;
pub mod config;
pub mod distance_matrix;
pub mod error;
pub mod event;
pub mod fasta;
pub mod models;
pub mod msa;
pub mod nj;
pub mod tree;

use bitvec::prelude::{BitVec, Lsb0, bitvec};
use std::collections::{HashMap, HashSet};

use crate::alphabet::{Alphabet, AlphabetEncoding, DNA, Protein};
use crate::config::SubstitutionModel;
pub use crate::config::{DistConfig, MSA, NJConfig, NJResult, SequenceObject};
use crate::distance_matrix::DistMat;
pub use crate::distance_matrix::DistanceResult;
pub use crate::error::NJError;
pub use crate::event::{LogLevel, NJEvent};
pub use crate::fasta::parse_fasta;
use crate::models::{
    JukesCantor, Kimura2P, KimuraProtein, ModelCalculation, PDiff, Poisson, TajimaNei, Tamura,
};
use crate::tree::{NameOrSupport, TreeNode};

/// Fills `out` with the leaf indices of all taxa in the subtree rooted at `node`.
///
/// Bits in `out` are set to `true` for each leaf encountered. The bit position
/// is looked up in `idx` by the leaf's name label. Returns `Err` if a leaf
/// has no name label (should not occur for well-formed NJ trees).
fn bitset_of(
    node: &TreeNode,
    idx: &HashMap<String, usize>,
    out: &mut BitVec<u8, Lsb0>,
) -> Result<(), String> {
    match &node.children {
        None => match &node.label {
            Some(NameOrSupport::Name(name)) => {
                let i = *idx
                    .get(name)
                    .ok_or_else(|| format!("leaf '{name}' not found in index map"))?;
                out.set(i, true);
                Ok(())
            }
            _ => Err("Leaf node without a name label".into()),
        },
        Some([l, r]) => {
            bitset_of(l, idx, out)?;
            bitset_of(r, idx, out)?;
            Ok(())
        }
    }
}

/// Recursively counts how many times each non-trivial clade appears in `tree`.
///
/// A clade is represented as a raw-byte encoding of a `BitVec` over the `n_taxa`
/// leaf indices. Only clades with `1 < size < n_taxa` (i.e. proper internal
/// clades) are counted. Each call increments the clade's entry in `counter` by 1.
/// Used by [`bootstrap_clade_counts`] to aggregate over bootstrap replicates.
fn count_clades(
    tree: &TreeNode,
    idx: &HashMap<String, usize>,
    n_taxa: usize,
    counter: &mut HashMap<Vec<u8>, usize>,
) -> Result<(), String> {
    if let Some([l, r]) = &tree.children {
        let mut bv = bitvec![u8, Lsb0; 0; n_taxa];
        bitset_of(tree, idx, &mut bv)?;

        let n = bv.count_ones();
        if n > 1 && n < n_taxa {
            counter
                .entry(bv.as_raw_slice().to_vec())
                .and_modify(|c| *c += 1)
                .or_insert(1);
        }

        count_clades(l, idx, n_taxa, counter)?;
        count_clades(r, idx, n_taxa, counter)?;
    }
    Ok(())
}

/// Builds a Rayon thread pool with `num_threads` workers.
///
/// When `num_threads` is `None`, Rayon uses its default (one thread per logical CPU).
/// Returns `Err` if the pool cannot be constructed.
#[cfg(feature = "parallel")]
pub(crate) fn build_thread_pool(num_threads: Option<usize>) -> Result<rayon::ThreadPool, String> {
    let mut builder = rayon::ThreadPoolBuilder::new();
    if let Some(n) = num_threads {
        builder = builder.num_threads(n);
    }
    builder.build().map_err(|e| e.to_string())
}

/// Runs a single bootstrap replicate end-to-end: resample alignment columns,
/// compute the distance matrix with model `M`, run Neighbor-Joining, and tally
/// the replicate's clades into a fresh counter.
///
/// Shared by the sequential and parallel bootstrap paths. Every fallible step is
/// propagated via `?` so a failed replicate becomes a returned `Err` rather than
/// a panic (which would otherwise poison a Rayon worker thread).
fn run_one_replicate<A, M>(
    msa: &MSA<A>,
    idx_map: &HashMap<String, usize>,
    n_taxa: usize,
) -> Result<HashMap<Vec<u8>, usize>, String>
where
    A: AlphabetEncoding + Send + Sync,
    A::Symbol: Send + Sync,
    M: ModelCalculation<A> + Send + Sync,
{
    let tree = msa.bootstrap()?.into_dist::<M>().neighbor_joining()?;
    let mut local = HashMap::new();
    count_clades(&tree, idx_map, n_taxa, &mut local)?;
    Ok(local)
}

/// Parallel bootstrap worker: runs all replicates on Rayon threads and sends
/// per-replicate clade maps over an MPSC channel. The main thread (caller)
/// merges results and fires `on_event` — keeping the callback on a single
/// thread with no `Sync` requirement.
///
/// Uses `std::thread::scope` so `msa` and `idx_map` can be borrowed into the
/// spawned thread without requiring `'static` lifetimes.
#[cfg(feature = "parallel")]
fn bootstrap_clade_counts_parallel<A, M>(
    msa: &MSA<A>,
    n_bootstrap_samples: usize,
    idx_map: &HashMap<String, usize>,
    n_taxa: usize,
    on_event: Option<&dyn Fn(NJEvent)>,
    num_threads: Option<usize>,
) -> Result<HashMap<Vec<u8>, usize>, String>
where
    A: AlphabetEncoding + Send + Sync,
    A::Symbol: Send + Sync,
    M: ModelCalculation<A> + Send + Sync,
{
    use rayon::iter::{IntoParallelIterator, ParallelIterator};
    use std::sync::mpsc;

    let pool = build_thread_pool(num_threads)?;
    let (tx, rx) = mpsc::channel::<Result<HashMap<Vec<u8>, usize>, String>>();
    let mut counter: HashMap<Vec<u8>, usize> = HashMap::new();

    std::thread::scope(|scope| -> Result<(), String> {
        // Spawn the Rayon work onto a background OS thread so the main thread
        // can consume from `rx` concurrently, giving real per-replicate progress.
        // `pool.install` installs the thread pool as the active pool for the
        // duration of the closure, bounding parallelism to `num_threads` workers.
        scope.spawn(|| {
            pool.install(|| {
                (0..n_bootstrap_samples)
                    .into_par_iter()
                    .for_each_with(tx, |sender, _| {
                        let result = run_one_replicate::<A, M>(msa, idx_map, n_taxa);
                        // Ignore send errors: only occurs if receiver was dropped,
                        // which cannot happen while we are in the recv loop below.
                        let _ = sender.send(result);
                    });
            });
        });

        // Main thread: receive results as they arrive, merge, and fire progress.
        for completed in 1..=n_bootstrap_samples {
            match rx.recv() {
                Ok(Ok(local)) => {
                    for (clade, count) in local {
                        *counter.entry(clade).or_insert(0) += count;
                    }
                }
                Ok(Err(e)) => return Err(e),
                Err(_) => return Err("bootstrap channel closed unexpectedly".into()),
            }
            if let Some(cb) = on_event {
                cb(NJEvent::BootstrapProgress {
                    completed,
                    total: n_bootstrap_samples,
                });
            }
        }
        Ok(())
    })?;

    Ok(counter)
}

/// Performs bootstrap sampling and counts clades across bootstrap trees.
///
/// When the `parallel` feature is enabled, replicates run on the Rayon thread
/// pool via [`bootstrap_clade_counts_parallel`]; the `on_event` callback is
/// still called after each replicate from the main thread. Without the feature
/// the loop is sequential and `on_event` fires in order.
fn bootstrap_clade_counts<A, M>(
    msa: &MSA<A>,
    n_bootstrap_samples: usize,
    on_event: Option<&dyn Fn(NJEvent)>,
    num_threads: Option<usize>,
) -> Result<Option<HashMap<Vec<u8>, usize>>, String>
where
    A: AlphabetEncoding + Send + Sync,
    A::Symbol: Send + Sync,
    M: ModelCalculation<A> + Send + Sync,
{
    if n_bootstrap_samples == 0 {
        return Ok(None);
    }
    if let Some(cb) = on_event {
        cb(NJEvent::BootstrapStarted {
            total: n_bootstrap_samples,
        });
    }
    let idx_map: HashMap<String, usize> = msa.to_index_map();
    let n_taxa = msa.n_sequences;

    #[cfg(feature = "parallel")]
    let counter = bootstrap_clade_counts_parallel::<A, M>(
        msa,
        n_bootstrap_samples,
        &idx_map,
        n_taxa,
        on_event,
        num_threads,
    )?;

    #[cfg(not(feature = "parallel"))]
    let counter = {
        let _ = num_threads;
        let mut c = HashMap::new();
        for i in 0..n_bootstrap_samples {
            let local = run_one_replicate::<A, M>(msa, &idx_map, n_taxa)?;
            for (clade, count) in local {
                *c.entry(clade).or_insert(0) += count;
            }
            if let Some(cb) = on_event {
                cb(NJEvent::BootstrapProgress {
                    completed: i + 1,
                    total: n_bootstrap_samples,
                });
            }
        }
        c
    };

    Ok(Some(counter))
}

/// Annotates internal nodes with bootstrap support values from `counts`.
///
/// For each internal node, computes its clade `BitVec`, looks up the count in
/// `counts`, normalises it to a percentage (`count * 100 / n_bootstrap_samples`),
/// and assigns a [`NameOrSupport::Support`] label if a matching entry is found.
/// Nodes whose clade was never observed in bootstrap replicates receive no label.
fn add_bootstrap_to_tree(
    node: &mut TreeNode,
    idx: &HashMap<String, usize>,
    n_taxa: usize,
    counts: &HashMap<Vec<u8>, usize>,
    n_bootstrap_samples: usize,
) -> Result<(), String> {
    if node.children.is_some() {
        let mut bv = bitvec![u8, Lsb0; 0; n_taxa];
        bitset_of(node, idx, &mut bv)?;

        let n = bv.count_ones();
        if n > 1
            && n < n_taxa
            && let Some(c) = counts.get(&bv.as_raw_slice().to_vec())
        {
            let pct = c * 100 / n_bootstrap_samples;
            node.label = Some(NameOrSupport::Support(pct));
        }

        if let Some([l, r]) = &mut node.children {
            add_bootstrap_to_tree(l, idx, n_taxa, counts, n_bootstrap_samples)?;
            add_bootstrap_to_tree(r, idx, n_taxa, counts, n_bootstrap_samples)?;
        }
    }
    Ok(())
}

/// Validates that the MSA is non-empty, all sequences have equal length, and
/// all identifiers are unique.
fn validate_msa(msa: &[SequenceObject]) -> Result<(), NJError> {
    if msa.is_empty() {
        return Err(NJError::EmptyMsa);
    }
    let expected_len = msa[0].sequence.len();
    if expected_len == 0 {
        return Err(NJError::EmptySequence);
    }
    let mut seen: HashSet<&str> = HashSet::with_capacity(msa.len());
    for s in msa {
        if s.sequence.len() != expected_len {
            return Err(NJError::SequenceLengthMismatch {
                expected: expected_len,
                got: s.sequence.len(),
                identifier: s.identifier.clone(),
            });
        }
        if !seen.insert(s.identifier.as_str()) {
            return Err(NJError::DuplicateIdentifier {
                identifier: s.identifier.clone(),
            });
        }
    }
    Ok(())
}

/// Heuristically detects whether the MSA contains DNA or protein sequences.
///
/// Returns [`Alphabet::DNA`] unless any sequence contains a byte that is not
/// in the DNA character set (case-insensitive), in which case
/// [`Alphabet::Protein`] is returned. The DNA set includes standard bases
/// `{A, C, G, T}`, uridine `U` (RNA), `N` (unknown), gap `-`, and all 11
/// IUPAC ambiguity codes `{R, Y, S, W, K, M, B, D, H, V}`.
fn detect_alphabet(msa: &[SequenceObject]) -> Alphabet {
    let mut is_protein = false;

    'outer: for seq in msa {
        for c in seq.sequence.bytes() {
            match c.to_ascii_uppercase() {
                b'A' | b'C' | b'G' | b'T' | b'U' | b'N' | b'-' | b'R' | b'Y' | b'S' | b'W'
                | b'K' | b'M' | b'B' | b'D' | b'H' | b'V' => { /* still possible DNA */ }
                _ => {
                    is_protein = true;
                    break 'outer;
                }
            }
        }
    }

    if is_protein {
        Alphabet::Protein
    } else {
        Alphabet::DNA
    }
}

/// Runs distance matrix computation with model `M` on alphabet `A`.
fn run_distance_matrix<A, M>(
    msa: MSA<A>,
    num_threads: Option<usize>,
) -> Result<DistanceResult, String>
where
    A: AlphabetEncoding + Send + Sync,
    A::Symbol: Send + Sync,
    M: ModelCalculation<A> + Send + Sync,
{
    #[cfg(feature = "parallel")]
    {
        let pool = build_thread_pool(num_threads)?;
        Ok(pool.install(|| msa.into_dist::<M>()).into_result())
    }
    #[cfg(not(feature = "parallel"))]
    {
        let _ = num_threads;
        Ok(msa.into_dist::<M>().into_result())
    }
}

/// Runs average distance computation with model `M` on alphabet `A`.
fn run_average_distance<A, M>(msa: MSA<A>, num_threads: Option<usize>) -> Result<f64, String>
where
    A: AlphabetEncoding + Send + Sync,
    A::Symbol: Send + Sync,
    M: ModelCalculation<A> + Send + Sync,
{
    #[cfg(feature = "parallel")]
    {
        let pool = build_thread_pool(num_threads)?;
        Ok(pool.install(|| msa.into_dist::<M>()).average())
    }
    #[cfg(not(feature = "parallel"))]
    {
        let _ = num_threads;
        Ok(msa.into_dist::<M>().average())
    }
}

/// Runs NJ with model `M` on alphabet `A` and returns an [`NJResult`].
///
/// If `n_bootstrap_samples > 0`, generates that many bootstrap replicates,
/// collects clade counts via [`bootstrap_clade_counts`], runs NJ on the
/// original distances, and annotates the tree before serialising to Newick.
/// When `include_distance_matrix` or `include_average_distance` is `true`,
/// the corresponding values are captured from the distance matrix before NJ
/// consumes it and included in the returned [`NJResult`].
fn run_nj<A, M>(
    msa: MSA<A>,
    n_bootstrap_samples: usize,
    on_event: Option<&dyn Fn(NJEvent)>,
    num_threads: Option<usize>,
    include_distance_matrix: bool,
    include_average_distance: bool,
) -> Result<NJResult, String>
where
    A: AlphabetEncoding + Send + Sync,
    A::Symbol: Send + Sync,
    M: ModelCalculation<A> + Send + Sync,
{
    let clade_counts =
        bootstrap_clade_counts::<A, M>(&msa, n_bootstrap_samples, on_event, num_threads)?;

    if let Some(cb) = on_event {
        cb(NJEvent::ComputingDistances);
    }

    #[cfg(feature = "parallel")]
    let dist = {
        let pool = build_thread_pool(num_threads)?;
        pool.install(|| msa.into_dist::<M>())
    };
    #[cfg(not(feature = "parallel"))]
    let dist = msa.into_dist::<M>();

    let distance_matrix = if include_distance_matrix {
        Some(dist.to_result())
    } else {
        None
    };
    let average_distance = if include_average_distance {
        Some(dist.average())
    } else {
        None
    };

    if let Some(cb) = on_event {
        cb(NJEvent::RunningNJ);
    }

    let mut main_tree = dist.neighbor_joining()?;

    let newick = match clade_counts {
        Some(counts) => {
            if let Some(cb) = on_event {
                cb(NJEvent::AnnotatingBootstrap);
            }
            let main_idx_map: HashMap<String, usize> = msa.to_index_map();
            add_bootstrap_to_tree(
                &mut main_tree,
                &main_idx_map,
                msa.n_sequences,
                &counts,
                n_bootstrap_samples,
            )?;
            main_tree.to_newick()
        }
        None => main_tree.to_newick(),
    };
    Ok(NJResult {
        newick,
        distance_matrix,
        average_distance,
    })
}

/// Emits a `Warning`-level [`NJEvent::Log`] for any sequence consisting entirely
/// of gap characters. Such sequences have no comparable sites, so every pairwise
/// distance involving them is `0.0` ([`crate::models`]); surfacing a warning makes
/// that otherwise-silent behaviour observable. No-op when `on_event` is `None`.
fn warn_degenerate_sequences(msa: &[SequenceObject], on_event: Option<&dyn Fn(NJEvent)>) {
    let Some(cb) = on_event else { return };
    for s in msa {
        if !s.sequence.is_empty() && s.sequence.bytes().all(|b| b == b'-') {
            cb(NJEvent::Log {
                level: LogLevel::Warning,
                message: format!(
                    "Sequence '{}' contains only gaps; its pairwise distances will be 0.0",
                    s.identifier
                ),
            });
        }
    }
}

/// Dispatches to a generic runner based on the runtime alphabet and substitution
/// model, enforcing the model–alphabet compatibility table in a single place.
///
/// `$pairs` is consumed to build the typed [`MSA`]; `$run` is the generic runner
/// (`run_nj`, `run_distance_matrix`, `run_average_distance`) and its trailing
/// arguments follow the `msa` it receives. Each runner returns `Result<_, String>`,
/// mapped here to [`NJError::AlgorithmFailure`]; unsupported (alphabet, model)
/// pairs yield [`NJError::IncompatibleModel`].
macro_rules! dispatch_run {
    ($alphabet:expr, $model:expr, $pairs:expr, $run:ident ( $($arg:expr),* $(,)? )) => {{
        let alphabet = $alphabet;
        let model = $model;
        match alphabet {
            Alphabet::DNA => {
                let msa = MSA::<DNA>::try_from_pairs($pairs).map_err(NJError::AlgorithmFailure)?;
                match model {
                    SubstitutionModel::PDiff => {
                        $run::<DNA, PDiff>(msa, $($arg),*).map_err(NJError::AlgorithmFailure)
                    }
                    SubstitutionModel::JukesCantor => {
                        $run::<DNA, JukesCantor>(msa, $($arg),*).map_err(NJError::AlgorithmFailure)
                    }
                    SubstitutionModel::Kimura2P => {
                        $run::<DNA, Kimura2P>(msa, $($arg),*).map_err(NJError::AlgorithmFailure)
                    }
                    SubstitutionModel::TajimaNei => {
                        $run::<DNA, TajimaNei>(msa, $($arg),*).map_err(NJError::AlgorithmFailure)
                    }
                    SubstitutionModel::Tamura => {
                        $run::<DNA, Tamura>(msa, $($arg),*).map_err(NJError::AlgorithmFailure)
                    }
                    SubstitutionModel::Poisson | SubstitutionModel::KimuraProtein => {
                        Err(NJError::IncompatibleModel {
                            model,
                            alphabet: Alphabet::DNA,
                        })
                    }
                }
            }
            Alphabet::Protein => {
                let msa =
                    MSA::<Protein>::try_from_pairs($pairs).map_err(NJError::AlgorithmFailure)?;
                match model {
                    SubstitutionModel::Poisson => {
                        $run::<Protein, Poisson>(msa, $($arg),*).map_err(NJError::AlgorithmFailure)
                    }
                    SubstitutionModel::KimuraProtein => {
                        $run::<Protein, KimuraProtein>(msa, $($arg),*)
                            .map_err(NJError::AlgorithmFailure)
                    }
                    SubstitutionModel::PDiff => {
                        $run::<Protein, PDiff>(msa, $($arg),*).map_err(NJError::AlgorithmFailure)
                    }
                    SubstitutionModel::JukesCantor
                    | SubstitutionModel::Kimura2P
                    | SubstitutionModel::TajimaNei
                    | SubstitutionModel::Tamura => Err(NJError::IncompatibleModel {
                        model,
                        alphabet: Alphabet::Protein,
                    }),
                }
            }
        }
    }};
}

/// Infers a phylogenetic tree from an aligned MSA and returns an [`NJResult`].
///
/// This is the single public entry point for the library. The alphabet is
/// auto-detected from the sequences unless `conf.alphabet` is set; `conf.substitution_model`
/// must be compatible with the alphabet (see the module-level compatibility
/// table). Returns `Err` for an empty MSA, duplicate identifiers, an
/// incompatible model, or any internal NJ failure.
///
/// Set [`NJConfig::return_distance_matrix`] and/or
/// [`NJConfig::return_average_distance`] to `true` to include those values in
/// the returned [`NJResult`] alongside the Newick tree.
///
/// **Degenerate inputs.** A single-sequence MSA produces a lone-leaf "tree".
/// Sequences consisting entirely of gaps have no comparable sites, so their
/// pairwise distances are `0.0` (see [`crate::models`]); a `Warning`-level
/// [`NJEvent::Log`] is emitted for each such sequence so the behaviour is not silent.
///
/// `on_event` is called with an [`NJEvent`] at each stage of the algorithm.
/// Bootstrap progress is reported via [`NJEvent::BootstrapProgress`] after
/// each replicate. Pass `None` if event reporting is not needed.
pub fn nj(conf: NJConfig, on_event: Option<Box<dyn Fn(NJEvent)>>) -> Result<NJResult, NJError> {
    let cb = on_event.as_deref();
    let num_threads = conf.num_threads;
    let include_distance_matrix = conf.return_distance_matrix;
    let include_average_distance = conf.return_average_distance;
    let n_bootstrap_samples = conf.n_bootstrap_samples;
    validate_msa(&conf.msa)?;
    let n_sites = conf.msa[0].sequence.len();
    if let Some(cb) = cb {
        cb(NJEvent::MsaValidated {
            n_sequences: conf.msa.len(),
            n_sites,
        });
    }
    warn_degenerate_sequences(&conf.msa, cb);
    let alphabet = conf.alphabet.unwrap_or_else(|| detect_alphabet(&conf.msa));
    if let Some(cb) = cb {
        cb(NJEvent::AlphabetDetected {
            alphabet: alphabet.clone(),
        });
    }
    let model = conf.substitution_model;
    let pairs = conf.msa.into_iter().map(|s| (s.identifier, s.sequence));
    dispatch_run!(
        alphabet,
        model,
        pairs,
        run_nj(
            n_bootstrap_samples,
            cb,
            num_threads,
            include_distance_matrix,
            include_average_distance,
        )
    )
}

/// Computes pairwise distances from an aligned MSA and returns a [`DistanceResult`].
///
/// The alphabet is auto-detected from the sequences unless `conf.alphabet` is set;
/// `conf.substitution_model` must be compatible with the alphabet (see the
/// module-level compatibility table). Returns `Err` for an empty MSA, incompatible
/// model, or mismatched sequence lengths. Does not run Neighbor-Joining or bootstrapping.
pub fn distance_matrix(conf: DistConfig) -> Result<DistanceResult, NJError> {
    let num_threads = conf.num_threads;
    validate_msa(&conf.msa)?;
    let alphabet = conf.alphabet.unwrap_or_else(|| detect_alphabet(&conf.msa));
    let model = conf.substitution_model;
    let pairs = conf.msa.into_iter().map(|s| (s.identifier, s.sequence));
    dispatch_run!(alphabet, model, pairs, run_distance_matrix(num_threads))
}

/// Computes the mean of all `n*(n-1)/2` unique pairwise distances.
///
/// Same alphabet auto-detection and model–alphabet compatibility as [`nj`].
/// Returns `0.0` for fewer than 2 taxa. Returns `Err` for an empty MSA,
/// incompatible model, or mismatched sequence lengths.
pub fn average_distance(conf: DistConfig) -> Result<f64, NJError> {
    let num_threads = conf.num_threads;
    validate_msa(&conf.msa)?;
    let alphabet = conf.alphabet.unwrap_or_else(|| detect_alphabet(&conf.msa));
    let model = conf.substitution_model;
    let pairs = conf.msa.into_iter().map(|s| (s.identifier, s.sequence));
    dispatch_run!(alphabet, model, pairs, run_average_distance(num_threads))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::DistConfig;
    use crate::models::SubstitutionModel;

    /// Builds a minimal [`NJConfig`] with both include flags set to the given values.
    fn nj_conf(pairs: &[(&str, &str)], include_dm: bool, include_avg: bool) -> NJConfig {
        NJConfig {
            msa: pairs
                .iter()
                .map(|(id, seq)| SequenceObject {
                    identifier: id.to_string(),
                    sequence: seq.to_string(),
                })
                .collect(),
            n_bootstrap_samples: 0,
            substitution_model: SubstitutionModel::PDiff,
            alphabet: None,
            num_threads: None,
            return_distance_matrix: include_dm,
            return_average_distance: include_avg,
        }
    }

    #[test]
    fn test_nj_wrapper_simple_tree() {
        // ACGTCG vs ACG-GC: pos 3 is gapped (excluded), 2 diffs out of 5 comparable
        // → distance = 0.4; two taxa → each branch = 0.2.
        let sequences = vec![
            SequenceObject {
                identifier: "A".into(),
                sequence: "ACGTCG".into(),
            },
            SequenceObject {
                identifier: "B".into(),
                sequence: "ACG-GC".into(),
            },
        ];
        let conf = NJConfig {
            msa: sequences,
            n_bootstrap_samples: 0,
            substitution_model: SubstitutionModel::PDiff,
            alphabet: None,
            num_threads: None,
            return_distance_matrix: false,
            return_average_distance: false,
        };
        let result = nj(conf, None).expect("NJ failed");
        assert_eq!(result.newick, "(A:0.200,B:0.200);");
    }

    #[test]
    fn test_nj_wrapper_adds_semicolon() {
        let sequences = vec![
            SequenceObject {
                identifier: "Seq0".into(),
                sequence: "A".into(),
            },
            SequenceObject {
                identifier: "Seq1".into(),
                sequence: "A".into(),
            },
        ];
        let conf = NJConfig {
            msa: sequences,
            n_bootstrap_samples: 0,
            substitution_model: SubstitutionModel::PDiff,
            alphabet: None,
            num_threads: None,
            return_distance_matrix: false,
            return_average_distance: false,
        };
        let out = nj(conf, None).unwrap();
        assert!(out.newick.ends_with(';'));
    }

    #[test]
    fn test_nj_deterministic_order() {
        let sequences = vec![
            SequenceObject {
                identifier: "Seq0".into(),
                sequence: "ACGTCG".into(),
            },
            SequenceObject {
                identifier: "Seq1".into(),
                sequence: "ACG-GC".into(),
            },
            SequenceObject {
                identifier: "Seq2".into(),
                sequence: "ACGCGT".into(),
            },
        ];
        let conf = NJConfig {
            msa: sequences,
            n_bootstrap_samples: 0,
            substitution_model: SubstitutionModel::PDiff,
            alphabet: None,
            num_threads: None,
            return_distance_matrix: false,
            return_average_distance: false,
        };

        let t1 = nj(conf.clone(), None).unwrap();
        let t2 = nj(conf, None).unwrap();
        assert_eq!(t1, t2);
    }

    #[test]
    fn test_nj_wrapper_empty_msa() {
        let conf = NJConfig {
            msa: vec![],
            n_bootstrap_samples: 0,
            substitution_model: SubstitutionModel::PDiff,
            alphabet: None,
            num_threads: None,
            return_distance_matrix: false,
            return_average_distance: false,
        };
        let result = nj(conf, None);
        assert!(result.is_err());
    }

    #[test]
    fn test_nj_wrapper_incorrect_model_for_alphabet() {
        let sequences = vec![
            SequenceObject {
                identifier: "Seq0".into(),
                sequence: "ACGTCG".into(),
            },
            SequenceObject {
                identifier: "Seq1".into(),
                sequence: "ACG-GC".into(),
            },
        ];
        let conf = NJConfig {
            msa: sequences,
            n_bootstrap_samples: 0,
            substitution_model: SubstitutionModel::Poisson, // protein model for DNA MSA
            alphabet: None,
            num_threads: None,
            return_distance_matrix: false,
            return_average_distance: false,
        };
        let result = nj(conf, None);
        assert!(result.is_err());
    }

    #[test]
    fn test_nj_wrapper_incorrect_model_for_protein() {
        let sequences = vec![
            SequenceObject {
                identifier: "Seq0".into(),
                sequence: "ACDEFGH".into(),
            },
            SequenceObject {
                identifier: "Seq1".into(),
                sequence: "ACD-FGH".into(),
            },
        ];
        let conf = NJConfig {
            msa: sequences,
            n_bootstrap_samples: 0,
            substitution_model: SubstitutionModel::JukesCantor, // DNA model for protein MSA
            alphabet: None,
            num_threads: None,
            return_distance_matrix: false,
            return_average_distance: false,
        };
        let result = nj(conf, None);
        assert!(result.is_err());
    }

    #[test]
    fn test_nj_rejects_duplicate_identifiers() {
        let result = nj(nj_conf(&[("A", "ACGT"), ("A", "ACGA")], false, false), None);
        assert!(matches!(result, Err(NJError::DuplicateIdentifier { .. })));
    }

    #[test]
    fn test_distance_matrix_rejects_duplicate_identifiers() {
        let conf = dist_conf(&[("A", "ACGT"), ("A", "ACGA")], SubstitutionModel::PDiff);
        assert!(matches!(
            distance_matrix(conf),
            Err(NJError::DuplicateIdentifier { .. })
        ));
    }

    #[test]
    fn test_lowercase_protein_nj_is_non_degenerate() {
        // Lowercase protein residues must encode like uppercase, so the two
        // distinct sequences yield a positive branch length rather than 0.
        let conf = NJConfig {
            msa: vec![
                SequenceObject {
                    identifier: "P0".into(),
                    sequence: "acdefghikl".into(),
                },
                SequenceObject {
                    identifier: "P1".into(),
                    sequence: "acdefgwxyl".into(),
                },
            ],
            n_bootstrap_samples: 0,
            substitution_model: SubstitutionModel::PDiff,
            alphabet: Some(Alphabet::Protein),
            num_threads: None,
            return_distance_matrix: true,
            return_average_distance: false,
        };
        let result = nj(conf, None).expect("NJ on lowercase protein should succeed");
        let dm = result.distance_matrix.unwrap();
        assert!(dm.matrix[0][1] > 0.0, "expected non-zero distance, got {dm:?}");
    }

    #[test]
    fn test_all_gap_sequence_emits_warning() {
        use std::sync::{Arc, Mutex};
        let warnings: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
        let sink = Arc::clone(&warnings);
        let cb: Box<dyn Fn(NJEvent)> = Box::new(move |event| {
            if let NJEvent::Log {
                level: LogLevel::Warning,
                message,
            } = event
            {
                sink.lock().unwrap().push(message);
            }
        });
        let conf = nj_conf(&[("A", "----"), ("B", "ACGA")], false, false);
        let _ = nj(conf, Some(cb)).unwrap();
        let collected = warnings.lock().unwrap();
        assert_eq!(collected.len(), 1, "expected exactly one all-gap warning");
        assert!(collected[0].contains("only gaps"), "got {:?}", collected[0]);
    }

    // --- NJResult optional fields ---

    #[test]
    fn test_nj_result_no_extras_by_default() {
        let result = nj(nj_conf(&[("A", "ACGT"), ("B", "ACGA")], false, false), None).unwrap();
        assert!(result.distance_matrix.is_none());
        assert!(result.average_distance.is_none());
    }

    #[test]
    fn test_nj_result_include_distance_matrix() {
        let result = nj(nj_conf(&[("A", "ACGT"), ("B", "ACGA")], true, false), None).unwrap();
        let dm = result
            .distance_matrix
            .expect("distance_matrix should be present");
        assert_eq!(dm.names, vec!["A", "B"]);
        assert_eq!(dm.matrix.len(), 2);
        assert_eq!(dm.matrix[0][0], 0.0);
        assert!((dm.matrix[0][1] - 0.25).abs() < 1e-12);
        assert!(result.average_distance.is_none());
    }

    #[test]
    fn test_nj_result_include_average_distance() {
        let result = nj(nj_conf(&[("A", "ACGT"), ("B", "ACGA")], false, true), None).unwrap();
        let avg = result
            .average_distance
            .expect("average_distance should be present");
        assert!((avg - 0.25).abs() < 1e-12);
        assert!(result.distance_matrix.is_none());
    }

    #[test]
    fn test_nj_result_include_both() {
        let result = nj(nj_conf(&[("A", "ACGT"), ("B", "ACGA")], true, true), None).unwrap();
        let dm = result
            .distance_matrix
            .as_ref()
            .expect("distance_matrix should be present");
        let avg = result
            .average_distance
            .expect("average_distance should be present");
        // distance matrix and average_distance must be consistent
        assert!((dm.matrix[0][1] - avg).abs() < 1e-12);
    }

    #[test]
    fn test_nj_result_distance_matrix_consistent_with_newick() {
        // Known tree: two taxa with 1/4 difference → branch lengths 0.125 each from NJ
        let result = nj(nj_conf(&[("A", "ACGT"), ("B", "ACGA")], true, false), None).unwrap();
        let dm = result.distance_matrix.unwrap();
        // full matrix should be symmetric with zero diagonal
        assert_eq!(dm.matrix[0][0], 0.0);
        assert_eq!(dm.matrix[1][1], 0.0);
        assert!((dm.matrix[0][1] - dm.matrix[1][0]).abs() < 1e-12);
    }

    // --- distance_matrix ---

    fn dist_conf(pairs: &[(&str, &str)], model: SubstitutionModel) -> DistConfig {
        DistConfig {
            msa: pairs
                .iter()
                .map(|(id, seq)| SequenceObject {
                    identifier: id.to_string(),
                    sequence: seq.to_string(),
                })
                .collect(),
            substitution_model: model,
            alphabet: None,
            num_threads: None,
        }
    }

    #[test]
    fn test_distance_matrix_names_and_shape() {
        let conf = dist_conf(&[("A", "ACGT"), ("B", "ACGA")], SubstitutionModel::PDiff);
        let result = distance_matrix(conf).unwrap();
        assert_eq!(result.names, vec!["A", "B"]);
        assert_eq!(result.matrix.len(), 2);
        assert_eq!(result.matrix[0].len(), 2);
        assert_eq!(result.matrix[1].len(), 2);
    }

    #[test]
    fn test_distance_matrix_diagonal_zero() {
        let conf = dist_conf(
            &[("A", "ACGT"), ("B", "ACGA"), ("C", "AGGT")],
            SubstitutionModel::PDiff,
        );
        let result = distance_matrix(conf).unwrap();
        for i in 0..3 {
            assert_eq!(result.matrix[i][i], 0.0);
        }
    }

    #[test]
    fn test_distance_matrix_symmetric() {
        let conf = dist_conf(
            &[("A", "ACGT"), ("B", "ACGA"), ("C", "AGGT")],
            SubstitutionModel::PDiff,
        );
        let result = distance_matrix(conf).unwrap();
        for i in 0..3 {
            for j in 0..3 {
                assert_eq!(result.matrix[i][j], result.matrix[j][i]);
            }
        }
    }

    #[test]
    fn test_distance_matrix_pdiff_known_value() {
        // one difference at position 3 (T vs A) out of 4 → 0.25
        let conf = dist_conf(&[("A", "ACGT"), ("B", "ACGA")], SubstitutionModel::PDiff);
        let result = distance_matrix(conf).unwrap();
        assert!((result.matrix[0][1] - 0.25).abs() < 1e-12);
        assert!((result.matrix[1][0] - 0.25).abs() < 1e-12);
    }

    #[test]
    fn test_distance_matrix_identical_sequences_zero() {
        let conf = dist_conf(&[("A", "ACGT"), ("B", "ACGT")], SubstitutionModel::PDiff);
        let result = distance_matrix(conf).unwrap();
        assert_eq!(result.matrix[0][1], 0.0);
    }

    #[test]
    fn test_distance_matrix_jukes_cantor_dna() {
        let conf = dist_conf(
            &[("A", "ACGT"), ("B", "ACGA"), ("C", "AGGT")],
            SubstitutionModel::JukesCantor,
        );
        let result = distance_matrix(conf).unwrap();
        // JC distance for p=0.25: -0.75 * ln(1 - 4/3 * 0.25)
        let expected = -0.75_f64 * (1.0_f64 - (4.0_f64 / 3.0) * 0.25).ln();
        assert!((result.matrix[0][1] - expected).abs() < 1e-10);
    }

    #[test]
    fn test_distance_matrix_kimura2p_dna() {
        let conf = dist_conf(
            &[("A", "ACGT"), ("B", "ACGA"), ("C", "AGGT")],
            SubstitutionModel::Kimura2P,
        );
        let result = distance_matrix(conf).unwrap();
        assert_eq!(result.names, vec!["A", "B", "C"]);
        assert!(result.matrix[0][0] == 0.0);
    }

    #[test]
    fn test_distance_matrix_poisson_protein() {
        let conf = dist_conf(
            &[("A", "ACDEFGH"), ("B", "ACDEFGK")],
            SubstitutionModel::Poisson,
        );
        let result = distance_matrix(conf).unwrap();
        // 1 diff (H vs K) out of 7: p=1/7, d=-ln(1-1/7)
        let expected = -(1.0_f64 - 1.0 / 7.0).ln();
        assert!((result.matrix[0][1] - expected).abs() < 1e-10);
    }

    #[test]
    fn test_distance_matrix_pdiff_protein() {
        let conf = dist_conf(
            &[("A", "ACDEFGH"), ("B", "ACDEFGK")],
            SubstitutionModel::PDiff,
        );
        let result = distance_matrix(conf).unwrap();
        assert!((result.matrix[0][1] - 1.0 / 7.0).abs() < 1e-12);
    }

    #[test]
    fn test_distance_matrix_empty_msa_errors() {
        let conf = DistConfig {
            msa: vec![],
            substitution_model: SubstitutionModel::PDiff,
            alphabet: None,
            num_threads: None,
        };
        assert!(distance_matrix(conf).is_err());
    }

    #[test]
    fn test_distance_matrix_incompatible_model_errors() {
        // Poisson on DNA
        let conf = dist_conf(&[("A", "ACGT"), ("B", "ACGA")], SubstitutionModel::Poisson);
        assert!(distance_matrix(conf).is_err());
        // JukesCantor on Protein
        let conf = dist_conf(
            &[("A", "ACDEFGH"), ("B", "ACDEFGK")],
            SubstitutionModel::JukesCantor,
        );
        assert!(distance_matrix(conf).is_err());
    }

    // --- average_distance ---

    #[test]
    fn test_average_distance_identical_sequences_zero() {
        let conf = dist_conf(&[("A", "ACGT"), ("B", "ACGT")], SubstitutionModel::PDiff);
        let avg = average_distance(conf).unwrap();
        assert_eq!(avg, 0.0);
    }

    #[test]
    fn test_average_distance_two_taxa_equals_pairwise() {
        // one difference out of 4 → 0.25
        let conf = dist_conf(&[("A", "ACGT"), ("B", "ACGA")], SubstitutionModel::PDiff);
        let avg = average_distance(conf).unwrap();
        assert!((avg - 0.25).abs() < 1e-12);
    }

    #[test]
    fn test_average_distance_three_taxa_known_value() {
        // A↔B: 1/4=0.25 (T→A), A↔C: 1/4=0.25 (C→G), B↔C: 2/4=0.5 → avg = 1/3
        let conf = dist_conf(
            &[("A", "ACGT"), ("B", "ACGA"), ("C", "AGGT")],
            SubstitutionModel::PDiff,
        );
        let avg = average_distance(conf).unwrap();
        assert!((avg - 1.0 / 3.0).abs() < 1e-12);
    }

    #[test]
    fn test_average_distance_jukes_cantor_dna() {
        let conf = dist_conf(
            &[("A", "ACGT"), ("B", "ACGA")],
            SubstitutionModel::JukesCantor,
        );
        let avg = average_distance(conf).unwrap();
        let expected = -0.75_f64 * (1.0_f64 - (4.0_f64 / 3.0) * 0.25).ln();
        assert!((avg - expected).abs() < 1e-10);
    }

    #[test]
    fn test_average_distance_empty_msa_errors() {
        let conf = DistConfig {
            msa: vec![],
            substitution_model: SubstitutionModel::PDiff,
            alphabet: None,
            num_threads: None,
        };
        assert!(average_distance(conf).is_err());
    }

    #[test]
    fn test_average_distance_incompatible_model_errors() {
        let conf = dist_conf(&[("A", "ACGT"), ("B", "ACGA")], SubstitutionModel::Poisson);
        assert!(average_distance(conf).is_err());
    }

    #[test]
    fn test_detect_alphabet_dna() {
        let msa = vec![
            SequenceObject {
                identifier: "Seq0".into(),
                sequence: "ACGTACGT".into(),
            },
            SequenceObject {
                identifier: "Seq1".into(),
                sequence: "ACG-ACGT".into(),
            },
        ];
        assert_eq!(detect_alphabet(&msa), Alphabet::DNA);
    }

    #[test]
    fn test_detect_alphabet_dna_iupac() {
        // IUPAC ambiguity codes and RNA U should be detected as DNA, not protein.
        let msa = vec![SequenceObject {
            identifier: "Seq0".into(),
            sequence: "ACGTRYWSMKHBDVNU".into(),
        }];
        assert_eq!(detect_alphabet(&msa), Alphabet::DNA);
    }

    #[test]
    fn test_detect_alphabet_protein() {
        let msa = vec![
            SequenceObject {
                identifier: "Seq0".into(),
                sequence: "ACDEFGHIK".into(),
            },
            SequenceObject {
                identifier: "Seq1".into(),
                sequence: "ACD-FGHIK".into(),
            },
        ];
        assert_eq!(detect_alphabet(&msa), Alphabet::Protein);
    }
}

#[cfg(all(test, feature = "parallel"))]
mod parallel_tests {
    use super::*;
    use crate::models::SubstitutionModel;
    use std::sync::Arc;
    use std::sync::atomic::{AtomicUsize, Ordering};

    fn three_seq_dna() -> Vec<SequenceObject> {
        vec![
            SequenceObject {
                identifier: "A".into(),
                sequence: "ACGTACGT".into(),
            },
            SequenceObject {
                identifier: "B".into(),
                sequence: "ACGCACGT".into(),
            },
            SequenceObject {
                identifier: "C".into(),
                sequence: "ACGTACGC".into(),
            },
        ]
    }

    fn base_conf(msa: Vec<SequenceObject>, n_bootstrap_samples: usize) -> NJConfig {
        NJConfig {
            msa,
            n_bootstrap_samples,
            substitution_model: SubstitutionModel::PDiff,
            alphabet: None,
            num_threads: None,
            return_distance_matrix: false,
            return_average_distance: false,
        }
    }

    #[test]
    fn test_parallel_bootstrap_returns_valid_newick() {
        let result = nj(base_conf(three_seq_dna(), 20), None).expect("parallel NJ failed");
        assert!(result.newick.ends_with(';'));
        assert!(result.newick.contains(':'));
    }

    #[test]
    fn test_parallel_progress_fires_exactly_n_times() {
        let n = 10_usize;
        let count = Arc::new(AtomicUsize::new(0));
        let count2 = count.clone();
        let cb: Box<dyn Fn(NJEvent)> = Box::new(move |event| {
            if let NJEvent::BootstrapProgress { .. } = event {
                count2.fetch_add(1, Ordering::SeqCst);
            }
        });
        nj(base_conf(three_seq_dna(), n), Some(cb)).unwrap();
        assert_eq!(count.load(Ordering::SeqCst), n);
    }

    #[test]
    fn test_parallel_progress_last_call_is_total() {
        let n = 8_usize;
        let last = Arc::new(AtomicUsize::new(0));
        let last2 = last.clone();
        let cb: Box<dyn Fn(NJEvent)> = Box::new(move |event| {
            if let NJEvent::BootstrapProgress { completed, .. } = event {
                last2.store(completed, Ordering::SeqCst);
            }
        });
        nj(base_conf(three_seq_dna(), n), Some(cb)).unwrap();
        assert_eq!(last.load(Ordering::SeqCst), n);
    }
}
