use obicompactvec::FilterMask;

/// Trait for kmer row filters.
///
/// `row` contains raw per-genome counts (or 0/1 for presence/absence data).
/// `n_genomes` equals `row.len()`.
pub trait KmerFilter: Send + Sync {
    fn passes(&self, row: &[u32], n_genomes: usize) -> bool;

    /// Express this filter as a [`FilterMask`] column-operation expression.
    ///
    /// Returns `Some(expr)` if the filter can be evaluated solely from matrix
    /// column aggregates (no per-kmer row scan needed).  Returns `None` if the
    /// filter requires row-level inspection.
    ///
    /// `threshold` semantics in the returned mask use `>= threshold`, matching
    /// [`obicompactvec::MatrixGroupOps`].  Implementations must add 1 to any
    /// row-level threshold that uses strict `>` comparison.
    fn column_mask_expr(&self, _n_genomes: usize) -> Option<FilterMask> {
        None
    }
}

/// True when `row` passes every filter in `filters`.
/// Returns `true` if `filters` is empty.
pub fn passes_all(filters: &[Box<dyn KmerFilter>], row: &[u32], n_genomes: usize) -> bool {
    filters.iter().all(|f| f.passes(row, n_genomes))
}

// ── Quorum filters ─────────────────────────────────────────────────────────────

fn present_count(row: &[u32], threshold: u32) -> usize {
    row.iter().filter(|&&v| v > threshold).count()
}

/// At least `frac` fraction of genomes contain this kmer (count > `threshold`).
pub struct MinGenomeFraction {
    pub frac: f64,
    pub threshold: u32,
}

impl KmerFilter for MinGenomeFraction {
    fn passes(&self, row: &[u32], n_genomes: usize) -> bool {
        let p = present_count(row, self.threshold);
        p as f64 / n_genomes as f64 >= self.frac
    }

    fn column_mask_expr(&self, n_genomes: usize) -> Option<FilterMask> {
        let t = self.threshold.checked_add(1)?;
        let min_count = (self.frac * n_genomes as f64).ceil() as usize;
        Some(FilterMask::PresenceGeq {
            indices: (0..n_genomes).collect(),
            threshold: t,
            min_count,
        })
    }
}

/// At most `frac` fraction of genomes contain this kmer (count > `threshold`).
pub struct MaxGenomeFraction {
    pub frac: f64,
    pub threshold: u32,
}

impl KmerFilter for MaxGenomeFraction {
    fn passes(&self, row: &[u32], n_genomes: usize) -> bool {
        let p = present_count(row, self.threshold);
        p as f64 / n_genomes as f64 <= self.frac
    }

    fn column_mask_expr(&self, n_genomes: usize) -> Option<FilterMask> {
        let t = self.threshold.checked_add(1)?;
        let max_count = (self.frac * n_genomes as f64).floor() as usize;
        Some(FilterMask::PresenceLeq {
            indices: (0..n_genomes).collect(),
            threshold: t,
            max_count,
        })
    }
}

/// At least `count` genomes contain this kmer (count > `threshold`).
pub struct MinGenomeCount {
    pub count: usize,
    pub threshold: u32,
}

impl KmerFilter for MinGenomeCount {
    fn passes(&self, row: &[u32], _n_genomes: usize) -> bool {
        present_count(row, self.threshold) >= self.count
    }

    fn column_mask_expr(&self, n_genomes: usize) -> Option<FilterMask> {
        let t = self.threshold.checked_add(1)?;
        Some(FilterMask::PresenceGeq {
            indices: (0..n_genomes).collect(),
            threshold: t,
            min_count: self.count,
        })
    }
}

/// At most `count` genomes contain this kmer (count > `threshold`).
pub struct MaxGenomeCount {
    pub count: usize,
    pub threshold: u32,
}

impl KmerFilter for MaxGenomeCount {
    fn passes(&self, row: &[u32], _n_genomes: usize) -> bool {
        present_count(row, self.threshold) <= self.count
    }

    fn column_mask_expr(&self, n_genomes: usize) -> Option<FilterMask> {
        let t = self.threshold.checked_add(1)?;
        Some(FilterMask::PresenceLeq {
            indices: (0..n_genomes).collect(),
            threshold: t,
            max_count: self.count,
        })
    }
}

// ── Total-count filters (count indexes only) ───────────────────────────────────

/// Sum of counts across all genomes >= `total`.
pub struct MinTotalCount {
    pub total: u32,
}

impl KmerFilter for MinTotalCount {
    fn passes(&self, row: &[u32], _n_genomes: usize) -> bool {
        row.iter().sum::<u32>() >= self.total
    }

    fn column_mask_expr(&self, n_genomes: usize) -> Option<FilterMask> {
        Some(FilterMask::SumGeq {
            indices: (0..n_genomes).collect(),
            min_sum: self.total,
        })
    }
}

/// Sum of counts across all genomes <= `total`.
pub struct MaxTotalCount {
    pub total: u32,
}

impl KmerFilter for MaxTotalCount {
    fn passes(&self, row: &[u32], _n_genomes: usize) -> bool {
        row.iter().sum::<u32>() <= self.total
    }

    fn column_mask_expr(&self, n_genomes: usize) -> Option<FilterMask> {
        Some(FilterMask::SumLeq {
            indices: (0..n_genomes).collect(),
            max_sum: self.total,
        })
    }
}

// ── Group-based quorum filter ─────────────────────────────────────────────────

/// Quorum filter operating on pre-classified genome groups.
///
/// `ingroup_idx` / `outgroup_idx` are column indices into the per-genome row.
/// When `ingroup_idx` is empty, no ingroup quorum is checked.
/// When `outgroup_idx` is empty, no outgroup quorum is checked.
pub struct GroupQuorumFilter {
    pub ingroup_idx:        Vec<usize>,
    pub outgroup_idx:       Vec<usize>,
    pub threshold:          u32,
    pub min_count:          usize,
    pub max_count:          usize,
    pub min_frac:           f64,
    pub max_frac:           f64,
    pub min_outgroup_count: usize,
    pub max_outgroup_count: usize,
    pub min_outgroup_frac:  f64,
    pub max_outgroup_frac:  f64,
}

impl GroupQuorumFilter {
    // Build PresenceGeq/PresenceLeq constraints for one group (ingroup or outgroup).
    fn group_mask_parts(
        indices: &[usize],
        threshold: u32,
        min_count: usize,
        max_count: usize,
        min_frac: f64,
        max_frac: f64,
        parts: &mut Vec<FilterMask>,
    ) {
        let n = indices.len();
        let geq = min_count.max((min_frac * n as f64).ceil() as usize);
        if geq > 0 {
            parts.push(FilterMask::PresenceGeq {
                indices: indices.to_vec(),
                threshold,
                min_count: geq,
            });
        }
        let leq = max_count.min((max_frac * n as f64).floor() as usize);
        if leq < n {
            parts.push(FilterMask::PresenceLeq {
                indices: indices.to_vec(),
                threshold,
                max_count: leq,
            });
        }
    }
}

impl KmerFilter for GroupQuorumFilter {
    fn passes(&self, row: &[u32], _n_genomes: usize) -> bool {
        if !self.ingroup_idx.is_empty() {
            let n = self.ingroup_idx.iter()
                .filter(|&&i| row.get(i).copied().unwrap_or(0) > self.threshold)
                .count();
            let denom = self.ingroup_idx.len();
            if n < self.min_count { return false; }
            if n > self.max_count { return false; }
            let frac = n as f64 / denom as f64;
            if frac < self.min_frac { return false; }
            if frac > self.max_frac { return false; }
        }
        if !self.outgroup_idx.is_empty() {
            let n = self.outgroup_idx.iter()
                .filter(|&&i| row.get(i).copied().unwrap_or(0) > self.threshold)
                .count();
            let denom = self.outgroup_idx.len();
            if n < self.min_outgroup_count { return false; }
            if n > self.max_outgroup_count { return false; }
            let frac = n as f64 / denom as f64;
            if frac < self.min_outgroup_frac { return false; }
            if frac > self.max_outgroup_frac { return false; }
        }
        true
    }

    fn column_mask_expr(&self, _n_genomes: usize) -> Option<FilterMask> {
        let t = self.threshold.checked_add(1)?;
        let mut parts: Vec<FilterMask> = Vec::new();
        if !self.ingroup_idx.is_empty() {
            Self::group_mask_parts(
                &self.ingroup_idx, t,
                self.min_count, self.max_count,
                self.min_frac, self.max_frac,
                &mut parts,
            );
        }
        if !self.outgroup_idx.is_empty() {
            Self::group_mask_parts(
                &self.outgroup_idx, t,
                self.min_outgroup_count, self.max_outgroup_count,
                self.min_outgroup_frac, self.max_outgroup_frac,
                &mut parts,
            );
        }
        Some(FilterMask::And(parts))
    }
}
