//! MaxScore disjunction scorer with block-level skip (WAND optimization).
//!
//! Partitions sub-scorers into essential and non-essential based on their
//! `max_score()` bounds. Essential scorers drive iteration; non-essential
//! are only evaluated when the essential sum could be competitive.
//!
//! See [[architecture-query-execution#WAND / MaxScore Optimization]] and
//! [[optimization-block-max-wand]].

use luci_core::{DocId, Scorer, TwoPhaseIterator, NO_MORE_DOCS};

/// MaxScore disjunction scorer.
///
/// Sub-scorers are sorted by `max_score()` ascending. The partition splits
/// them into non-essential (head, low max_score) and essential (tail).
pub struct WANDScorer {
    /// Sub-scorers sorted by max_score ascending.
    /// `scorers[0..partition_idx]` = non-essential
    /// `scorers[partition_idx..]` = essential
    scorers: Vec<Box<dyn Scorer>>,

    /// Cached doc_ids for each scorer (avoids repeated vtable dispatch).
    /// Updated whenever a scorer advances.
    doc_ids: Vec<DocId>,

    /// Cumulative max_score prefix sums.
    /// `max_score_prefix[i]` = sum of max_score for `scorers[0..=i]`.
    max_score_prefix: Vec<f32>,

    /// Index into scorers: `scorers[partition_idx..]` are essential.
    partition_idx: usize,

    /// Current minimum competitive score (from collector feedback).
    min_competitive_score: f32,

    /// Current document state.
    current: DocId,
    current_score: f32,
}

impl WANDScorer {
    pub fn new(mut scorers: Vec<Box<dyn Scorer>>) -> Self {
        assert!(!scorers.is_empty(), "WAND needs at least one scorer");

        // Remove exhausted scorers
        scorers.retain(|s| s.doc_id() != NO_MORE_DOCS);

        // Sort by max_score ascending (lowest max contribution first)
        scorers.sort_by(|a, b| {
            a.max_score()
                .partial_cmp(&b.max_score())
                .unwrap_or(std::cmp::Ordering::Equal)
        });

        // Build prefix sums
        let mut max_score_prefix = Vec::with_capacity(scorers.len());
        let mut cumulative = 0.0f32;
        for s in &scorers {
            cumulative += s.max_score();
            max_score_prefix.push(cumulative);
        }

        let doc_ids: Vec<DocId> = scorers.iter().map(|s| s.doc_id()).collect();

        let mut wand = Self {
            scorers,
            doc_ids,
            max_score_prefix,
            partition_idx: 0, // all essential initially
            min_competitive_score: 0.0,
            current: NO_MORE_DOCS,
            current_score: 0.0,
        };

        // Position on the first competitive doc
        if !wand.scorers.is_empty() {
            wand.find_next_competitive();
        }

        wand
    }

    /// Recompute the essential/non-essential partition based on
    /// `min_competitive_score`.
    fn update_partition(&mut self) {
        self.partition_idx = 0;
        for i in 0..self.scorers.len() {
            if self.max_score_prefix[i] >= self.min_competitive_score {
                break;
            }
            self.partition_idx = i + 1;
        }
    }

    /// Non-essential max score sum (prefix sum up to partition_idx - 1).
    fn non_essential_max_sum(&self) -> f32 {
        if self.partition_idx == 0 {
            0.0
        } else {
            self.max_score_prefix[self.partition_idx - 1]
        }
    }

    /// Find the minimum doc ID among essential scorers using cached doc_ids.
    #[inline(always)]
    fn min_essential_doc(&self) -> DocId {
        let mut min_doc = NO_MORE_DOCS;
        for &d in &self.doc_ids[self.partition_idx..] {
            if d < min_doc {
                min_doc = d;
            }
        }
        min_doc
    }

    /// Find the next document that could be competitive and set `current`.
    /// Uses cached doc_ids to avoid redundant vtable dispatches.
    fn find_next_competitive(&mut self) {
        let n = self.scorers.len();

        loop {
            if self.partition_idx >= n {
                self.current = NO_MORE_DOCS;
                self.current_score = 0.0;
                return;
            }

            let min_doc = self.min_essential_doc();
            if min_doc == NO_MORE_DOCS {
                self.current = NO_MORE_DOCS;
                self.current_score = 0.0;
                return;
            }

            // Sum scores from essential scorers at min_doc (using cached doc_ids)
            let mut essential_score = 0.0f32;
            for i in self.partition_idx..n {
                if self.doc_ids[i] == min_doc {
                    essential_score += self.scorers[i].score();
                }
            }

            // Block-max pre-check on non-essential scorers
            let mut non_essential_upper = 0.0f32;
            for i in 0..self.partition_idx {
                non_essential_upper += self.scorers[i].block_max_score(min_doc);
            }

            let potentially_competitive =
                self.min_competitive_score == 0.0
                    || essential_score + non_essential_upper >= self.min_competitive_score;

            let mut total_score = essential_score;

            if potentially_competitive && self.partition_idx > 0 {
                for i in 0..self.partition_idx {
                    let advanced = self.scorers[i].advance(min_doc);
                    self.doc_ids[i] = advanced;
                    if advanced == min_doc {
                        total_score += self.scorers[i].score();
                    }
                }
            }

            // Advance essential scorers past min_doc (update cached doc_ids)
            for i in self.partition_idx..n {
                if self.doc_ids[i] == min_doc {
                    self.doc_ids[i] = self.scorers[i].next();
                }
            }

            if potentially_competitive && total_score > 0.0 {
                self.current = min_doc;
                self.current_score = total_score;
                return;
            }
        }
    }
}

impl Scorer for WANDScorer {
    fn doc_id(&self) -> DocId {
        self.current
    }

    fn next(&mut self) -> DocId {
        if self.current == NO_MORE_DOCS {
            return NO_MORE_DOCS;
        }
        self.find_next_competitive();
        self.current
    }

    fn advance(&mut self, target: DocId) -> DocId {
        if self.current >= target {
            return self.current;
        }

        // Advance all scorers to at least target (update cached doc_ids)
        for i in 0..self.scorers.len() {
            if self.doc_ids[i] < target {
                self.doc_ids[i] = self.scorers[i].advance(target);
            }
        }

        self.find_next_competitive();
        self.current
    }

    fn score(&mut self) -> f32 {
        self.current_score
    }

    fn two_phase(&mut self) -> Option<&mut dyn TwoPhaseIterator> {
        None
    }

    fn max_score(&self) -> f32 {
        self.max_score_prefix.last().copied().unwrap_or(0.0)
    }

    fn set_min_competitive_score(&mut self, min_score: f32) {
        if min_score > self.min_competitive_score {
            self.min_competitive_score = min_score;
            self.update_partition();

            if self.partition_idx >= self.scorers.len() {
                self.current = NO_MORE_DOCS;
                self.current_score = 0.0;
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    struct VecScorer {
        docs: Vec<(DocId, f32)>,
        pos: usize,
        max: f32,
    }

    impl VecScorer {
        fn new(docs: Vec<(u32, f32)>, max: f32) -> Box<dyn Scorer> {
            Box::new(Self {
                docs: docs.into_iter().map(|(id, s)| (DocId::new(id), s)).collect(),
                pos: 0,
                max,
            })
        }
        fn current(&self) -> (DocId, f32) {
            if self.pos < self.docs.len() {
                self.docs[self.pos]
            } else {
                (NO_MORE_DOCS, 0.0)
            }
        }
    }

    impl Scorer for VecScorer {
        fn doc_id(&self) -> DocId { self.current().0 }
        fn next(&mut self) -> DocId {
            if self.pos < self.docs.len() { self.pos += 1; }
            self.current().0
        }
        fn advance(&mut self, target: DocId) -> DocId {
            while self.pos < self.docs.len() && self.docs[self.pos].0 < target {
                self.pos += 1;
            }
            self.current().0
        }
        fn score(&mut self) -> f32 { self.current().1 }
        fn two_phase(&mut self) -> Option<&mut dyn TwoPhaseIterator> { None }
        fn max_score(&self) -> f32 { self.max }
    }

    /// Gold standard: WANDScorer produces identical results to exhaustive
    /// collection when min_competitive_score is not set.
    #[test]
    fn wand_matches_exhaustive() {
        let s1 = VecScorer::new(vec![(0, 1.0), (2, 1.0), (4, 1.0)], 1.0);
        let s2 = VecScorer::new(vec![(1, 2.0), (2, 2.0), (3, 2.0)], 2.0);

        let mut wand = WANDScorer::new(vec![s1, s2]);

        let mut results = Vec::new();
        while wand.doc_id() != NO_MORE_DOCS {
            results.push((wand.doc_id().as_u32(), wand.score()));
            wand.next();
        }

        assert_eq!(results, vec![
            (0, 1.0), (1, 2.0), (2, 3.0), (3, 2.0), (4, 1.0)
        ]);
    }

    #[test]
    fn wand_three_scorers() {
        let s1 = VecScorer::new(vec![(5, 1.0)], 1.0);
        let s2 = VecScorer::new(vec![(5, 2.0)], 2.0);
        let s3 = VecScorer::new(vec![(5, 3.0)], 3.0);

        let mut wand = WANDScorer::new(vec![s1, s2, s3]);
        assert_eq!(wand.doc_id(), DocId::new(5));
        assert_eq!(wand.score(), 6.0);
        assert_eq!(wand.next(), NO_MORE_DOCS);
    }

    #[test]
    fn wand_advance() {
        let s1 = VecScorer::new(vec![(0, 1.0), (5, 1.0), (10, 1.0)], 1.0);
        let s2 = VecScorer::new(vec![(3, 2.0), (7, 2.0), (10, 2.0)], 2.0);

        let mut wand = WANDScorer::new(vec![s1, s2]);
        assert_eq!(wand.advance(DocId::new(5)), DocId::new(5));
        assert_eq!(wand.score(), 1.0);
        assert_eq!(wand.advance(DocId::new(10)), DocId::new(10));
        assert_eq!(wand.score(), 3.0);
    }

    #[test]
    fn wand_partition_moves_with_threshold() {
        // s1 has max_score=1.0, s2 has max_score=5.0
        let s1 = VecScorer::new(vec![(0, 0.5), (1, 0.5)], 1.0);
        let s2 = VecScorer::new(vec![(0, 4.0), (1, 4.0)], 5.0);

        let mut wand = WANDScorer::new(vec![s1, s2]);
        // Initially all essential (partition_idx=0)
        assert_eq!(wand.partition_idx, 0);

        // Set threshold > s1.max_score but < s1.max + s2.max
        wand.set_min_competitive_score(2.0);
        // s1 (max_score=1.0) becomes non-essential since prefix[0]=1.0 < 2.0
        assert_eq!(wand.partition_idx, 1);
    }

    #[test]
    fn wand_all_non_essential_exhausts() {
        let s1 = VecScorer::new(vec![(0, 0.5)], 1.0);
        let s2 = VecScorer::new(vec![(0, 1.0)], 2.0);

        let mut wand = WANDScorer::new(vec![s1, s2]);
        assert_ne!(wand.doc_id(), NO_MORE_DOCS);

        // Set threshold higher than total max_score
        wand.set_min_competitive_score(10.0);
        assert_eq!(wand.doc_id(), NO_MORE_DOCS);
    }

    #[test]
    fn wand_non_overlapping() {
        let s1 = VecScorer::new(vec![(0, 1.0), (2, 1.0)], 1.0);
        let s2 = VecScorer::new(vec![(1, 2.0), (3, 2.0)], 2.0);

        let mut wand = WANDScorer::new(vec![s1, s2]);
        let mut docs = Vec::new();
        while wand.doc_id() != NO_MORE_DOCS {
            docs.push(wand.doc_id().as_u32());
            wand.next();
        }
        assert_eq!(docs, vec![0, 1, 2, 3]);
    }

    #[test]
    fn wand_skips_non_competitive_docs() {
        // 3 scorers, set high threshold → only docs scoring above threshold survive
        let s1 = VecScorer::new(vec![(0, 1.0), (1, 1.0), (2, 1.0)], 1.0);
        let s2 = VecScorer::new(vec![(1, 2.0), (2, 2.0)], 2.0);
        let s3 = VecScorer::new(vec![(2, 3.0)], 3.0);

        let mut wand = WANDScorer::new(vec![s1, s2, s3]);

        // Doc 0 (score=1.0) was found during construction before any threshold.
        assert_eq!(wand.doc_id(), DocId::new(0));
        assert_eq!(wand.score(), 1.0);

        // Set threshold so only doc 2 (score=6.0) is competitive going forward.
        // Doc 1 (score=3.0) should be skipped.
        wand.set_min_competitive_score(5.0);
        wand.next();

        assert_eq!(wand.doc_id(), DocId::new(2));
        assert_eq!(wand.score(), 6.0);
        assert_eq!(wand.next(), NO_MORE_DOCS);
    }

    #[test]
    fn wand_max_score() {
        let s1 = VecScorer::new(vec![(0, 1.0)], 1.5);
        let s2 = VecScorer::new(vec![(0, 2.0)], 2.5);

        let wand = WANDScorer::new(vec![s1, s2]);
        assert_eq!(wand.max_score(), 4.0);
    }
}
