fast-nnt 0.4.0

use fixedbitset::FixedBitSet;
use std::cmp::Ordering;
use std::fmt::{self, Display};
use std::hash::{Hash, Hasher};

/// Bipartition of a 1-based taxon set into two non-empty parts A and B.
/// Indices are 1..=ntax (bit 0 unused); the constructor swaps parts so the
/// one whose first set bit is smaller becomes A.
#[derive(Debug, Clone)]
pub struct BiPartition {
    pub a: FixedBitSet,
    pub b: FixedBitSet,
}

impl BiPartition {
    /// Both parts must be non-empty and disjoint; together they define the
    /// taxa universe.
    pub fn new(a: FixedBitSet, b: FixedBitSet) -> Self {
        let ca = cardinality(&a);
        let cb = cardinality(&b);
        debug_assert!(ca > 0 && cb > 0, "invalid split: |A|={ca}, |B|={cb}");
        let a_first = first_set_at_or_after(&a, 1).unwrap_or(usize::MAX);
        let b_first = first_set_at_or_after(&b, 1).unwrap_or(usize::MAX);

        if a_first < b_first {
            Self { a, b }
        } else {
            Self { a: b, b: a }
        }
    }

    /// Total number of taxa: |A| + |B|.
    pub fn ntax(&self) -> usize {
        cardinality(&self.a) + cardinality(&self.b)
    }

    /// Size of the smaller side: min(|A|, |B|).
    pub fn size(&self) -> usize {
        cardinality(&self.a).min(cardinality(&self.b))
    }

    /// Part A (1-based).
    pub fn get_a(&self) -> &FixedBitSet {
        &self.a
    }

    /// Part B (1-based).
    pub fn get_b(&self) -> &FixedBitSet {
        &self.b
    }

    /// Split part that contains taxon `t` (1-based); A if neither contains it.
    pub fn part_containing(&self, t: usize) -> &FixedBitSet {
        if self.b.contains(t) { &self.b } else { &self.a }
    }

    /// The part that does *not* contain `t` (1-based).
    pub fn part_not_containing(&self, t: usize) -> &FixedBitSet {
        if !self.a.contains(t) {
            &self.a
        } else {
            &self.b
        }
    }

    /// The smaller part; on tie, the one whose first set bit (≥1) is smaller.
    pub fn smaller_part(&self) -> &FixedBitSet {
        let ca = cardinality(&self.a);
        let cb = cardinality(&self.b);
        match ca.cmp(&cb) {
            Ordering::Less => &self.a,
            Ordering::Greater => &self.b,
            Ordering::Equal => {
                let af = first_set_at_or_after(&self.a, 1).unwrap_or(usize::MAX);
                let bf = first_set_at_or_after(&self.b, 1).unwrap_or(usize::MAX);
                if af < bf { &self.a } else { &self.b }
            }
        }
    }

    /// Is this a trivial split (smaller side has cardinality 1)?
    pub fn is_trivial(&self) -> bool {
        cardinality(self.smaller_part()) == 1
    }
}

// ---- Equality / Hash / Ordering ----

impl PartialEq for BiPartition {
    fn eq(&self, other: &Self) -> bool {
        bitset_eq(&self.a, &other.a) && bitset_eq(&self.b, &other.b)
    }
}
impl Eq for BiPartition {}

impl Hash for BiPartition {
    fn hash<H: Hasher>(&self, state: &mut H) {
        // Hash by set elements (stable across implementations)
        for i in self.a.ones() {
            i.hash(state);
        }
        0usize.hash(state); // separator
        for i in self.b.ones() {
            i.hash(state);
        }
    }
}

impl PartialOrd for BiPartition {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}
impl Ord for BiPartition {
    fn cmp(&self, other: &Self) -> Ordering {
        match compare_bitsets(&self.a, &other.a) {
            Ordering::Equal => compare_bitsets(&self.b, &other.b),
            ord => ord,
        }
    }
}

impl Display for BiPartition {
    /// Formats as "{i1,i2,...} | {j1,j2,...}" with 1-based indices ascending.
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fn part_to_string(bs: &FixedBitSet) -> String {
            let mut first = true;
            let mut s = String::new();
            for i in bs.ones() {
                if i == 0 {
                    continue;
                } // ignore 0 to keep 1-based view
                if !first {
                    s.push(',');
                } else {
                    first = false;
                }
                s.push_str(&i.to_string());
            }
            s
        }
        write!(
            f,
            "{{{}}} | {{{}}}",
            part_to_string(&self.a),
            part_to_string(&self.b)
        )
    }
}

// ---- Small helpers (BitSet-like utilities) ----

fn cardinality(bs: &FixedBitSet) -> usize {
    bs.ones().filter(|&i| i != 0).count()
}

fn first_set_at_or_after(bs: &FixedBitSet, start: usize) -> Option<usize> {
    bs.ones().find(|&i| i >= start)
}

fn bitset_eq(a: &FixedBitSet, b: &FixedBitSet) -> bool {
    // FixedBitSet equality is structural (len + bits); for safety, compare membership.
    // (In practice, `==` would also work if lengths are kept consistent.)
    let a_len = a.ones().count();
    let b_len = b.ones().count();
    if a_len != b_len {
        return false;
    }
    // Same members?
    for i in a.ones() {
        if !b.contains(i) {
            return false;
        }
    }
    true
}

/// Lexicographic compare of bit sets by ascending element list (ignoring 0).
fn compare_bitsets(a: &FixedBitSet, b: &FixedBitSet) -> Ordering {
    let mut ia = a.ones().filter(|&i| i != 0);
    let mut ib = b.ones().filter(|&i| i != 0);
    loop {
        match (ia.next(), ib.next()) {
            (Some(x), Some(y)) => {
                if x != y {
                    return x.cmp(&y);
                }
            }
            (None, Some(_)) => return Ordering::Less,
            (Some(_), None) => return Ordering::Greater,
            (None, None) => return Ordering::Equal,
        }
    }
}