birdnet_onnx/

rangefilter.rs

//! Range filter for location and date-based species filtering
//!
//! # Usage Examples
//!
//! ## Basic Filtering
//!
//! ```ignore
//! use birdnet_onnx::{Classifier, RangeFilter};
//!
//! let classifier = Classifier::builder()
//!     .model_path("birdnet.onnx")
//!     .labels_path("labels.txt")
//!     .build()?;
//!
//! let range_filter = RangeFilter::builder()
//!     .model_path("meta_model.onnx")
//!     .from_classifier_labels(classifier.labels())
//!     .threshold(0.01)
//!     .build()?;
//!
//! // Get predictions
//! let result = classifier.predict(&audio_segment)?;
//!
//! // Get location scores
//! let location_scores = range_filter.predict(60.17, 24.94, 6, 15)?;
//!
//! // Filter predictions
//! let filtered = range_filter.filter_predictions(
//!     &result.predictions,
//!     &location_scores,
//!     false, // don't rerank
//! );
//! ```
//!
//! ## Batch Processing
//!
//! ```ignore
//! // Calculate location scores once
//! let location_scores = range_filter.predict(lat, lon, month, day)?;
//!
//! // Process multiple audio segments
//! let mut predictions_batch = Vec::new();
//! for segment in audio_segments {
//!     let result = classifier.predict(&segment)?;
//!     predictions_batch.push(result.predictions);
//! }
//!
//! // Filter all at once
//! let filtered_batch = range_filter.filter_batch_predictions(
//!     predictions_batch,
//!     &location_scores,
//!     true, // rerank by location score
//! );
//! ```

use crate::error::{Error, Result};
use crate::labels::parse_labels;
use crate::types::{LabelFormat, LocationScore, Prediction};
use ndarray::Array2;
use ort::session::Session;
use ort::value::Value;
use std::sync::{Arc, Mutex};

/// Calculate week number for `BirdNET` meta model (48-week year, 4 weeks per month).
///
/// `BirdNET` assumes each month has exactly 4 weeks, creating a 48-week year.
/// Week calculation: weeksFromMonths = (month - 1) * 4; weekInMonth = (day - 1) / 7 + 1
///
/// # Arguments
/// * `month` - Month number (1-12)
/// * `day` - Day of month (1-31)
///
/// # Returns
/// Week number as f32 (typically 1-48, but can exceed 48 for days 29-31)
#[must_use]
#[allow(clippy::cast_precision_loss)]
pub const fn calculate_week(month: u32, day: u32) -> f32 {
    let weeks_from_months = (month - 1) * 4;
    let week_in_month = (day - 1) / 7 + 1;
    (weeks_from_months + week_in_month) as f32
}

/// Validate geographic coordinates.
///
/// # Arguments
/// * `latitude` - Latitude in degrees (-90 to 90)
/// * `longitude` - Longitude in degrees (-180 to 180)
///
/// # Errors
/// Returns `Error::InvalidCoordinates` if values are out of range
pub fn validate_coordinates(latitude: f32, longitude: f32) -> Result<()> {
    if !(-90.0..=90.0).contains(&latitude) {
        return Err(Error::InvalidCoordinates {
            latitude,
            longitude,
            reason: format!("latitude must be in range [-90, 90], got {latitude}"),
        });
    }
    if !(-180.0..=180.0).contains(&longitude) {
        return Err(Error::InvalidCoordinates {
            latitude,
            longitude,
            reason: format!("longitude must be in range [-180, 180], got {longitude}"),
        });
    }
    Ok(())
}

/// Validate date parameters for `BirdNET` calendar.
///
/// # Arguments
/// * `month` - Month number (1-12)
/// * `day` - Day of month (1-31)
///
/// # Errors
/// Returns `Error::InvalidDate` if values are out of range
pub fn validate_date(month: u32, day: u32) -> Result<()> {
    if !(1..=12).contains(&month) {
        return Err(Error::InvalidDate {
            month,
            day,
            reason: format!("month must be in range [1, 12], got {month}"),
        });
    }
    if !(1..=31).contains(&day) {
        return Err(Error::InvalidDate {
            month,
            day,
            reason: format!("day must be in range [1, 31], got {day}"),
        });
    }
    Ok(())
}

/// Labels source for builder
#[derive(Debug)]
enum Labels {
    Path(String),
    InMemory(Vec<String>),
}

/// Builder for constructing a `RangeFilter`
#[derive(Debug)]
pub struct RangeFilterBuilder {
    model_path: Option<String>,
    labels: Option<Labels>,
    execution_providers: Vec<ort::execution_providers::ExecutionProviderDispatch>,
    threshold: f32,
}

impl Default for RangeFilterBuilder {
    fn default() -> Self {
        Self::new()
    }
}

impl RangeFilterBuilder {
    /// Create a new range filter builder
    #[must_use]
    pub const fn new() -> Self {
        Self {
            model_path: None,
            labels: None,
            execution_providers: Vec::new(),
            threshold: 0.01,
        }
    }

    /// Set the path to the ONNX meta model file (required)
    #[must_use]
    pub fn model_path(mut self, path: impl Into<String>) -> Self {
        self.model_path = Some(path.into());
        self
    }

    /// Set the path to the labels file (required, must match model output size)
    #[must_use]
    pub fn labels_path(mut self, path: impl Into<String>) -> Self {
        self.labels = Some(Labels::Path(path.into()));
        self
    }

    /// Set species labels directly (required, must match model output size)
    #[must_use]
    pub fn labels(mut self, labels: Vec<String>) -> Self {
        self.labels = Some(Labels::InMemory(labels));
        self
    }

    /// Use labels from an existing Classifier.
    ///
    /// This is a convenience method that copies labels from a classifier,
    /// ensuring they stay in sync with the main model.
    #[must_use]
    pub fn from_classifier_labels(mut self, labels: &[String]) -> Self {
        self.labels = Some(Labels::InMemory(labels.to_vec()));
        self
    }

    /// Add an execution provider (GPU, CPU, etc.)
    #[must_use]
    pub fn execution_provider(
        mut self,
        provider: impl Into<ort::execution_providers::ExecutionProviderDispatch>,
    ) -> Self {
        self.execution_providers.push(provider.into());
        self
    }

    /// Set minimum score threshold (default: 0.01)
    #[must_use]
    pub const fn threshold(mut self, threshold: f32) -> Self {
        self.threshold = threshold;
        self
    }

    /// Build the range filter
    ///
    /// # Errors
    /// Returns error if model path or labels not set, or if model loading fails
    pub fn build(self) -> Result<RangeFilter> {
        let model_path = self.model_path.ok_or(Error::ModelPathRequired)?;
        let labels_source = self.labels.ok_or(Error::LabelsRequired)?;

        // Load labels from file or use in-memory vector
        let labels = match labels_source {
            Labels::Path(path) => {
                // Read and parse labels file (text format: one label per line)
                let content = std::fs::read_to_string(&path).map_err(|e| Error::LabelLoad {
                    path: path.clone(),
                    reason: e.to_string(),
                })?;
                parse_labels(&content, LabelFormat::Text)?
            }
            Labels::InMemory(labels) => labels,
        };

        // Build session with execution providers
        let mut session_builder = Session::builder().map_err(Error::ModelLoad)?;

        if !self.execution_providers.is_empty() {
            session_builder = session_builder
                .with_execution_providers(self.execution_providers)
                .map_err(Error::ModelLoad)?;
        }

        let session = session_builder
            .commit_from_file(&model_path)
            .map_err(Error::ModelLoad)?;

        // Validate label count matches model output
        let output_shapes = extract_output_shapes(&session)?;

        // Meta model should have exactly one output
        if output_shapes.len() != 1 {
            return Err(Error::ModelDetection {
                reason: format!("meta model expects 1 output, got {}", output_shapes.len()),
            });
        }

        let expected = extract_last_dim(&output_shapes[0])?;
        if labels.len() != expected {
            return Err(Error::LabelCount {
                expected,
                got: labels.len(),
            });
        }

        Ok(RangeFilter {
            inner: Arc::new(RangeFilterInner {
                session: Mutex::new(session),
                labels,
                threshold: self.threshold,
            }),
        })
    }
}

/// Extract output tensor shapes from session
fn extract_output_shapes(session: &Session) -> Result<Vec<Vec<i64>>> {
    session
        .outputs
        .iter()
        .map(|output| {
            let shape = output
                .output_type
                .tensor_shape()
                .ok_or_else(|| Error::ModelDetection {
                    reason: "output is not a tensor".to_string(),
                })?;
            Ok(shape.iter().copied().collect())
        })
        .collect()
}

/// Extract last dimension from output shape
fn extract_last_dim(shape: &[i64]) -> Result<usize> {
    let value = shape.last().copied().ok_or_else(|| Error::ModelDetection {
        reason: "empty output shape".to_string(),
    })?;

    usize::try_from(value).map_err(|_| Error::ModelDetection {
        reason: format!("invalid dimension: {value}"),
    })
}

/// Filter multiple prediction sets with the same location scores.
///
/// This is a helper for batch processing - runs filtering on each
/// prediction set using the same location scores.
fn filter_batch_predictions_impl(
    predictions_batch: Vec<Vec<crate::types::Prediction>>,
    location_scores: &[LocationScore],
    threshold: f32,
    rerank: bool,
) -> Vec<Vec<crate::types::Prediction>> {
    predictions_batch
        .into_iter()
        .map(|preds| filter_predictions_impl(&preds, location_scores, threshold, rerank))
        .collect()
}

/// Filter predictions based on meta model location scores
///
/// # Arguments
/// * `predictions` - Original predictions from audio analysis
/// * `location_scores` - Location-based species scores from meta model
/// * `threshold` - Minimum location score threshold
/// * `rerank` - Whether to rerank by location score (multiply confidence by location score)
///
/// # Returns
/// Filtered predictions, optionally reranked by location score
fn filter_predictions_impl(
    predictions: &[Prediction],
    location_scores: &[LocationScore],
    threshold: f32,
    rerank: bool,
) -> Vec<Prediction> {
    // Build lookup map from species to location score
    let location_map: std::collections::HashMap<&str, f32> = location_scores
        .iter()
        .map(|score| (score.species.as_str(), score.score))
        .collect();

    // Filter and optionally rerank predictions
    let mut filtered: Vec<Prediction> = predictions
        .iter()
        .filter_map(|pred| {
            let location_score = location_map.get(pred.species.as_str()).copied();
            match location_score {
                Some(score) if score >= threshold => {
                    // Species in meta model with score >= threshold: keep and optionally rerank
                    let confidence = if rerank {
                        pred.confidence * score
                    } else {
                        pred.confidence
                    };
                    Some(Prediction {
                        species: pred.species.clone(),
                        confidence,
                        index: pred.index,
                    })
                }
                Some(_) => {
                    // Species in meta model with score < threshold: filter out
                    None
                }
                None => {
                    // Species NOT in meta model: keep unchanged
                    Some(Prediction {
                        species: pred.species.clone(),
                        confidence: pred.confidence,
                        index: pred.index,
                    })
                }
            }
        })
        .collect();

    // Re-sort by confidence descending if reranked
    if rerank {
        filtered.sort_unstable_by(|a, b| b.confidence.total_cmp(&a.confidence));
    }

    filtered
}

/// Internal state for `RangeFilter`
struct RangeFilterInner {
    session: Mutex<Session>,
    labels: Vec<String>,
    threshold: f32,
}

/// Thread-safe range filter for location-based species filtering
#[derive(Clone)]
pub struct RangeFilter {
    inner: Arc<RangeFilterInner>,
}

impl std::fmt::Debug for RangeFilter {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("RangeFilter")
            .field("labels_count", &self.inner.labels.len())
            .field("threshold", &self.inner.threshold)
            .finish_non_exhaustive()
    }
}

impl RangeFilter {
    /// Create a new range filter builder
    #[must_use]
    pub const fn builder() -> RangeFilterBuilder {
        RangeFilterBuilder::new()
    }

    /// Get species probability scores for given location and date
    ///
    /// # Arguments
    /// * `latitude` - Latitude in degrees (-90 to 90)
    /// * `longitude` - Longitude in degrees (-180 to 180)
    /// * `month` - Month number (1-12)
    /// * `day` - Day of month (1-31)
    ///
    /// # Returns
    /// Vector of `LocationScore` sorted by score (descending)
    ///
    /// # Errors
    /// Returns error if:
    /// - Coordinates are invalid (latitude not in [-90, 90] or longitude not in [-180, 180])
    /// - Date parameters are invalid (month not in [1, 12] or day not in [1, 31])
    /// - Session lock is poisoned
    /// - ONNX inference fails
    #[allow(clippy::significant_drop_tightening)]
    pub fn predict(
        &self,
        latitude: f32,
        longitude: f32,
        month: u32,
        day: u32,
    ) -> Result<Vec<LocationScore>> {
        // Validate coordinates
        validate_coordinates(latitude, longitude)?;

        // Validate date parameters
        validate_date(month, day)?;

        // Calculate week number
        let week = calculate_week(month, day);

        // Create input tensor [1, 3] with [latitude, longitude, week]
        let input_data = vec![latitude, longitude, week];
        let input_array = Array2::from_shape_vec((1, 3), input_data).map_err(|e| {
            Error::RangeFilterInference(format!("failed to create input array: {e}"))
        })?;

        let input_value = Value::from_array(input_array).map_err(|e| {
            Error::RangeFilterInference(format!("failed to create input tensor: {e}"))
        })?;

        // Run inference with locked session
        let mut session = self
            .inner
            .session
            .lock()
            .map_err(|e| Error::RangeFilterInference(format!("session lock poisoned: {e}")))?;

        let outputs = session
            .run(ort::inputs![input_value])
            .map_err(|e| Error::RangeFilterInference(e.to_string()))?;

        // Extract output tensor (build validates exactly one output exists)
        let tensor = outputs.values().next().ok_or_else(|| {
            Error::RangeFilterInference("model returned no output tensors".to_string())
        })?;

        let (_, data) = tensor
            .try_extract_tensor::<f32>()
            .map_err(|e| Error::RangeFilterInference(e.to_string()))?;

        // Build scores above threshold
        let mut scores: Vec<LocationScore> = data
            .iter()
            .enumerate()
            .filter_map(|(i, &score)| {
                if score >= self.inner.threshold && i < self.inner.labels.len() {
                    Some(LocationScore {
                        species: self.inner.labels[i].clone(),
                        score,
                        index: i,
                    })
                } else {
                    None
                }
            })
            .collect();

        // Sort by score descending
        scores.sort_unstable_by(|a, b| b.score.total_cmp(&a.score));

        Ok(scores)
    }

    /// Filter predictions based on location scores from meta model
    ///
    /// # Arguments
    /// * `predictions` - Original predictions from audio analysis
    /// * `location_scores` - Location-based species scores (from `predict`)
    /// * `rerank` - Whether to rerank by multiplying confidence by location score
    ///
    /// # Returns
    /// Filtered predictions, optionally reranked by location score
    ///
    /// # Example
    /// ```no_run
    /// # use birdnet_onnx::{RangeFilter, Prediction};
    /// # fn example(filter: &RangeFilter, predictions: Vec<Prediction>) -> birdnet_onnx::Result<()> {
    /// // Get location scores for a specific place and time
    /// let location_scores = filter.predict(45.0, -122.0, 6, 15)?;
    ///
    /// // Filter predictions to only include species likely at this location
    /// let filtered = filter.filter_predictions(&predictions, &location_scores, false);
    /// # Ok(())
    /// # }
    /// ```
    #[must_use]
    pub fn filter_predictions(
        &self,
        predictions: &[Prediction],
        location_scores: &[LocationScore],
        rerank: bool,
    ) -> Vec<Prediction> {
        filter_predictions_impl(predictions, location_scores, self.inner.threshold, rerank)
    }

    /// Filter multiple prediction sets using location scores.
    ///
    /// This is a convenience method for batch processing multiple audio files
    /// from the same location. Predict location scores once, then apply to
    /// multiple prediction sets.
    ///
    /// # Arguments
    /// * `predictions_batch` - Vector of prediction vectors to filter
    /// * `location_scores` - Location scores from `predict()`
    /// * `rerank` - Whether to rerank each prediction set
    ///
    /// # Returns
    /// Vector of filtered prediction vectors
    ///
    /// # Example
    /// ```ignore
    /// let location_scores = range_filter.predict(lat, lon, month, day)?;
    ///
    /// let mut predictions_batch = Vec::new();
    /// for segment in audio_segments {
    ///     let result = classifier.predict(&segment)?;
    ///     predictions_batch.push(result.predictions);
    /// }
    ///
    /// let filtered_batch = range_filter.filter_batch_predictions(
    ///     predictions_batch,
    ///     &location_scores,
    ///     rerank,
    /// );
    /// ```
    #[must_use]
    pub fn filter_batch_predictions(
        &self,
        predictions_batch: Vec<Vec<crate::types::Prediction>>,
        location_scores: &[LocationScore],
        rerank: bool,
    ) -> Vec<Vec<crate::types::Prediction>> {
        filter_batch_predictions_impl(
            predictions_batch,
            location_scores,
            self.inner.threshold,
            rerank,
        )
    }
}

#[cfg(test)]
mod tests {
    #![allow(clippy::unwrap_used)]
    #![allow(clippy::float_cmp)]
    use super::*;

    #[test]
    fn test_calculate_week_january_first() {
        // January 1st = month 1, day 1
        // weeksFromMonths = (1 - 1) * 4 = 0
        // weekInMonth = (1 - 1) / 7 + 1 = 1
        // week = 0 + 1 = 1
        let week = calculate_week(1, 1);
        assert_eq!(week, 1.0);
    }

    #[test]
    fn test_calculate_week_january_eighth() {
        // January 8th = month 1, day 8
        // weeksFromMonths = 0
        // weekInMonth = (8 - 1) / 7 + 1 = 2
        // week = 0 + 2 = 2
        let week = calculate_week(1, 8);
        assert_eq!(week, 2.0);
    }

    #[test]
    fn test_calculate_week_february_first() {
        // February 1st = month 2, day 1
        // weeksFromMonths = (2 - 1) * 4 = 4
        // weekInMonth = (1 - 1) / 7 + 1 = 1
        // week = 4 + 1 = 5
        let week = calculate_week(2, 1);
        assert_eq!(week, 5.0);
    }

    #[test]
    fn test_calculate_week_december_last() {
        // December 31st = month 12, day 31
        // weeksFromMonths = (12 - 1) * 4 = 44
        // weekInMonth = (31 - 1) / 7 + 1 = 5
        // week = 44 + 5 = 49
        // Note: BirdNET uses 48-week year, but days 29-31 can exceed 48
        let week = calculate_week(12, 31);
        assert_eq!(week, 49.0);
    }

    #[test]
    fn test_validate_coordinates_valid() {
        assert!(validate_coordinates(45.0, -122.0).is_ok());
        assert!(validate_coordinates(0.0, 0.0).is_ok());
        assert!(validate_coordinates(-90.0, -180.0).is_ok());
        assert!(validate_coordinates(90.0, 180.0).is_ok());
    }

    #[test]
    fn test_validate_coordinates_invalid_latitude() {
        let result = validate_coordinates(91.0, 0.0);
        assert!(result.is_err());
        assert!(matches!(
            result.unwrap_err(),
            Error::InvalidCoordinates { .. }
        ));
    }

    #[test]
    fn test_validate_coordinates_invalid_longitude() {
        let result = validate_coordinates(0.0, 181.0);
        assert!(result.is_err());
    }

    #[test]
    fn test_validate_date_valid() {
        assert!(validate_date(1, 1).is_ok());
        assert!(validate_date(6, 15).is_ok());
        assert!(validate_date(12, 31).is_ok());
    }

    #[test]
    fn test_validate_date_invalid_month_zero() {
        let result = validate_date(0, 1);
        assert!(result.is_err());
        assert!(matches!(result.unwrap_err(), Error::InvalidDate { .. }));
    }

    #[test]
    fn test_validate_date_invalid_month_thirteen() {
        let result = validate_date(13, 1);
        assert!(result.is_err());
        assert!(matches!(result.unwrap_err(), Error::InvalidDate { .. }));
    }

    #[test]
    fn test_validate_date_invalid_day_zero() {
        let result = validate_date(1, 0);
        assert!(result.is_err());
        assert!(matches!(result.unwrap_err(), Error::InvalidDate { .. }));
    }

    #[test]
    fn test_validate_date_invalid_day_thirty_two() {
        let result = validate_date(1, 32);
        assert!(result.is_err());
        assert!(matches!(result.unwrap_err(), Error::InvalidDate { .. }));
    }

    #[test]
    fn test_range_filter_builder_missing_model_path() {
        let result = RangeFilter::builder().build();
        assert!(result.is_err());
        assert!(matches!(result.unwrap_err(), Error::ModelPathRequired));
    }

    #[test]
    fn test_range_filter_builder_missing_labels() {
        let result = RangeFilter::builder().model_path("/tmp/model.onnx").build();
        assert!(result.is_err());
        assert!(matches!(result.unwrap_err(), Error::LabelsRequired));
    }

    #[test]
    fn test_filter_predictions_above_threshold() {
        // Setup test data
        let predictions = vec![
            Prediction {
                species: "Species A".to_string(),
                confidence: 0.8,
                index: 0,
            },
            Prediction {
                species: "Species B".to_string(),
                confidence: 0.3,
                index: 1,
            },
            Prediction {
                species: "Species C".to_string(),
                confidence: 0.05,
                index: 2,
            },
        ];

        let location_scores = vec![
            LocationScore {
                species: "Species A".to_string(),
                score: 0.9,
                index: 0,
            },
            LocationScore {
                species: "Species B".to_string(),
                score: 0.02,
                index: 1,
            },
            LocationScore {
                species: "Species C".to_string(),
                score: 0.5,
                index: 2,
            },
        ];

        let threshold = 0.03;
        let rerank = false;

        // Call filter_predictions_impl (will fail - not implemented yet)
        let filtered = filter_predictions_impl(&predictions, &location_scores, threshold, rerank);

        // Species B should be filtered out (score 0.02 < threshold 0.03)
        assert_eq!(filtered.len(), 2);
        assert_eq!(filtered[0].species, "Species A");
        assert_eq!(filtered[1].species, "Species C");
    }

    #[test]
    fn test_filter_predictions_with_rerank() {
        // Setup test data with different confidence and location scores
        let predictions = vec![
            Prediction {
                species: "Species A".to_string(),
                confidence: 0.9, // High confidence
                index: 0,
            },
            Prediction {
                species: "Species B".to_string(),
                confidence: 0.8, // Medium-high confidence
                index: 1,
            },
            Prediction {
                species: "Species C".to_string(),
                confidence: 0.7, // Medium confidence
                index: 2,
            },
        ];

        let location_scores = vec![
            LocationScore {
                species: "Species A".to_string(),
                score: 0.5, // Medium location score
                index: 0,
            },
            LocationScore {
                species: "Species B".to_string(),
                score: 0.9, // High location score
                index: 1,
            },
            LocationScore {
                species: "Species C".to_string(),
                score: 0.6, // Medium location score
                index: 2,
            },
        ];

        let threshold = 0.03;
        let rerank = true;

        let filtered = filter_predictions_impl(&predictions, &location_scores, threshold, rerank);

        // All should pass threshold
        assert_eq!(filtered.len(), 3);

        // After reranking (confidence * location_score):
        // Species A: 0.9 * 0.5 = 0.45
        // Species B: 0.8 * 0.9 = 0.72 (highest)
        // Species C: 0.7 * 0.6 = 0.42
        // Should be sorted: B, A, C
        assert_eq!(filtered[0].species, "Species B");
        assert_eq!(filtered[1].species, "Species A");
        assert_eq!(filtered[2].species, "Species C");

        // Verify reranked scores
        assert!((filtered[0].confidence - 0.72).abs() < 0.001);
        assert!((filtered[1].confidence - 0.45).abs() < 0.001);
        assert!((filtered[2].confidence - 0.42).abs() < 0.001);
    }

    #[test]
    fn test_filter_predictions_species_not_in_meta_model() {
        // Setup test data where some predictions are not in meta model
        let predictions = vec![
            Prediction {
                species: "Species A".to_string(),
                confidence: 0.8,
                index: 0,
            },
            Prediction {
                species: "Species B".to_string(),
                confidence: 0.7,
                index: 1,
            },
            Prediction {
                species: "Species D".to_string(), // Not in meta model
                confidence: 0.9,
                index: 3,
            },
        ];

        let location_scores = vec![
            LocationScore {
                species: "Species A".to_string(),
                score: 0.9,
                index: 0,
            },
            LocationScore {
                species: "Species C".to_string(), // Not in predictions
                score: 0.8,
                index: 2,
            },
        ];

        let threshold = 0.03;
        let rerank = false;

        let filtered = filter_predictions_impl(&predictions, &location_scores, threshold, rerank);

        // Species A (in meta, score >= threshold): KEEP
        // Species B (NOT in meta model): KEEP unchanged
        // Species D (NOT in meta model): KEEP unchanged
        assert_eq!(filtered.len(), 3);
        assert_eq!(filtered[0].species, "Species A");
        assert_eq!(filtered[0].confidence, 0.8);
        assert_eq!(filtered[1].species, "Species B");
        assert_eq!(filtered[1].confidence, 0.7);
        assert_eq!(filtered[2].species, "Species D");
        assert_eq!(filtered[2].confidence, 0.9);
    }

    #[test]
    fn test_builder_from_classifier_labels() {
        // This test verifies the builder can accept a label reference
        // We can't test with a real Classifier without a model file,
        // so we test the builder configuration

        let labels = vec!["Species A".to_string(), "Species B".to_string()];
        let builder = RangeFilterBuilder::new().from_classifier_labels(&labels);

        // Verify labels were set (we'll need to expose this for testing)
        assert!(matches!(builder.labels, Some(Labels::InMemory(_))));
    }

    #[test]
    fn test_filter_batch_predictions() {
        use crate::types::Prediction;

        let batch1 = vec![Prediction {
            species: "Species A".to_string(),
            confidence: 0.8,
            index: 0,
        }];
        let batch2 = vec![Prediction {
            species: "Species B".to_string(),
            confidence: 0.6,
            index: 1,
        }];

        let predictions_batch = vec![batch1, batch2];

        let location_scores = vec![
            LocationScore {
                species: "Species A".to_string(),
                score: 0.9,
                index: 0,
            },
            LocationScore {
                species: "Species B".to_string(),
                score: 0.05,
                index: 1,
            },
        ];

        let threshold = 0.1;
        let results =
            filter_batch_predictions_impl(predictions_batch, &location_scores, threshold, false);

        assert_eq!(results.len(), 2);
        assert_eq!(results[0].len(), 1); // Species A kept
        assert_eq!(results[1].len(), 0); // Species B filtered
    }
}