use std::io;

use crate::tempbitvec::{TempBitVec, TempBitVecBuilder};
use crate::tempintvec::TempCompactIntVec;

// ── ColGroup ──────────────────────────────────────────────────────────────────

/// A named subset of columns, identified by their indices within the matrix.
///
/// Defined once at the index level; the same indices are valid across all
/// partitions and layers because the column structure (samples / genomes) is
/// identical everywhere — only the row space (kmer slots) is partitioned.
pub struct ColGroup {
    pub name:    String,
    pub indices: Vec<usize>,
}

impl ColGroup {
    pub fn new(name: impl Into<String>, indices: Vec<usize>) -> Self {
        Self { name: name.into(), indices }
    }
}

// ── MatrixGroupOps ────────────────────────────────────────────────────────────

/// Per-matrix group aggregations.
///
/// `partial_group_presence_count`, `partial_group_sum`, `partial_group_any`,
/// `partial_group_min`, `partial_group_max` are the primitives; each impl must
/// provide all five.
///
/// `partial_group_all` and `partial_group_none` have default implementations
/// derived from `partial_group_presence_count` and should rarely need overriding.
pub trait MatrixGroupOps {
    /// Per-slot count of group columns whose value ≥ `threshold`.
    fn partial_group_presence_count(&self, g: &ColGroup, threshold: u32) -> io::Result<TempCompactIntVec>;

    /// Per-slot sum of values across all group columns.
    fn partial_group_sum(&self, g: &ColGroup) -> io::Result<TempCompactIntVec>;

    /// Per-slot OR: 1 if any group column has value ≥ `threshold`.
    fn partial_group_any(&self, g: &ColGroup, threshold: u32) -> io::Result<TempBitVec>;

    /// Per-slot min value across all group columns (0 if group is empty).
    fn partial_group_min(&self, g: &ColGroup) -> io::Result<TempCompactIntVec>;

    /// Per-slot max value across all group columns (0 if group is empty).
    fn partial_group_max(&self, g: &ColGroup) -> io::Result<TempCompactIntVec>;

    /// Per-slot AND: 1 if ALL group columns have value ≥ `threshold`.
    fn partial_group_all(&self, g: &ColGroup, threshold: u32) -> io::Result<TempBitVec> {
        let counts = self.partial_group_presence_count(g, threshold)?;
        let n = counts.len();
        let n_required = g.indices.len() as u32;
        let mut b = TempBitVecBuilder::new(n)?;
        b.or_where(counts.view(), |v| v >= n_required);
        b.freeze()
    }

    /// Per-slot NOR: 1 if NO group column has value ≥ `threshold`.
    fn partial_group_none(&self, g: &ColGroup, threshold: u32) -> io::Result<TempBitVec> {
        let counts = self.partial_group_presence_count(g, threshold)?;
        let n = counts.len();
        let mut b = TempBitVecBuilder::new(n)?;
        b.or_where(counts.view(), |v| v == 0);
        b.freeze()
    }
}

// ── FilterMask — expression tree for column-based slot filters ────────────────

/// A composable filter expression that can be evaluated against a matrix
/// using only column operations (no MPHF lookup per kmer).
///
/// `threshold` semantics follow [`MatrixGroupOps::partial_group_presence_count`]:
/// a slot contributes to the count when its value is **≥ threshold**.
/// To match the row-level filter (`value > t`), callers should pass `t + 1`.
#[derive(Debug, Clone)]
pub enum FilterMask {
    /// Slot passes if count of columns in `indices` with value ≥ `threshold` is ≥ `min_count`.
    PresenceGeq { indices: Vec<usize>, threshold: u32, min_count: usize },
    /// Slot passes if count of columns in `indices` with value ≥ `threshold` is ≤ `max_count`.
    PresenceLeq { indices: Vec<usize>, threshold: u32, max_count: usize },
    /// Slot passes if sum of values across `indices` columns is ≥ `min_sum`.
    SumGeq { indices: Vec<usize>, min_sum: u32 },
    /// Slot passes if sum of values across `indices` columns is ≤ `max_sum`.
    SumLeq { indices: Vec<usize>, max_sum: u32 },
    /// Slot passes if it passes all sub-expressions. Empty `And` is always true.
    And(Vec<FilterMask>),
}

/// Evaluate a [`FilterMask`] against `mat`, returning a per-slot `TempBitVec`
/// where bit=1 means the slot passes the filter.
pub fn eval_filter_mask(expr: &FilterMask, mat: &dyn MatrixGroupOps, n: usize) -> io::Result<TempBitVec> {
    match expr {
        FilterMask::PresenceGeq { indices, threshold, min_count } => {
            let g = ColGroup::new("", indices.clone());
            let counts = mat.partial_group_presence_count(&g, *threshold)?;
            let mut b = TempBitVecBuilder::new(n)?;
            let mc = *min_count as u32;
            b.or_where(counts.view(), |v| v >= mc);
            b.freeze()
        }
        FilterMask::PresenceLeq { indices, threshold, max_count } => {
            let g = ColGroup::new("", indices.clone());
            let counts = mat.partial_group_presence_count(&g, *threshold)?;
            let mut b = TempBitVecBuilder::new(n)?;
            let mc = *max_count as u32;
            b.or_where(counts.view(), |v| v <= mc);
            b.freeze()
        }
        FilterMask::SumGeq { indices, min_sum } => {
            let g = ColGroup::new("", indices.clone());
            let sums = mat.partial_group_sum(&g)?;
            let mut b = TempBitVecBuilder::new(n)?;
            let ms = *min_sum;
            b.or_where(sums.view(), |v| v >= ms);
            b.freeze()
        }
        FilterMask::SumLeq { indices, max_sum } => {
            let g = ColGroup::new("", indices.clone());
            let sums = mat.partial_group_sum(&g)?;
            let mut b = TempBitVecBuilder::new(n)?;
            let ms = *max_sum;
            b.or_where(sums.view(), |v| v <= ms);
            b.freeze()
        }
        FilterMask::And(parts) => {
            let mut b = TempBitVecBuilder::new_ones(n)?;
            for part in parts {
                let m = eval_filter_mask(part, mat, n)?;
                b.and(m.view());
            }
            b.freeze()
        }
    }
}
