Skip to main content

ipfrs_semantic/
cross_modal_reranker.rs

1//! Cross-Modal Reranker — fuses text (BM25) and vector similarity signals
2//! to produce a single ranked list from multi-modal retrieval candidates.
3//!
4//! # Overview
5//!
6//! [`CrossModalReranker`] is the primary entry point.  Call
7//! [`CrossModalReranker::rerank`] with a candidate list and an optional query
8//! text / query embedding.  The ranker:
9//!
10//! 1. Computes BM25 features against each candidate's `text_snippet`.
11//! 2. Computes cosine / dot-product / L2 features against each candidate's
12//!    `embedding`.
13//! 3. Fuses the per-modality scores with the configured [`CmrFusionStrategy`].
14//! 4. Optionally normalises scores to `[0, 1]`, filters by
15//!    `min_score_threshold`, and keeps only the top-k results.
16//!
17//! # Example
18//!
19//! ```rust
20//! use ipfrs_semantic::cross_modal_reranker::{
21//!     CrossModalReranker, CmrFusionStrategy, RerankerConfig, RerankerCandidate,
22//! };
23//!
24//! let mut reranker = CrossModalReranker::new(RerankerConfig::default());
25//!
26//! let candidates = vec![
27//!     RerankerCandidate::new("doc1", Some("rust systems programming"), None),
28//!     RerankerCandidate::new("doc2", Some("python machine learning"), None),
29//! ];
30//!
31//! let results = reranker.rerank(candidates, Some("rust"), None).unwrap();
32//! assert!(!results.is_empty());
33//! ```
34
35use std::collections::HashMap;
36use std::fmt;
37
38// ─────────────────────────────────────────────────────────────────────────────
39// Error
40// ─────────────────────────────────────────────────────────────────────────────
41
42/// Errors produced by [`CrossModalReranker`].
43#[derive(Debug, Clone, PartialEq)]
44pub enum RerankerError {
45    /// No candidates were supplied to [`CrossModalReranker::rerank`].
46    NoCandidates,
47    /// Query and candidate embedding dimension mismatch.
48    IncompatibleDimensions {
49        /// Expected dimension (query).
50        expected: usize,
51        /// Dimension received (candidate).
52        got: usize,
53    },
54    /// A weight value is outside the valid range `[0, ∞)`.
55    InvalidWeight(f64),
56    /// General configuration problem.
57    ConfigurationError(String),
58}
59
60impl fmt::Display for RerankerError {
61    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
62        match self {
63            RerankerError::NoCandidates => {
64                write!(f, "no candidates provided for reranking")
65            }
66            RerankerError::IncompatibleDimensions { expected, got } => {
67                write!(
68                    f,
69                    "embedding dimension mismatch: expected {expected}, got {got}"
70                )
71            }
72            RerankerError::InvalidWeight(w) => {
73                write!(f, "invalid weight value: {w}")
74            }
75            RerankerError::ConfigurationError(msg) => {
76                write!(f, "configuration error: {msg}")
77            }
78        }
79    }
80}
81
82impl std::error::Error for RerankerError {}
83
84// ─────────────────────────────────────────────────────────────────────────────
85// Score per modality
86// ─────────────────────────────────────────────────────────────────────────────
87
88/// Score contribution from a single retrieval modality.
89#[derive(Debug, Clone)]
90pub struct ModalityScore {
91    /// Name of the modality (e.g. `"text"`, `"vector"`).
92    pub modality: String,
93    /// Raw, un-normalised score.
94    pub raw_score: f64,
95    /// Score after min-max normalisation across all candidates.
96    pub normalized_score: f64,
97    /// Weight applied to this modality during fusion.
98    pub weight: f64,
99}
100
101impl ModalityScore {
102    /// Create a new [`ModalityScore`] with `normalized_score` equal to `raw_score`.
103    pub fn new(modality: impl Into<String>, raw_score: f64, weight: f64) -> Self {
104        Self {
105            modality: modality.into(),
106            raw_score,
107            normalized_score: raw_score,
108            weight,
109        }
110    }
111}
112
113// ─────────────────────────────────────────────────────────────────────────────
114// Candidate
115// ─────────────────────────────────────────────────────────────────────────────
116
117/// A single document / chunk that can be reranked.
118#[derive(Debug, Clone)]
119pub struct RerankerCandidate {
120    /// Unique identifier.
121    pub id: String,
122    /// Optional text used for BM25 scoring.
123    pub text_snippet: Option<String>,
124    /// Optional dense embedding used for vector scoring.
125    pub embedding: Option<Vec<f64>>,
126    /// Per-modality breakdown (populated after `rerank`).
127    pub modality_scores: Vec<ModalityScore>,
128    /// Fused final score (populated after `rerank`).
129    pub final_score: f64,
130    /// 1-based rank in the result list (populated after `rerank`).
131    pub rank: usize,
132}
133
134impl RerankerCandidate {
135    /// Convenience constructor.
136    pub fn new(
137        id: impl Into<String>,
138        text_snippet: Option<&str>,
139        embedding: Option<Vec<f64>>,
140    ) -> Self {
141        Self {
142            id: id.into(),
143            text_snippet: text_snippet.map(str::to_owned),
144            embedding,
145            modality_scores: Vec::new(),
146            final_score: 0.0,
147            rank: 0,
148        }
149    }
150}
151
152// ─────────────────────────────────────────────────────────────────────────────
153// BM25 / text features
154// ─────────────────────────────────────────────────────────────────────────────
155
156/// Features derived from BM25-style keyword matching.
157#[derive(Debug, Clone)]
158pub struct TextFeatures {
159    /// Per-term `(term, tf_idf_contribution)` pairs.
160    pub term_frequency: Vec<(String, f64)>,
161    /// Aggregate BM25 score.
162    pub bm25_score: f64,
163    /// `+0.5` bonus when the whole query is an exact sub-string of the text.
164    pub exact_match_bonus: f64,
165    /// Length penalty ∈ `(-∞, 1.0]`; penalises documents longer than average.
166    pub length_penalty: f64,
167}
168
169// ─────────────────────────────────────────────────────────────────────────────
170// Vector features
171// ─────────────────────────────────────────────────────────────────────────────
172
173/// Features derived from dense vector comparison.
174#[derive(Debug, Clone)]
175pub struct VectorFeatures {
176    /// Cosine similarity ∈ `[-1, 1]`.
177    pub cosine_similarity: f64,
178    /// Raw dot product.
179    pub dot_product: f64,
180    /// Euclidean (L2) distance ∈ `[0, ∞)`.
181    pub l2_distance: f64,
182    /// `1 / (1 + l2_distance)` — bounded similarity in `(0, 1]`.
183    pub euclidean_normalized: f64,
184}
185
186// ─────────────────────────────────────────────────────────────────────────────
187// Fusion strategy
188// ─────────────────────────────────────────────────────────────────────────────
189
190/// How scores from different modalities are combined into a single ranking.
191///
192/// Note: This type is named `CmrFusionStrategy` in the public crate API to
193/// avoid conflicting with `FusionStrategy` already exported from
194/// `multimodal_search`.  Within this module it is also re-exported as the
195/// canonical name.
196#[derive(Debug, Clone)]
197pub enum CmrFusionStrategy {
198    /// Weighted linear combination: `Σ weight_i * score_i`.
199    /// The `Vec` contains `(modality_name, weight)` pairs.
200    LinearCombination(Vec<(String, f64)>),
201    /// Reciprocal Rank Fusion: `score = Σ 1 / (k + rank_i)`.
202    ReciprocalRankFusion(f64),
203    /// Borda count — rank-based voting.
204    Borda,
205    /// Keep the maximum individual modality score.
206    MaxScore,
207    /// Pre-trained / supplied weight vector applied to ordered modality scores.
208    LearnedWeights(Vec<f64>),
209}
210
211impl Default for CmrFusionStrategy {
212    fn default() -> Self {
213        CmrFusionStrategy::LinearCombination(vec![
214            ("text".to_string(), 0.4),
215            ("vector".to_string(), 0.6),
216        ])
217    }
218}
219
220// ─────────────────────────────────────────────────────────────────────────────
221// Config
222// ─────────────────────────────────────────────────────────────────────────────
223
224/// Configuration for [`CrossModalReranker`].
225#[derive(Debug, Clone)]
226pub struct RerankerConfig {
227    /// Score fusion strategy.
228    pub fusion_strategy: CmrFusionStrategy,
229    /// Global text-modality weight (used when `fusion_strategy` is not
230    /// `LinearCombination`).
231    pub text_weight: f64,
232    /// Global vector-modality weight.
233    pub vector_weight: f64,
234    /// BM25 term-frequency saturation constant (typical: 1.2–2.0).
235    pub bm25_k1: f64,
236    /// BM25 length normalisation constant (typical: 0.75).
237    pub bm25_b: f64,
238    /// Normalise final scores to `[0, 1]`.
239    pub normalize_scores: bool,
240    /// Discard candidates whose final score is below this threshold.
241    pub min_score_threshold: f64,
242    /// Maximum number of results to return.  `0` means no limit.
243    pub top_k: usize,
244}
245
246impl Default for RerankerConfig {
247    fn default() -> Self {
248        Self {
249            fusion_strategy: CmrFusionStrategy::default(),
250            text_weight: 0.4,
251            vector_weight: 0.6,
252            bm25_k1: 1.5,
253            bm25_b: 0.75,
254            normalize_scores: true,
255            min_score_threshold: 0.0,
256            top_k: 100,
257        }
258    }
259}
260
261impl RerankerConfig {
262    /// Validate all weight fields, returning an error on invalid values.
263    fn validate(&self) -> Result<(), RerankerError> {
264        if self.text_weight < 0.0 || self.text_weight.is_nan() {
265            return Err(RerankerError::InvalidWeight(self.text_weight));
266        }
267        if self.vector_weight < 0.0 || self.vector_weight.is_nan() {
268            return Err(RerankerError::InvalidWeight(self.vector_weight));
269        }
270        if self.bm25_k1 < 0.0 || self.bm25_k1.is_nan() {
271            return Err(RerankerError::ConfigurationError(
272                "bm25_k1 must be non-negative".to_string(),
273            ));
274        }
275        if !(0.0..=1.0).contains(&self.bm25_b) {
276            return Err(RerankerError::ConfigurationError(
277                "bm25_b must be in [0, 1]".to_string(),
278            ));
279        }
280        // Validate LinearCombination weights
281        if let CmrFusionStrategy::LinearCombination(ref pairs) = self.fusion_strategy {
282            for (_, w) in pairs {
283                if *w < 0.0 || w.is_nan() {
284                    return Err(RerankerError::InvalidWeight(*w));
285                }
286            }
287        }
288        if let CmrFusionStrategy::ReciprocalRankFusion(k) = self.fusion_strategy {
289            if k <= 0.0 || k.is_nan() {
290                return Err(RerankerError::ConfigurationError(
291                    "RRF k must be positive".to_string(),
292                ));
293            }
294        }
295        Ok(())
296    }
297}
298
299// ─────────────────────────────────────────────────────────────────────────────
300// Stats
301// ─────────────────────────────────────────────────────────────────────────────
302
303/// Operational statistics for [`CrossModalReranker`].
304#[derive(Debug, Clone, Default)]
305pub struct RerankerStats {
306    /// Total candidates processed across all `rerank` calls.
307    pub candidates_reranked: u64,
308    /// Average displacement of rank position (abs(new_rank - old_rank)).
309    pub avg_rank_displacement: f64,
310    /// Distinct modality names observed.
311    pub modalities_used: Vec<String>,
312    /// Number of fusion operations performed.
313    pub fusion_calls: u64,
314}
315
316// ─────────────────────────────────────────────────────────────────────────────
317// Tokenizer
318// ─────────────────────────────────────────────────────────────────────────────
319
320/// Simple whitespace tokenizer with alphabetic filtering.
321/// Uses the same semantics as the FNV-1a-compatible form in the spec.
322fn tokenize(text: &str) -> Vec<String> {
323    text.split_whitespace()
324        .map(|w| {
325            w.to_lowercase()
326                .trim_matches(|c: char| !c.is_alphabetic())
327                .to_string()
328        })
329        .filter(|w| !w.is_empty())
330        .collect()
331}
332
333// ─────────────────────────────────────────────────────────────────────────────
334// Pure PRNG (for tests, not used in production logic)
335// ─────────────────────────────────────────────────────────────────────────────
336
337/// XorShift-64 PRNG — integer step.
338#[allow(dead_code)]
339fn xorshift64(state: &mut u64) -> u64 {
340    let mut x = *state;
341    x ^= x << 13;
342    x ^= x >> 7;
343    x ^= x << 17;
344    *state = x;
345    x
346}
347
348/// XorShift-64 PRNG — float step returning a value in `[0, 1)`.
349#[allow(dead_code)]
350fn xorshift_f64(state: &mut u64) -> f64 {
351    (xorshift64(state) >> 11) as f64 / (1u64 << 53) as f64
352}
353
354// ─────────────────────────────────────────────────────────────────────────────
355// CrossModalReranker
356// ─────────────────────────────────────────────────────────────────────────────
357
358/// Cross-modal reranker that fuses BM25 text scores with dense vector scores.
359pub struct CrossModalReranker {
360    config: RerankerConfig,
361    stats: RerankerStats,
362}
363
364impl CrossModalReranker {
365    /// Create a new reranker with the supplied configuration.
366    pub fn new(config: RerankerConfig) -> Self {
367        Self {
368            config,
369            stats: RerankerStats::default(),
370        }
371    }
372
373    /// Replace the current configuration.
374    pub fn update_config(&mut self, config: RerankerConfig) {
375        self.config = config;
376    }
377
378    /// Return a snapshot of the current operational statistics.
379    pub fn stats(&self) -> RerankerStats {
380        self.stats.clone()
381    }
382
383    // ──────────────────────────────────────────────────────────────────────
384    // BM25 text features
385    // ──────────────────────────────────────────────────────────────────────
386
387    /// Compute BM25-derived text features for a single `(query, text)` pair.
388    ///
389    /// `avg_doc_len` is the average document token count across the corpus;
390    /// pass `0.0` (or any non-positive value) to fall back to `1.0`.
391    pub fn compute_text_features(&self, query: &str, text: &str, avg_doc_len: f64) -> TextFeatures {
392        let avg_doc_len = if avg_doc_len > 0.0 { avg_doc_len } else { 1.0 };
393
394        let query_tokens = tokenize(query);
395        let doc_tokens = tokenize(text);
396        let doc_len = doc_tokens.len() as f64;
397
398        // Term frequency map for the document
399        let mut tf_map: HashMap<String, f64> = HashMap::new();
400        for tok in &doc_tokens {
401            *tf_map.entry(tok.clone()).or_insert(0.0) += 1.0;
402        }
403
404        let k1 = self.config.bm25_k1;
405        let b = self.config.bm25_b;
406
407        // N = 1 (single-document estimate)
408        // IDF = ln((N - df + 0.5) / (df + 0.5) + 1) where df = 1 if term present, 0 otherwise
409        let n = 1.0_f64;
410        let mut term_contributions: Vec<(String, f64)> = Vec::new();
411        let mut bm25_total = 0.0_f64;
412
413        for qt in &query_tokens {
414            let freq = tf_map.get(qt).copied().unwrap_or(0.0);
415            let df = if freq > 0.0 { 1.0 } else { 0.0 };
416
417            let idf = ((n - df + 0.5) / (df + 0.5) + 1.0).ln();
418            let tf_norm = (freq * (k1 + 1.0)) / (freq + k1 * (1.0 - b + b * doc_len / avg_doc_len));
419
420            let contribution = idf * tf_norm;
421            bm25_total += contribution;
422            term_contributions.push((qt.clone(), contribution));
423        }
424
425        let exact_match_bonus =
426            if !query.is_empty() && text.to_lowercase().contains(&query.to_lowercase()) {
427                0.5
428            } else {
429                0.0
430            };
431
432        let length_penalty = 1.0 - 0.1 * ((doc_len / avg_doc_len) - 1.0).max(0.0);
433
434        TextFeatures {
435            term_frequency: term_contributions,
436            bm25_score: bm25_total,
437            exact_match_bonus,
438            length_penalty,
439        }
440    }
441
442    // ──────────────────────────────────────────────────────────────────────
443    // Vector features
444    // ──────────────────────────────────────────────────────────────────────
445
446    /// Compute vector-space features between a query embedding and a candidate
447    /// embedding.
448    ///
449    /// Returns [`RerankerError::IncompatibleDimensions`] if the slices have
450    /// different lengths.
451    pub fn compute_vector_features(
452        query: &[f64],
453        candidate: &[f64],
454    ) -> Result<VectorFeatures, RerankerError> {
455        if query.len() != candidate.len() {
456            return Err(RerankerError::IncompatibleDimensions {
457                expected: query.len(),
458                got: candidate.len(),
459            });
460        }
461
462        let mut dot = 0.0_f64;
463        let mut norm_q = 0.0_f64;
464        let mut norm_c = 0.0_f64;
465        let mut sq_diff = 0.0_f64;
466
467        for (q, c) in query.iter().zip(candidate.iter()) {
468            dot += q * c;
469            norm_q += q * q;
470            norm_c += c * c;
471            let d = q - c;
472            sq_diff += d * d;
473        }
474
475        let norm_q = norm_q.sqrt();
476        let norm_c = norm_c.sqrt();
477        let denom = norm_q * norm_c;
478
479        let cosine_similarity = if denom > 0.0 { dot / denom } else { 0.0 };
480        let l2_distance = sq_diff.sqrt();
481        let euclidean_normalized = 1.0 / (1.0 + l2_distance);
482
483        Ok(VectorFeatures {
484            cosine_similarity,
485            dot_product: dot,
486            l2_distance,
487            euclidean_normalized,
488        })
489    }
490
491    // ──────────────────────────────────────────────────────────────────────
492    // Reciprocal Rank Fusion
493    // ──────────────────────────────────────────────────────────────────────
494
495    /// Merge multiple ranked lists using Reciprocal Rank Fusion.
496    ///
497    /// `rank_lists` is a list of ranked ID lists (index 0 = rank 1).
498    /// Returns `(id, rrf_score)` pairs sorted by descending score.
499    pub fn reciprocal_rank_fusion(rank_lists: Vec<Vec<String>>, k: f64) -> Vec<(String, f64)> {
500        let mut scores: HashMap<String, f64> = HashMap::new();
501
502        for list in &rank_lists {
503            for (rank_zero_based, id) in list.iter().enumerate() {
504                let rank = (rank_zero_based + 1) as f64;
505                *scores.entry(id.clone()).or_insert(0.0) += 1.0 / (k + rank);
506            }
507        }
508
509        let mut result: Vec<(String, f64)> = scores.into_iter().collect();
510        result.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
511        result
512    }
513
514    // ──────────────────────────────────────────────────────────────────────
515    // Main rerank entry point
516    // ──────────────────────────────────────────────────────────────────────
517
518    /// Rerank the given candidates by fusing text and/or vector similarity
519    /// with the configured [`CmrFusionStrategy`].
520    ///
521    /// * `query_text`      — used to compute BM25 features (optional).
522    /// * `query_embedding` — used to compute vector features (optional).
523    ///
524    /// The returned list is sorted descending by `final_score`, filtered
525    /// by `min_score_threshold`, and limited to `top_k` entries.
526    pub fn rerank(
527        &mut self,
528        mut candidates: Vec<RerankerCandidate>,
529        query_text: Option<&str>,
530        query_embedding: Option<&[f64]>,
531    ) -> Result<Vec<RerankerCandidate>, RerankerError> {
532        if candidates.is_empty() {
533            return Err(RerankerError::NoCandidates);
534        }
535
536        self.config.validate()?;
537
538        // Validate embedding dimensions up-front
539        if let Some(qe) = query_embedding {
540            for c in &candidates {
541                if let Some(ce) = &c.embedding {
542                    if ce.len() != qe.len() {
543                        return Err(RerankerError::IncompatibleDimensions {
544                            expected: qe.len(),
545                            got: ce.len(),
546                        });
547                    }
548                }
549            }
550        }
551
552        // Compute average document length for BM25
553        let avg_doc_len = {
554            let texts: Vec<usize> = candidates
555                .iter()
556                .filter_map(|c| c.text_snippet.as_ref())
557                .map(|t| tokenize(t).len())
558                .collect();
559            if texts.is_empty() {
560                1.0
561            } else {
562                texts.iter().sum::<usize>() as f64 / texts.len() as f64
563            }
564        };
565
566        // Assign modality scores to each candidate
567        for cand in candidates.iter_mut() {
568            cand.modality_scores.clear();
569
570            // ── Text ──
571            if let (Some(qt), Some(snippet)) = (query_text, cand.text_snippet.as_deref()) {
572                let tf = self.compute_text_features(qt, snippet, avg_doc_len);
573                let text_score = (tf.bm25_score + tf.exact_match_bonus) * tf.length_penalty;
574                cand.modality_scores.push(ModalityScore::new(
575                    "text",
576                    text_score,
577                    self.config.text_weight,
578                ));
579            }
580
581            // ── Vector ──
582            if let (Some(qe), Some(ce)) = (query_embedding, cand.embedding.as_deref()) {
583                // Dimension already validated above
584                let vf = Self::compute_vector_features(qe, ce)?;
585                cand.modality_scores.push(ModalityScore::new(
586                    "vector",
587                    vf.cosine_similarity,
588                    self.config.vector_weight,
589                ));
590            }
591        }
592
593        // Fuse scores
594        self.apply_fusion(&mut candidates)?;
595
596        // Sort descending
597        candidates.sort_by(|a, b| {
598            b.final_score
599                .partial_cmp(&a.final_score)
600                .unwrap_or(std::cmp::Ordering::Equal)
601        });
602
603        // Normalise if requested (before filtering / truncation)
604        if self.config.normalize_scores {
605            Self::normalize_scores(&mut candidates);
606        }
607
608        // Assign pre-filter ranks to compute displacement later
609        let original_ranks: Vec<(String, usize)> = candidates
610            .iter()
611            .enumerate()
612            .map(|(i, c)| (c.id.clone(), i + 1))
613            .collect();
614
615        // Filter by min_score_threshold
616        candidates.retain(|c| c.final_score >= self.config.min_score_threshold);
617
618        // Limit to top_k
619        if self.config.top_k > 0 && candidates.len() > self.config.top_k {
620            candidates.truncate(self.config.top_k);
621        }
622
623        // Assign final 1-based ranks
624        for (i, c) in candidates.iter_mut().enumerate() {
625            c.rank = i + 1;
626        }
627
628        // Update stats
629        let total = candidates.len() as u64;
630        let displacement: f64 = candidates
631            .iter()
632            .map(|c| {
633                original_ranks
634                    .iter()
635                    .find(|(id, _)| id == &c.id)
636                    .map(|(_, orig)| (c.rank as i64 - *orig as i64).unsigned_abs() as f64)
637                    .unwrap_or(0.0)
638            })
639            .sum::<f64>()
640            / total.max(1) as f64;
641
642        self.stats.candidates_reranked += total;
643        self.stats.fusion_calls += 1;
644        // Rolling average of displacement
645        if self.stats.fusion_calls == 1 {
646            self.stats.avg_rank_displacement = displacement;
647        } else {
648            let n = self.stats.fusion_calls as f64;
649            self.stats.avg_rank_displacement =
650                (self.stats.avg_rank_displacement * (n - 1.0) + displacement) / n;
651        }
652
653        // Track modalities
654        for c in &candidates {
655            for ms in &c.modality_scores {
656                if !self.stats.modalities_used.contains(&ms.modality) {
657                    self.stats.modalities_used.push(ms.modality.clone());
658                }
659            }
660        }
661
662        Ok(candidates)
663    }
664
665    // ──────────────────────────────────────────────────────────────────────
666    // Internal: per-strategy fusion
667    // ──────────────────────────────────────────────────────────────────────
668
669    fn apply_fusion(&self, candidates: &mut [RerankerCandidate]) -> Result<(), RerankerError> {
670        match &self.config.fusion_strategy {
671            CmrFusionStrategy::LinearCombination(pairs) => {
672                let weight_map: HashMap<&str, f64> =
673                    pairs.iter().map(|(k, v)| (k.as_str(), *v)).collect();
674
675                for cand in candidates.iter_mut() {
676                    let score: f64 = cand
677                        .modality_scores
678                        .iter()
679                        .map(|ms| {
680                            let w = weight_map
681                                .get(ms.modality.as_str())
682                                .copied()
683                                .unwrap_or(ms.weight);
684                            w * ms.raw_score
685                        })
686                        .sum();
687                    cand.final_score = score;
688                }
689            }
690
691            CmrFusionStrategy::ReciprocalRankFusion(k) => {
692                let k = *k;
693                // Build per-modality rank lists
694                let modality_names: Vec<String> = {
695                    let mut names: Vec<String> = Vec::new();
696                    for c in candidates.iter() {
697                        for ms in &c.modality_scores {
698                            if !names.contains(&ms.modality) {
699                                names.push(ms.modality.clone());
700                            }
701                        }
702                    }
703                    names
704                };
705
706                // Per-modality sorted ID lists (by raw_score desc)
707                let rank_lists: Vec<Vec<String>> = modality_names
708                    .iter()
709                    .map(|m| {
710                        let mut scored: Vec<(String, f64)> = candidates
711                            .iter()
712                            .filter_map(|c| {
713                                c.modality_scores
714                                    .iter()
715                                    .find(|ms| &ms.modality == m)
716                                    .map(|ms| (c.id.clone(), ms.raw_score))
717                            })
718                            .collect();
719                        scored.sort_by(|a, b| {
720                            b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)
721                        });
722                        scored.into_iter().map(|(id, _)| id).collect()
723                    })
724                    .collect();
725
726                let rrf_scores = Self::reciprocal_rank_fusion(rank_lists, k);
727                let score_map: HashMap<&str, f64> =
728                    rrf_scores.iter().map(|(id, s)| (id.as_str(), *s)).collect();
729
730                for cand in candidates.iter_mut() {
731                    cand.final_score = score_map.get(cand.id.as_str()).copied().unwrap_or(0.0);
732                }
733            }
734
735            CmrFusionStrategy::Borda => {
736                // Borda count: for each modality rank list, candidate at rank r
737                // receives (N - r) points.
738                let n = candidates.len();
739                let modality_names: Vec<String> = {
740                    let mut names: Vec<String> = Vec::new();
741                    for c in candidates.iter() {
742                        for ms in &c.modality_scores {
743                            if !names.contains(&ms.modality) {
744                                names.push(ms.modality.clone());
745                            }
746                        }
747                    }
748                    names
749                };
750
751                let mut borda_totals: HashMap<String, f64> =
752                    candidates.iter().map(|c| (c.id.clone(), 0.0)).collect();
753
754                for m in &modality_names {
755                    let mut scored: Vec<(String, f64)> = candidates
756                        .iter()
757                        .filter_map(|c| {
758                            c.modality_scores
759                                .iter()
760                                .find(|ms| &ms.modality == m)
761                                .map(|ms| (c.id.clone(), ms.raw_score))
762                        })
763                        .collect();
764                    scored
765                        .sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
766
767                    for (rank_zero, (id, _)) in scored.iter().enumerate() {
768                        let points = (n - rank_zero) as f64;
769                        if let Some(total) = borda_totals.get_mut(id) {
770                            *total += points;
771                        }
772                    }
773                }
774
775                for cand in candidates.iter_mut() {
776                    cand.final_score = borda_totals.get(&cand.id).copied().unwrap_or(0.0);
777                }
778            }
779
780            CmrFusionStrategy::MaxScore => {
781                for cand in candidates.iter_mut() {
782                    cand.final_score = cand
783                        .modality_scores
784                        .iter()
785                        .map(|ms| ms.raw_score)
786                        .fold(f64::NEG_INFINITY, f64::max);
787                    if cand.final_score.is_infinite() {
788                        cand.final_score = 0.0;
789                    }
790                }
791            }
792
793            CmrFusionStrategy::LearnedWeights(weights) => {
794                for cand in candidates.iter_mut() {
795                    let score: f64 = cand
796                        .modality_scores
797                        .iter()
798                        .enumerate()
799                        .map(|(i, ms)| {
800                            let w = weights.get(i).copied().unwrap_or(1.0);
801                            w * ms.raw_score
802                        })
803                        .sum();
804                    cand.final_score = score;
805                }
806            }
807        }
808
809        Ok(())
810    }
811
812    // ──────────────────────────────────────────────────────────────────────
813    // Internal: score normalisation
814    // ──────────────────────────────────────────────────────────────────────
815
816    fn normalize_scores(candidates: &mut [RerankerCandidate]) {
817        if candidates.is_empty() {
818            return;
819        }
820        let min_s = candidates
821            .iter()
822            .map(|c| c.final_score)
823            .fold(f64::INFINITY, f64::min);
824        let max_s = candidates
825            .iter()
826            .map(|c| c.final_score)
827            .fold(f64::NEG_INFINITY, f64::max);
828
829        let range = max_s - min_s;
830        if range < f64::EPSILON {
831            for c in candidates.iter_mut() {
832                c.final_score = 1.0;
833            }
834            return;
835        }
836        for c in candidates.iter_mut() {
837            c.final_score = (c.final_score - min_s) / range;
838        }
839
840        // Also normalise per-modality scores
841        // (collect modality names first to avoid borrow conflicts)
842        let modality_names: Vec<String> = {
843            let mut names: Vec<String> = Vec::new();
844            for c in candidates.iter() {
845                for ms in &c.modality_scores {
846                    if !names.contains(&ms.modality) {
847                        names.push(ms.modality.clone());
848                    }
849                }
850            }
851            names
852        };
853
854        for m in &modality_names {
855            let min_r = candidates
856                .iter()
857                .flat_map(|c| c.modality_scores.iter())
858                .filter(|ms| &ms.modality == m)
859                .map(|ms| ms.raw_score)
860                .fold(f64::INFINITY, f64::min);
861            let max_r = candidates
862                .iter()
863                .flat_map(|c| c.modality_scores.iter())
864                .filter(|ms| &ms.modality == m)
865                .map(|ms| ms.raw_score)
866                .fold(f64::NEG_INFINITY, f64::max);
867
868            let r = max_r - min_r;
869            for c in candidates.iter_mut() {
870                for ms in c.modality_scores.iter_mut() {
871                    if &ms.modality == m {
872                        ms.normalized_score = if r < f64::EPSILON {
873                            1.0
874                        } else {
875                            (ms.raw_score - min_r) / r
876                        };
877                    }
878                }
879            }
880        }
881    }
882}
883
884// ─────────────────────────────────────────────────────────────────────────────
885// Tests
886// ─────────────────────────────────────────────────────────────────────────────
887
888#[cfg(test)]
889mod tests {
890    use super::*;
891
892    // ── helpers ──────────────────────────────────────────────────────────────
893
894    fn make_text_candidate(id: &str, text: &str) -> RerankerCandidate {
895        RerankerCandidate::new(id, Some(text), None)
896    }
897
898    fn make_vec_candidate(id: &str, embedding: Vec<f64>) -> RerankerCandidate {
899        RerankerCandidate::new(id, None, Some(embedding))
900    }
901
902    fn make_full_candidate(id: &str, text: &str, embedding: Vec<f64>) -> RerankerCandidate {
903        RerankerCandidate::new(id, Some(text), Some(embedding))
904    }
905
906    fn default_reranker() -> CrossModalReranker {
907        CrossModalReranker::new(RerankerConfig::default())
908    }
909
910    // ── tokenizer ────────────────────────────────────────────────────────────
911
912    #[test]
913    fn test_tokenize_basic() {
914        let tokens = tokenize("Hello, World!");
915        assert_eq!(tokens, vec!["hello", "world"]);
916    }
917
918    #[test]
919    fn test_tokenize_empty() {
920        assert!(tokenize("").is_empty());
921    }
922
923    #[test]
924    fn test_tokenize_punctuation_stripped() {
925        let tokens = tokenize("rust, systems, programming.");
926        assert_eq!(tokens, vec!["rust", "systems", "programming"]);
927    }
928
929    #[test]
930    fn test_tokenize_lowercase() {
931        let tokens = tokenize("Rust PROGRAMMING");
932        assert!(tokens
933            .iter()
934            .all(|t| t.chars().all(|c| c.is_lowercase() || !c.is_alphabetic())));
935    }
936
937    // ── PRNG ─────────────────────────────────────────────────────────────────
938
939    #[test]
940    fn test_xorshift64_not_zero_after_seed() {
941        let mut state: u64 = 12345;
942        let v = xorshift64(&mut state);
943        assert_ne!(v, 0);
944    }
945
946    #[test]
947    fn test_xorshift_f64_in_range() {
948        let mut state: u64 = 99999;
949        for _ in 0..1000 {
950            let v = xorshift_f64(&mut state);
951            assert!((0.0..1.0).contains(&v), "value out of range: {v}");
952        }
953    }
954
955    #[test]
956    fn test_xorshift_f64_deterministic() {
957        let mut s1: u64 = 42;
958        let mut s2: u64 = 42;
959        assert_eq!(xorshift_f64(&mut s1), xorshift_f64(&mut s2));
960    }
961
962    // ── BM25 text features ────────────────────────────────────────────────────
963
964    #[test]
965    fn test_bm25_empty_query() {
966        let r = default_reranker();
967        let tf = r.compute_text_features("", "some text here", 4.0);
968        assert_eq!(tf.bm25_score, 0.0);
969    }
970
971    #[test]
972    fn test_bm25_empty_document() {
973        let r = default_reranker();
974        let tf = r.compute_text_features("rust", "", 4.0);
975        assert_eq!(tf.bm25_score, 0.0);
976    }
977
978    #[test]
979    fn test_bm25_term_present_vs_absent() {
980        let r = default_reranker();
981        let tf_present = r.compute_text_features("rust", "rust systems", 2.0);
982        let tf_absent = r.compute_text_features("rust", "python systems", 2.0);
983        assert!(tf_present.bm25_score > tf_absent.bm25_score);
984    }
985
986    #[test]
987    fn test_bm25_exact_match_bonus() {
988        let r = default_reranker();
989        let tf_exact =
990            r.compute_text_features("rust programming", "I love rust programming a lot", 5.0);
991        let tf_partial =
992            r.compute_text_features("rust programming", "I love rust and programming", 5.0);
993        assert!(
994            tf_exact.exact_match_bonus > tf_partial.exact_match_bonus,
995            "exact match should have bonus: exact={}, partial={}",
996            tf_exact.exact_match_bonus,
997            tf_partial.exact_match_bonus
998        );
999    }
1000
1001    #[test]
1002    fn test_bm25_exact_match_bonus_value() {
1003        let r = default_reranker();
1004        let tf = r.compute_text_features("hello world", "hello world this is a test", 5.0);
1005        assert!((tf.exact_match_bonus - 0.5).abs() < 1e-10);
1006    }
1007
1008    #[test]
1009    fn test_bm25_no_exact_match_bonus() {
1010        let r = default_reranker();
1011        let tf = r.compute_text_features("hello world", "goodbye everyone", 5.0);
1012        assert_eq!(tf.exact_match_bonus, 0.0);
1013    }
1014
1015    #[test]
1016    fn test_bm25_length_penalty_short_doc() {
1017        let r = default_reranker();
1018        // doc shorter than avg → no penalty
1019        let tf = r.compute_text_features("a", "a", 100.0);
1020        assert!((tf.length_penalty - 1.0).abs() < 1e-10);
1021    }
1022
1023    #[test]
1024    fn test_bm25_length_penalty_long_doc() {
1025        let r = default_reranker();
1026        let long_text = "word ".repeat(100);
1027        let tf = r.compute_text_features("word", long_text.trim(), 10.0);
1028        assert!(tf.length_penalty < 1.0);
1029    }
1030
1031    #[test]
1032    fn test_bm25_term_frequency_populated() {
1033        let r = default_reranker();
1034        let tf = r.compute_text_features("rust python", "rust is great", 3.0);
1035        assert!(!tf.term_frequency.is_empty());
1036    }
1037
1038    #[test]
1039    fn test_bm25_zero_avg_doc_len_fallback() {
1040        let r = default_reranker();
1041        let tf = r.compute_text_features("hello", "hello world", 0.0);
1042        // should not panic, bm25_score should be finite
1043        assert!(tf.bm25_score.is_finite());
1044    }
1045
1046    #[test]
1047    fn test_bm25_custom_k1_b() {
1048        let config = RerankerConfig {
1049            bm25_k1: 2.0,
1050            bm25_b: 0.5,
1051            ..Default::default()
1052        };
1053        let r = CrossModalReranker::new(config);
1054        let tf = r.compute_text_features("rust", "rust systems rust", 3.0);
1055        assert!(tf.bm25_score > 0.0);
1056    }
1057
1058    // ── vector features ───────────────────────────────────────────────────────
1059
1060    #[test]
1061    fn test_vector_features_identical() {
1062        let v = vec![1.0, 0.0, 0.0];
1063        let vf = CrossModalReranker::compute_vector_features(&v, &v)
1064            .expect("test: identical vectors should compute without error");
1065        assert!((vf.cosine_similarity - 1.0).abs() < 1e-10);
1066        assert!(vf.l2_distance.abs() < 1e-10);
1067        assert!((vf.euclidean_normalized - 1.0).abs() < 1e-10);
1068    }
1069
1070    #[test]
1071    fn test_vector_features_orthogonal() {
1072        let q = vec![1.0, 0.0];
1073        let c = vec![0.0, 1.0];
1074        let vf = CrossModalReranker::compute_vector_features(&q, &c)
1075            .expect("test: orthogonal vectors should compute without error");
1076        assert!(vf.cosine_similarity.abs() < 1e-10);
1077    }
1078
1079    #[test]
1080    fn test_vector_features_opposite() {
1081        let q = vec![1.0, 0.0];
1082        let c = vec![-1.0, 0.0];
1083        let vf = CrossModalReranker::compute_vector_features(&q, &c)
1084            .expect("test: opposite vectors should compute without error");
1085        assert!((vf.cosine_similarity + 1.0).abs() < 1e-10);
1086    }
1087
1088    #[test]
1089    fn test_vector_features_dimension_mismatch() {
1090        let q = vec![1.0, 2.0, 3.0];
1091        let c = vec![1.0, 2.0];
1092        let err = CrossModalReranker::compute_vector_features(&q, &c)
1093            .expect_err("test: dimension mismatch should return error");
1094        assert_eq!(
1095            err,
1096            RerankerError::IncompatibleDimensions {
1097                expected: 3,
1098                got: 2
1099            }
1100        );
1101    }
1102
1103    #[test]
1104    fn test_vector_features_zero_vector() {
1105        let q = vec![0.0, 0.0];
1106        let c = vec![1.0, 0.0];
1107        let vf = CrossModalReranker::compute_vector_features(&q, &c)
1108            .expect("test: zero query vector should compute without error");
1109        // zero query → cosine = 0
1110        assert_eq!(vf.cosine_similarity, 0.0);
1111    }
1112
1113    #[test]
1114    fn test_vector_features_dot_product() {
1115        let q = vec![1.0, 2.0, 3.0];
1116        let c = vec![4.0, 5.0, 6.0];
1117        let vf = CrossModalReranker::compute_vector_features(&q, &c)
1118            .expect("test: dot product computation should succeed");
1119        assert!((vf.dot_product - 32.0).abs() < 1e-10);
1120    }
1121
1122    #[test]
1123    fn test_vector_features_l2_distance() {
1124        let q = vec![0.0, 0.0];
1125        let c = vec![3.0, 4.0];
1126        let vf = CrossModalReranker::compute_vector_features(&q, &c)
1127            .expect("test: L2 distance computation should succeed");
1128        assert!((vf.l2_distance - 5.0).abs() < 1e-10);
1129    }
1130
1131    #[test]
1132    fn test_vector_features_euclidean_normalized_bounded() {
1133        let mut state: u64 = 777;
1134        let q: Vec<f64> = (0..8).map(|_| xorshift_f64(&mut state)).collect();
1135        let c: Vec<f64> = (0..8).map(|_| xorshift_f64(&mut state)).collect();
1136        let vf = CrossModalReranker::compute_vector_features(&q, &c)
1137            .expect("test: euclidean_normalized computation should succeed");
1138        assert!(vf.euclidean_normalized > 0.0);
1139        assert!(vf.euclidean_normalized <= 1.0);
1140    }
1141
1142    // ── RRF ───────────────────────────────────────────────────────────────────
1143
1144    #[test]
1145    fn test_rrf_single_list() {
1146        let lists = vec![vec!["a".to_string(), "b".to_string(), "c".to_string()]];
1147        let scores = CrossModalReranker::reciprocal_rank_fusion(lists, 60.0);
1148        // a at rank 1 → 1/61
1149        let a_score = scores
1150            .iter()
1151            .find(|(id, _)| id == "a")
1152            .expect("test: 'a' must be in RRF scores")
1153            .1;
1154        let b_score = scores
1155            .iter()
1156            .find(|(id, _)| id == "b")
1157            .expect("test: 'b' must be in RRF scores")
1158            .1;
1159        assert!(a_score > b_score);
1160    }
1161
1162    #[test]
1163    fn test_rrf_two_lists_consensus() {
1164        let lists = vec![
1165            vec!["a".to_string(), "b".to_string()],
1166            vec!["a".to_string(), "b".to_string()],
1167        ];
1168        let scores = CrossModalReranker::reciprocal_rank_fusion(lists, 60.0);
1169        let a = scores
1170            .iter()
1171            .find(|(id, _)| id == "a")
1172            .expect("test: 'a' must be in RRF scores")
1173            .1;
1174        let b = scores
1175            .iter()
1176            .find(|(id, _)| id == "b")
1177            .expect("test: 'b' must be in RRF scores")
1178            .1;
1179        assert!(a > b);
1180    }
1181
1182    #[test]
1183    fn test_rrf_rank_disagreement() {
1184        let lists = vec![
1185            vec!["a".to_string(), "b".to_string()],
1186            vec!["b".to_string(), "a".to_string()],
1187        ];
1188        let scores = CrossModalReranker::reciprocal_rank_fusion(lists, 60.0);
1189        // Equal ranking → scores should be equal
1190        let a = scores
1191            .iter()
1192            .find(|(id, _)| id == "a")
1193            .expect("test: 'a' must be in RRF scores")
1194            .1;
1195        let b = scores
1196            .iter()
1197            .find(|(id, _)| id == "b")
1198            .expect("test: 'b' must be in RRF scores")
1199            .1;
1200        assert!((a - b).abs() < 1e-10, "a={a}, b={b}");
1201    }
1202
1203    #[test]
1204    fn test_rrf_custom_k() {
1205        let lists = vec![vec!["x".to_string()]];
1206        let s1 = CrossModalReranker::reciprocal_rank_fusion(lists.clone(), 10.0);
1207        let s2 = CrossModalReranker::reciprocal_rank_fusion(lists, 100.0);
1208        // smaller k → larger score
1209        let v1 = s1[0].1;
1210        let v2 = s2[0].1;
1211        assert!(v1 > v2);
1212    }
1213
1214    #[test]
1215    fn test_rrf_empty_lists() {
1216        let scores = CrossModalReranker::reciprocal_rank_fusion(vec![], 60.0);
1217        assert!(scores.is_empty());
1218    }
1219
1220    #[test]
1221    fn test_rrf_sorted_descending() {
1222        let lists = vec![
1223            vec!["c".to_string(), "b".to_string(), "a".to_string()],
1224            vec!["a".to_string(), "b".to_string(), "c".to_string()],
1225        ];
1226        let scores = CrossModalReranker::reciprocal_rank_fusion(lists, 60.0);
1227        for w in scores.windows(2) {
1228            assert!(w[0].1 >= w[1].1);
1229        }
1230    }
1231
1232    // ── text-only rerank ──────────────────────────────────────────────────────
1233
1234    #[test]
1235    fn test_text_only_rerank_ordering() {
1236        let mut r = default_reranker();
1237        let candidates = vec![
1238            make_text_candidate("doc1", "python machine learning"),
1239            make_text_candidate("doc2", "rust systems programming rust"),
1240        ];
1241        let results = r
1242            .rerank(candidates, Some("rust"), None)
1243            .expect("test: rerank should succeed");
1244        assert_eq!(results[0].id, "doc2");
1245    }
1246
1247    #[test]
1248    fn test_text_only_rerank_ranks_assigned() {
1249        let mut r = default_reranker();
1250        let candidates = vec![
1251            make_text_candidate("a", "foo"),
1252            make_text_candidate("b", "foo bar"),
1253            make_text_candidate("c", "foo bar baz"),
1254        ];
1255        let results = r
1256            .rerank(candidates, Some("foo"), None)
1257            .expect("test: rerank should succeed");
1258        for (i, res) in results.iter().enumerate() {
1259            assert_eq!(res.rank, i + 1);
1260        }
1261    }
1262
1263    #[test]
1264    fn test_text_only_empty_query_still_returns() {
1265        let mut r = default_reranker();
1266        let candidates = vec![make_text_candidate("a", "hello world")];
1267        let results = r
1268            .rerank(candidates, Some(""), None)
1269            .expect("test: rerank should succeed");
1270        assert_eq!(results.len(), 1);
1271    }
1272
1273    // ── vector-only rerank ────────────────────────────────────────────────────
1274
1275    #[test]
1276    fn test_vector_only_rerank_ordering() {
1277        let mut r = default_reranker();
1278        let query = vec![1.0_f64, 0.0];
1279        let close = make_vec_candidate("close", vec![0.99, 0.14]);
1280        let far = make_vec_candidate("far", vec![0.0, 1.0]);
1281        let results = r
1282            .rerank(vec![far, close], None, Some(&query))
1283            .expect("test: rerank should succeed");
1284        assert_eq!(results[0].id, "close");
1285    }
1286
1287    #[test]
1288    fn test_vector_only_rerank_scores_finite() {
1289        let mut r = default_reranker();
1290        let mut state: u64 = 1234;
1291        let q: Vec<f64> = (0..16).map(|_| xorshift_f64(&mut state)).collect();
1292        let candidates: Vec<RerankerCandidate> = (0..5)
1293            .map(|i| {
1294                let emb: Vec<f64> = (0..16).map(|_| xorshift_f64(&mut state)).collect();
1295                make_vec_candidate(&format!("doc{i}"), emb)
1296            })
1297            .collect();
1298        let results = r
1299            .rerank(candidates, None, Some(&q))
1300            .expect("test: rerank should succeed");
1301        for res in &results {
1302            assert!(res.final_score.is_finite());
1303        }
1304    }
1305
1306    #[test]
1307    fn test_vector_only_dimension_mismatch_error() {
1308        let mut r = default_reranker();
1309        let q = vec![1.0, 2.0, 3.0];
1310        let cand = make_vec_candidate("bad", vec![1.0, 2.0]);
1311        let err = r
1312            .rerank(vec![cand], None, Some(&q))
1313            .expect_err("test: rerank should return error for dimension mismatch");
1314        assert!(matches!(err, RerankerError::IncompatibleDimensions { .. }));
1315    }
1316
1317    // ── cross-modal fusion ────────────────────────────────────────────────────
1318
1319    #[test]
1320    fn test_cross_modal_fusion_both_modalities_present() {
1321        let mut r = default_reranker();
1322        let q_text = "rust programming";
1323        let q_emb = vec![1.0_f64, 0.0];
1324        let candidates = vec![
1325            make_full_candidate("doc1", "rust programming language", vec![0.99, 0.14]),
1326            make_full_candidate("doc2", "python data science", vec![0.0, 1.0]),
1327        ];
1328        let results = r
1329            .rerank(candidates, Some(q_text), Some(&q_emb))
1330            .expect("test: rerank should succeed");
1331        assert_eq!(results[0].id, "doc1");
1332    }
1333
1334    #[test]
1335    fn test_cross_modal_modality_scores_populated() {
1336        let mut r = default_reranker();
1337        let candidates = vec![make_full_candidate("doc1", "hello world", vec![1.0, 0.0])];
1338        let q_emb = vec![1.0, 0.0];
1339        let results = r
1340            .rerank(candidates, Some("hello"), Some(&q_emb))
1341            .expect("test: rerank should succeed");
1342        assert!(!results[0].modality_scores.is_empty());
1343        let has_text = results[0]
1344            .modality_scores
1345            .iter()
1346            .any(|ms| ms.modality == "text");
1347        let has_vec = results[0]
1348            .modality_scores
1349            .iter()
1350            .any(|ms| ms.modality == "vector");
1351        assert!(has_text);
1352        assert!(has_vec);
1353    }
1354
1355    // ── FusionStrategy: LinearCombination ────────────────────────────────────
1356
1357    #[test]
1358    fn test_linear_combination_weights() {
1359        let config = RerankerConfig {
1360            fusion_strategy: CmrFusionStrategy::LinearCombination(vec![
1361                ("text".to_string(), 0.9),
1362                ("vector".to_string(), 0.1),
1363            ]),
1364            normalize_scores: false,
1365            ..Default::default()
1366        };
1367        let mut r = CrossModalReranker::new(config);
1368        let candidates = vec![
1369            make_full_candidate("doc1", "rust is great", vec![0.0, 1.0]),
1370            make_full_candidate("doc2", "python is fine", vec![1.0, 0.0]),
1371        ];
1372        let q_emb = vec![1.0_f64, 0.0];
1373        // doc2 has better vector score, doc1 has better text score
1374        // heavy text weight → doc1 should win
1375        let results = r
1376            .rerank(candidates, Some("rust"), Some(&q_emb))
1377            .expect("test: rerank should succeed");
1378        assert_eq!(results[0].id, "doc1");
1379    }
1380
1381    // ── FusionStrategy: ReciprocalRankFusion ─────────────────────────────────
1382
1383    #[test]
1384    fn test_rrf_fusion_strategy() {
1385        let config = RerankerConfig {
1386            fusion_strategy: CmrFusionStrategy::ReciprocalRankFusion(60.0),
1387            normalize_scores: false,
1388            ..Default::default()
1389        };
1390        let mut r = CrossModalReranker::new(config);
1391        let candidates = vec![
1392            make_full_candidate("doc1", "rust rust rust", vec![0.9, 0.0]),
1393            make_full_candidate("doc2", "python", vec![0.1, 0.0]),
1394        ];
1395        let q_emb = vec![1.0_f64, 0.0];
1396        let results = r
1397            .rerank(candidates, Some("rust"), Some(&q_emb))
1398            .expect("test: rerank should succeed");
1399        assert!(!results.is_empty());
1400    }
1401
1402    // ── FusionStrategy: Borda ─────────────────────────────────────────────────
1403
1404    #[test]
1405    fn test_borda_fusion_strategy() {
1406        let config = RerankerConfig {
1407            fusion_strategy: CmrFusionStrategy::Borda,
1408            normalize_scores: false,
1409            ..Default::default()
1410        };
1411        let mut r = CrossModalReranker::new(config);
1412        let candidates = vec![
1413            make_full_candidate("doc1", "rust rust", vec![0.8, 0.2]),
1414            make_full_candidate("doc2", "python", vec![0.2, 0.8]),
1415        ];
1416        let q_emb = vec![1.0_f64, 0.0];
1417        let results = r
1418            .rerank(candidates, Some("rust"), Some(&q_emb))
1419            .expect("test: rerank should succeed");
1420        // doc1 ranks first in text; doc2 first in vector. Borda should give
1421        // each 2 points in one list and 1 in the other → equal Borda score.
1422        // Just verify no panic and scores are non-negative.
1423        for res in &results {
1424            assert!(res.final_score >= 0.0);
1425        }
1426    }
1427
1428    #[test]
1429    fn test_borda_scores_non_negative() {
1430        let config = RerankerConfig {
1431            fusion_strategy: CmrFusionStrategy::Borda,
1432            normalize_scores: false,
1433            ..Default::default()
1434        };
1435        let mut r = CrossModalReranker::new(config);
1436        let candidates = (0..5)
1437            .map(|i| make_text_candidate(&format!("d{i}"), &"word ".repeat(i + 1)))
1438            .collect();
1439        let results = r
1440            .rerank(candidates, Some("word"), None)
1441            .expect("test: rerank should succeed");
1442        for res in &results {
1443            assert!(res.final_score >= 0.0);
1444        }
1445    }
1446
1447    // ── FusionStrategy: MaxScore ──────────────────────────────────────────────
1448
1449    #[test]
1450    fn test_max_score_fusion() {
1451        let config = RerankerConfig {
1452            fusion_strategy: CmrFusionStrategy::MaxScore,
1453            normalize_scores: false,
1454            ..Default::default()
1455        };
1456        let mut r = CrossModalReranker::new(config);
1457        let candidates = vec![
1458            make_full_candidate("doc1", "rust rust rust rust", vec![0.2, 0.0]),
1459            make_full_candidate("doc2", "python", vec![0.99, 0.0]),
1460        ];
1461        let q_emb = vec![1.0_f64, 0.0];
1462        let results = r
1463            .rerank(candidates, Some("rust"), Some(&q_emb))
1464            .expect("test: rerank should succeed");
1465        // max score: doc1 should win if BM25 is higher; doc2 if vector is higher
1466        assert!(!results.is_empty());
1467        for res in &results {
1468            assert!(res.final_score.is_finite());
1469        }
1470    }
1471
1472    // ── FusionStrategy: LearnedWeights ───────────────────────────────────────
1473
1474    #[test]
1475    fn test_learned_weights_fusion() {
1476        let config = RerankerConfig {
1477            fusion_strategy: CmrFusionStrategy::LearnedWeights(vec![2.0, 1.0]),
1478            normalize_scores: false,
1479            ..Default::default()
1480        };
1481        let mut r = CrossModalReranker::new(config);
1482        let candidates = vec![
1483            make_full_candidate("doc1", "rust programming", vec![0.5, 0.0]),
1484            make_full_candidate("doc2", "java programming", vec![0.8, 0.0]),
1485        ];
1486        let q_emb = vec![1.0_f64, 0.0];
1487        let results = r
1488            .rerank(candidates, Some("rust"), Some(&q_emb))
1489            .expect("test: rerank should succeed");
1490        assert!(!results.is_empty());
1491    }
1492
1493    // ── normalization ─────────────────────────────────────────────────────────
1494
1495    #[test]
1496    fn test_normalize_scores_in_range() {
1497        let config = RerankerConfig {
1498            normalize_scores: true,
1499            ..Default::default()
1500        };
1501        let mut r = CrossModalReranker::new(config);
1502        let mut state: u64 = 54321;
1503        let q: Vec<f64> = (0..4).map(|_| xorshift_f64(&mut state)).collect();
1504        let candidates: Vec<RerankerCandidate> = (0..8)
1505            .map(|i| {
1506                let emb: Vec<f64> = (0..4).map(|_| xorshift_f64(&mut state)).collect();
1507                make_full_candidate(&format!("d{i}"), "some text here", emb)
1508            })
1509            .collect();
1510        let results = r
1511            .rerank(candidates, Some("text"), Some(&q))
1512            .expect("test: rerank should succeed");
1513        for res in &results {
1514            assert!(
1515                (0.0..=1.0).contains(&res.final_score),
1516                "score={}",
1517                res.final_score
1518            );
1519        }
1520    }
1521
1522    #[test]
1523    fn test_normalize_scores_disabled() {
1524        let config = RerankerConfig {
1525            normalize_scores: false,
1526            ..Default::default()
1527        };
1528        let mut r = CrossModalReranker::new(config);
1529        let candidates = vec![
1530            make_text_candidate("d1", "hello world hello"),
1531            make_text_candidate("d2", "foo bar"),
1532        ];
1533        let results = r
1534            .rerank(candidates, Some("hello"), None)
1535            .expect("test: rerank should succeed");
1536        // Scores may exceed 1
1537        for res in &results {
1538            assert!(res.final_score.is_finite());
1539        }
1540    }
1541
1542    #[test]
1543    fn test_normalize_all_equal_scores() {
1544        let config = RerankerConfig {
1545            normalize_scores: true,
1546            ..Default::default()
1547        };
1548        let mut r = CrossModalReranker::new(config);
1549        let candidates = vec![
1550            make_vec_candidate("d1", vec![1.0, 0.0]),
1551            make_vec_candidate("d2", vec![1.0, 0.0]),
1552        ];
1553        let q = vec![1.0, 0.0];
1554        let results = r
1555            .rerank(candidates, None, Some(&q))
1556            .expect("test: rerank should succeed");
1557        for res in &results {
1558            assert!((0.0..=1.0).contains(&res.final_score));
1559        }
1560    }
1561
1562    // ── top_k filtering ───────────────────────────────────────────────────────
1563
1564    #[test]
1565    fn test_top_k_limits_results() {
1566        let config = RerankerConfig {
1567            top_k: 3,
1568            normalize_scores: false,
1569            ..Default::default()
1570        };
1571        let mut r = CrossModalReranker::new(config);
1572        let candidates = (0..10)
1573            .map(|i| make_text_candidate(&format!("d{i}"), &format!("word {i}")))
1574            .collect();
1575        let results = r
1576            .rerank(candidates, Some("word"), None)
1577            .expect("test: rerank should succeed");
1578        assert!(results.len() <= 3);
1579    }
1580
1581    #[test]
1582    fn test_top_k_zero_returns_all() {
1583        let config = RerankerConfig {
1584            top_k: 0,
1585            normalize_scores: false,
1586            min_score_threshold: 0.0,
1587            ..Default::default()
1588        };
1589        let mut r = CrossModalReranker::new(config);
1590        let candidates = (0..5)
1591            .map(|i| make_text_candidate(&format!("d{i}"), &format!("word {i}")))
1592            .collect();
1593        let results = r
1594            .rerank(candidates, Some("word"), None)
1595            .expect("test: rerank should succeed");
1596        assert_eq!(results.len(), 5);
1597    }
1598
1599    // ── min_score_threshold ───────────────────────────────────────────────────
1600
1601    #[test]
1602    fn test_min_score_threshold_filters_low() {
1603        let config = RerankerConfig {
1604            normalize_scores: true,
1605            min_score_threshold: 0.5,
1606            top_k: 0,
1607            ..Default::default()
1608        };
1609        let mut r = CrossModalReranker::new(config);
1610        let candidates = vec![
1611            make_text_candidate("high", "rust rust rust rust"),
1612            make_text_candidate("low", "java"),
1613        ];
1614        let results = r
1615            .rerank(candidates, Some("rust"), None)
1616            .expect("test: rerank should succeed");
1617        for res in &results {
1618            assert!(
1619                res.final_score >= 0.5,
1620                "score below threshold: {}",
1621                res.final_score
1622            );
1623        }
1624    }
1625
1626    #[test]
1627    fn test_min_score_threshold_zero_keeps_all() {
1628        let config = RerankerConfig {
1629            normalize_scores: false,
1630            min_score_threshold: 0.0,
1631            top_k: 0,
1632            ..Default::default()
1633        };
1634        let mut r = CrossModalReranker::new(config);
1635        let candidates = (0..4)
1636            .map(|i| make_text_candidate(&format!("d{i}"), &format!("text {i}")))
1637            .collect();
1638        let results = r
1639            .rerank(candidates, Some("text"), None)
1640            .expect("test: rerank should succeed");
1641        assert_eq!(results.len(), 4);
1642    }
1643
1644    // ── error cases ───────────────────────────────────────────────────────────
1645
1646    #[test]
1647    fn test_error_no_candidates() {
1648        let mut r = default_reranker();
1649        let err = r
1650            .rerank(vec![], Some("query"), None)
1651            .expect_err("test: rerank should return error for empty candidates");
1652        assert_eq!(err, RerankerError::NoCandidates);
1653    }
1654
1655    #[test]
1656    fn test_error_incompatible_dimensions() {
1657        let mut r = default_reranker();
1658        let cand = make_vec_candidate("d1", vec![1.0, 2.0]);
1659        let q = vec![1.0, 2.0, 3.0];
1660        let err = r
1661            .rerank(vec![cand], None, Some(&q))
1662            .expect_err("test: rerank should return error for incompatible dimensions");
1663        assert!(matches!(err, RerankerError::IncompatibleDimensions { .. }));
1664    }
1665
1666    #[test]
1667    fn test_error_invalid_weight_negative() {
1668        let config = RerankerConfig {
1669            text_weight: -1.0,
1670            ..Default::default()
1671        };
1672        let mut r = CrossModalReranker::new(config);
1673        let cand = make_text_candidate("d1", "hello");
1674        let err = r
1675            .rerank(vec![cand], Some("hello"), None)
1676            .expect_err("test: rerank should return error for negative text_weight");
1677        assert!(matches!(err, RerankerError::InvalidWeight(_)));
1678    }
1679
1680    #[test]
1681    fn test_error_invalid_linear_weight() {
1682        let config = RerankerConfig {
1683            fusion_strategy: CmrFusionStrategy::LinearCombination(vec![("text".to_string(), -0.1)]),
1684            ..Default::default()
1685        };
1686        let mut r = CrossModalReranker::new(config);
1687        let cand = make_text_candidate("d1", "hello");
1688        let err = r
1689            .rerank(vec![cand], Some("hello"), None)
1690            .expect_err("test: rerank should return error for negative linear combination weight");
1691        assert!(matches!(err, RerankerError::InvalidWeight(_)));
1692    }
1693
1694    #[test]
1695    fn test_error_invalid_rrf_k() {
1696        let config = RerankerConfig {
1697            fusion_strategy: CmrFusionStrategy::ReciprocalRankFusion(-1.0),
1698            ..Default::default()
1699        };
1700        let mut r = CrossModalReranker::new(config);
1701        let cand = make_text_candidate("d1", "hello");
1702        let err = r
1703            .rerank(vec![cand], Some("hello"), None)
1704            .expect_err("test: rerank should return error for invalid RRF k value");
1705        assert!(matches!(err, RerankerError::ConfigurationError(_)));
1706    }
1707
1708    #[test]
1709    fn test_error_display() {
1710        let e = RerankerError::NoCandidates;
1711        assert!(!format!("{e}").is_empty());
1712        let e2 = RerankerError::IncompatibleDimensions {
1713            expected: 4,
1714            got: 3,
1715        };
1716        assert!(format!("{e2}").contains("4"));
1717    }
1718
1719    // ── stats ─────────────────────────────────────────────────────────────────
1720
1721    #[test]
1722    fn test_stats_fusion_calls_incremented() {
1723        let mut r = default_reranker();
1724        assert_eq!(r.stats().fusion_calls, 0);
1725        let _ = r.rerank(vec![make_text_candidate("a", "hi")], Some("hi"), None);
1726        assert_eq!(r.stats().fusion_calls, 1);
1727        let _ = r.rerank(vec![make_text_candidate("b", "bye")], Some("bye"), None);
1728        assert_eq!(r.stats().fusion_calls, 2);
1729    }
1730
1731    #[test]
1732    fn test_stats_candidates_reranked_accumulates() {
1733        let mut r = default_reranker();
1734        let c1 = vec![make_text_candidate("a", "a"), make_text_candidate("b", "b")];
1735        let c2 = vec![make_text_candidate("c", "c")];
1736        let _ = r.rerank(c1, Some("q"), None);
1737        let _ = r.rerank(c2, Some("q"), None);
1738        assert_eq!(r.stats().candidates_reranked, 3);
1739    }
1740
1741    #[test]
1742    fn test_stats_modalities_tracked() {
1743        let mut r = default_reranker();
1744        let cands = vec![make_full_candidate("d1", "hello", vec![1.0, 0.0])];
1745        let q_emb = vec![1.0, 0.0];
1746        let _ = r.rerank(cands, Some("hello"), Some(&q_emb));
1747        let stats = r.stats();
1748        assert!(stats.modalities_used.contains(&"text".to_string()));
1749        assert!(stats.modalities_used.contains(&"vector".to_string()));
1750    }
1751
1752    // ── update_config ─────────────────────────────────────────────────────────
1753
1754    #[test]
1755    fn test_update_config() {
1756        let mut r = default_reranker();
1757        let new_cfg = RerankerConfig {
1758            top_k: 5,
1759            ..Default::default()
1760        };
1761        r.update_config(new_cfg.clone());
1762        assert_eq!(r.config.top_k, 5);
1763    }
1764
1765    // ── ModalityScore ─────────────────────────────────────────────────────────
1766
1767    #[test]
1768    fn test_modality_score_new() {
1769        let ms = ModalityScore::new("text", 0.8, 0.4);
1770        assert_eq!(ms.modality, "text");
1771        assert!((ms.raw_score - 0.8).abs() < 1e-10);
1772        assert!((ms.normalized_score - 0.8).abs() < 1e-10);
1773        assert!((ms.weight - 0.4).abs() < 1e-10);
1774    }
1775
1776    // ── single candidate ──────────────────────────────────────────────────────
1777
1778    #[test]
1779    fn test_single_candidate_gets_rank_one() {
1780        let mut r = default_reranker();
1781        let cand = make_text_candidate("solo", "only document");
1782        let results = r
1783            .rerank(vec![cand], Some("document"), None)
1784            .expect("test: rerank should succeed");
1785        assert_eq!(results[0].rank, 1);
1786    }
1787
1788    #[test]
1789    fn test_single_candidate_score_normalised_to_one() {
1790        let config = RerankerConfig {
1791            normalize_scores: true,
1792            ..Default::default()
1793        };
1794        let mut r = CrossModalReranker::new(config);
1795        let cand = make_text_candidate("solo", "only document");
1796        let results = r
1797            .rerank(vec![cand], Some("document"), None)
1798            .expect("test: rerank should succeed");
1799        assert!((results[0].final_score - 1.0).abs() < 1e-10);
1800    }
1801
1802    // ── misc edge cases ───────────────────────────────────────────────────────
1803
1804    #[test]
1805    fn test_candidates_with_no_matching_modality_get_zero_score() {
1806        // No query text or embedding → all candidates should have final_score 0
1807        let config = RerankerConfig {
1808            normalize_scores: false,
1809            ..Default::default()
1810        };
1811        let mut r = CrossModalReranker::new(config);
1812        let candidates = vec![
1813            make_text_candidate("d1", "hello"),
1814            make_text_candidate("d2", "world"),
1815        ];
1816        // Neither query_text nor query_embedding provided
1817        let results = r
1818            .rerank(candidates, None, None)
1819            .expect("test: rerank should succeed");
1820        for res in &results {
1821            // All modality_scores will be empty → linear combo sums to 0
1822            assert_eq!(res.final_score, 0.0);
1823        }
1824    }
1825
1826    #[test]
1827    fn test_reranker_candidate_new() {
1828        let c = RerankerCandidate::new("id", Some("text"), Some(vec![1.0]));
1829        assert_eq!(c.id, "id");
1830        assert_eq!(c.text_snippet.as_deref(), Some("text"));
1831        assert_eq!(c.embedding, Some(vec![1.0]));
1832        assert!(c.modality_scores.is_empty());
1833        assert_eq!(c.rank, 0);
1834    }
1835
1836    #[test]
1837    fn test_large_candidate_set_no_panic() {
1838        let mut state: u64 = 314159;
1839        let mut r = default_reranker();
1840        let q: Vec<f64> = (0..32).map(|_| xorshift_f64(&mut state)).collect();
1841        let candidates: Vec<RerankerCandidate> = (0..200)
1842            .map(|i| {
1843                let emb: Vec<f64> = (0..32).map(|_| xorshift_f64(&mut state)).collect();
1844                make_full_candidate(&format!("d{i}"), "some query terms here", emb)
1845            })
1846            .collect();
1847        let results = r
1848            .rerank(candidates, Some("query terms"), Some(&q))
1849            .expect("test: rerank should succeed");
1850        assert!(results.len() <= 100); // default top_k = 100
1851    }
1852}