//! IOS — Information-Optimal Selection (v2: full-dominance fix)
//!
//! Two novel algorithms that compose into a single selection pass:
//!
//! 1. **Submodular Diversity Selection (SDS)**
//!    Standard knapsack treats value(A ∪ B) = value(A) + value(B).
//!    Reality: value(A ∪ B) ≤ value(A) + value(B) — information has
//!    diminishing returns. SDS penalizes redundancy using SimHash
//!    Hamming distance as a proxy for content overlap.
//!
//!    Algorithm: Lazy greedy with diversity penalty.
//!    Approximation: the (1 - 1/e) ≈ 0.63 ratio (Feige 1998 for cardinality;
//!    Sviridenko 2004 for knapsack) applies to *monotone* submodular
//!    objectives. A subtractive redundancy penalty can violate monotonicity
//!    — f(S ∪ {x}) may fall below f(S) for a near-duplicate x. For that
//!    regime, the best known polynomial-time ratio under a knapsack is
//!    0.325 (Chekuri, Vondrák, Zenklusen 2011 via continuous greedy +
//!    rounding). This implementation's empirical performance is measured in
//!    `bench/compare.py`; no tight worst-case ratio is claimed here.
//!
//! 2. **Multi-Resolution Knapsack (MRK)**
//!    Each fragment has up to 3 representations:
//!   - Full: ~100% information, ~100% tokens
//!   - Skeleton: ~70% information, ~20% tokens
//!   - Reference: ~15% information, ~2% tokens
//!
//!    This is the Multiple Choice Knapsack Problem (MCKP).
//!    Combined with SDS, each candidate is a (fragment, resolution)
//!    pair with resolution-adjusted value and diversity penalty.
//!
use crate::dedup::hamming_distance;
use crate::fragment::{compute_relevance, ContextFragment};
use std::collections::HashMap;

/// Resolution level for a selected fragment.
///
/// Hierarchical Context Synthesis: four abstraction levels that mirror
/// how engineers actually hold code in working memory.
///   Full      → raw code (100% info, 100% tokens)
///   Skeleton  → signatures + structure (70% info, 20% tokens)
///   Belief    → vault knowledge graph summary (50% info, 10-15% tokens)
///   Reference → file path only (15% info, 2% tokens)
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Resolution {
    /// Full content — maximum information, maximum tokens
    Full,
    /// Skeleton — signatures + structure, ~20% tokens
    Skeleton,
    /// Belief — vault-compiled summary with wiki-linked context, ~10-15% tokens
    Belief,
    /// Reference — file path + function name only, ~2% tokens
    Reference,
}

/// Configurable information retention factors for each resolution level.
/// These control the value/cost trade-off in multi-resolution knapsack.
/// Tunable via tuning_config.json → autotune daemon.
///
/// Belief factor (0.50): vault beliefs capture ~50% of a file's information
/// value at ~10-15% token cost. This is the key ratio for Hierarchical
/// Context Synthesis — beliefs provide architectural understanding that
/// makes raw code fragments cheaper to comprehend.
pub struct InfoFactors {
    pub skeleton: f64,  // default 0.70
    pub belief: f64,    // default 0.50
    pub reference: f64, // default 0.15
}

impl Default for InfoFactors {
    fn default() -> Self {
        InfoFactors {
            skeleton: 0.70,
            belief: 0.50,
            reference: 0.15,
        }
    }
}

impl Resolution {
    /// Information retention factor for this resolution level.
    fn info_factor(&self, factors: &InfoFactors) -> f64 {
        match self {
            Resolution::Full => 1.0,
            Resolution::Skeleton => factors.skeleton,
            Resolution::Belief => factors.belief,
            Resolution::Reference => factors.reference,
        }
    }

    pub fn as_str(&self) -> &'static str {
        match self {
            Resolution::Full => "full",
            Resolution::Skeleton => "skeleton",
            Resolution::Belief => "belief",
            Resolution::Reference => "reference",
        }
    }
}

/// A candidate item for the SDS+MRK optimizer.
/// Each fragment generates 1-4 candidates (one per resolution).
struct Candidate {
    frag_idx: usize, // Index into the fragments array
    resolution: Resolution,
    token_cost: u32, // Tokens for this resolution
    base_value: f64, // relevance × info_factor (before diversity penalty)
    simhash: u64,    // For diversity computation
}

/// Result of the IOS selection.
pub struct SdsResult {
    /// (fragment_index, chosen_resolution) pairs
    pub selections: Vec<(usize, Resolution)>,
    pub total_tokens: u32,
    pub(crate) _total_value: f64,
    pub diversity_score: f64, // Average pairwise diversity of selected set
}

/// Compute the diversity factor for a candidate given the current selected set.
///
/// diversity = 1 - max_similarity(candidate, selected_set)
///
/// Similarity is estimated from SimHash Hamming distance:
///   sim(a, b) = 1 - hamming(a, b) / 64
///
/// When the selected set is empty, diversity = 1.0 (no penalty).
///
/// Returns a value in [0, 1] where:
///   1.0 = completely novel information
///   0.0 = identical to something already selected
#[inline]
fn diversity_factor(candidate_hash: u64, selected_hashes: &[u64]) -> f64 {
    if selected_hashes.is_empty() {
        return 1.0;
    }

    let max_sim = selected_hashes
        .iter()
        .map(|&h| {
            let dist = hamming_distance(candidate_hash, h);
            // Similarity: 0 distance = identical = similarity 1.0
            (1.0 - dist as f64 / 64.0).max(0.0)
        })
        .fold(0.0_f64, f64::max);

    // Diversity = 1 - max_similarity
    // But don't penalize below diversity_floor — even similar fragments have SOME new info
    1.0 - max_sim
}

/// Compute average pairwise diversity from SimHash fingerprints.
///
/// diversity = mean over all pairs of (hamming_distance / 64).
/// Returns 1.0 when ≤ 1 hash (trivially diverse).
fn compute_pairwise_diversity(hashes: &[u64]) -> f64 {
    if hashes.len() <= 1 {
        return 1.0;
    }
    let n = hashes.len();
    let mut pair_count = 0usize;
    let mut diversity_sum = 0.0;
    for i in 0..n {
        for j in (i + 1)..n {
            let dist = hamming_distance(hashes[i], hashes[j]);
            diversity_sum += dist as f64 / 64.0;
            pair_count += 1;
        }
    }
    if pair_count > 0 {
        (diversity_sum / pair_count as f64 * 10000.0).round() / 10000.0
    } else {
        1.0
    }
}

/// IOS: Information-Optimal Selection
///
/// Combines Submodular Diversity Selection with Multi-Resolution Knapsack
/// in a single greedy pass.
///
/// Algorithm:
///   1. Generate candidates: each fragment × {full, skeleton, reference}
///   2. Separate pinned fragments (always full resolution, always included)
///   3. Greedy loop (greedy-by-density with diversity penalty):
///      - compute marginal_value = base_value × diversity_factor(hash)
///      - select candidate with highest marginal_value / token_cost
///      - remove all other resolutions of the same fragment
///      - update selected_hashes; repeat until budget exhausted
///
/// Complexity: O(N × K) where N = candidates, K = selected count
/// Typically K << N, so this is effectively O(N log N) after initial sort.
#[allow(clippy::too_many_arguments)]
pub fn ios_select(
    fragments: &[ContextFragment],
    token_budget: u32,
    w_recency: f64,
    w_frequency: f64,
    w_semantic: f64,
    w_entropy: f64,
    feedback_mults: &HashMap<String, f64>,
    enable_diversity: bool,
    enable_multi_resolution: bool,
    info_factors: &InfoFactors,
    diversity_floor: f64,
    min_candidate_value: f64,
) -> SdsResult {
    if fragments.is_empty() {
        return SdsResult {
            selections: vec![],
            total_tokens: 0,
            _total_value: 0.0,
            diversity_score: 1.0,
        };
    }

    // ── Phase 1: Budget-aware pinned fragment selection ──
    //
    // Old behavior: unconditionally include ALL pinned fragments, then use
    // remaining budget for everything else. This breaks when pinned tokens
    // exceed the budget (e.g. monorepos with 90 pinned files = 167K tokens
    // on an 8K budget → remaining_budget = 0 → zero query-relevant selection).
    //
    // New behavior: pinned fragments get a STRONG relevance boost (10x) but
    // are still subject to budget constraint. When total pinned tokens fit
    // within 50% of budget, include all (fast path). When they don't, treat
    // them as high-priority candidates in the greedy selection.
    //
    // This preserves the intent (critical files are almost always included)
    // while preventing budget blowout in large monorepos.
    let pinned_cap = token_budget / 2; // Reserve at least 50% for query-relevant fragments

    let mut all_pinned: Vec<(usize, u32)> = Vec::new(); // (index, token_count)
    let mut total_pinned_tokens: u32 = 0;
    for (i, frag) in fragments.iter().enumerate() {
        if frag.is_pinned {
            all_pinned.push((i, frag.token_count));
            total_pinned_tokens += frag.token_count;
        }
    }

    let mut pinned: Vec<(usize, Resolution)> = Vec::new();
    let mut pinned_tokens: u32 = 0;
    let mut pinned_hashes: Vec<u64> = Vec::new();
    let mut demoted_pinned: Vec<usize> = Vec::new(); // Pinned fragments that didn't fit → become high-priority candidates

    if total_pinned_tokens <= pinned_cap {
        // Fast path: all pinned fragments fit within cap
        for &(i, _tc) in &all_pinned {
            pinned.push((i, Resolution::Full));
            pinned_tokens += fragments[i].token_count;
            pinned_hashes.push(fragments[i].simhash);
        }
    } else {
        // Budget pressure: select top pinned fragments by relevance density,
        // demote the rest to high-priority candidates.
        // Sort pinned by entropy score (descending) as a proxy for importance.
        let mut scored_pinned: Vec<(usize, f64)> = all_pinned
            .iter()
            .map(|&(i, _)| {
                let frag = &fragments[i];
                let fm = feedback_mults
                    .get(&frag.fragment_id)
                    .copied()
                    .unwrap_or(1.0);
                let score =
                    compute_relevance(frag, w_recency, w_frequency, w_semantic, w_entropy, fm);
                (i, score)
            })
            .collect();
        scored_pinned
            .sort_unstable_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));

        for &(i, score) in &scored_pinned {
            let tc = fragments[i].token_count;
            if pinned_tokens + tc <= pinned_cap && score > 0.0 {
                pinned.push((i, Resolution::Full));
                pinned_tokens += tc;
                pinned_hashes.push(fragments[i].simhash);
            } else {
                demoted_pinned.push(i);
            }
        }
    }

    let remaining_budget = token_budget.saturating_sub(pinned_tokens);

    // ── Phase 2: Generate candidates ──
    // Demoted pinned fragments enter as candidates with a 5x relevance boost.
    // This ensures they're strongly preferred without being unconditional.
    let demoted_set: std::collections::HashSet<usize> = demoted_pinned.iter().copied().collect();
    let pinned_included_set: std::collections::HashSet<usize> =
        pinned.iter().map(|&(i, _)| i).collect();
    let mut candidates: Vec<Candidate> = Vec::new();

    for (i, frag) in fragments.iter().enumerate() {
        // Skip fragments already included as pinned
        if pinned_included_set.contains(&i) {
            continue;
        }

        let fm = feedback_mults
            .get(&frag.fragment_id)
            .copied()
            .unwrap_or(1.0);
        let mut relevance =
            compute_relevance(frag, w_recency, w_frequency, w_semantic, w_entropy, fm);

        // Demoted pinned fragments get 5x boost — strongly preferred but budget-constrained
        if demoted_set.contains(&i) {
            relevance *= 5.0;
        }

        if relevance <= 0.0 || frag.token_count == 0 {
            continue;
        }

        // Full resolution — always available
        candidates.push(Candidate {
            frag_idx: i,
            resolution: Resolution::Full,
            token_cost: frag.token_count,
            base_value: relevance * Resolution::Full.info_factor(info_factors),
            simhash: frag.simhash,
        });

        if enable_multi_resolution {
            // Skeleton resolution — only if skeleton was extracted
            if let Some(skel_tc) = frag.skeleton_token_count {
                let skel_value = relevance * Resolution::Skeleton.info_factor(info_factors);
                if skel_tc < frag.token_count && skel_value >= min_candidate_value {
                    candidates.push(Candidate {
                        frag_idx: i,
                        resolution: Resolution::Skeleton,
                        token_cost: skel_tc,
                        base_value: skel_value,
                        simhash: frag.simhash,
                    });
                }
            }

            // Belief resolution — only if vault belief was loaded during ingest.
            // Beliefs sit between Skeleton (structural) and Reference (path-only):
            // they capture semantic understanding at ~10-15% token cost.
            // The IOS greedy selector will pick the resolution with the best
            // value/cost ratio, so beliefs naturally win when budget is tight
            // but the file is important enough that a bare reference isn't enough.
            if let Some(belief_tc) = frag.belief_token_count {
                let belief_value = relevance * Resolution::Belief.info_factor(info_factors);
                if belief_tc < frag.token_count && belief_value >= min_candidate_value {
                    candidates.push(Candidate {
                        frag_idx: i,
                        resolution: Resolution::Belief,
                        token_cost: belief_tc,
                        base_value: belief_value,
                        simhash: frag.simhash,
                    });
                }
            }

            // Reference resolution — always available, very cheap
            // Cost: ~5 tokens for "file:source.py" reference line
            let ref_tokens = (frag.source.len() as u32 / 4).clamp(3, 10);
            let ref_value = relevance * Resolution::Reference.info_factor(info_factors);
            if ref_value >= min_candidate_value {
                candidates.push(Candidate {
                    frag_idx: i,
                    resolution: Resolution::Reference,
                    token_cost: ref_tokens,
                    base_value: ref_value,
                    simhash: frag.simhash,
                });
            }
        }
    }

    if candidates.is_empty() || remaining_budget == 0 {
        let _total_value: f64 = pinned
            .iter()
            .map(|&(i, _)| {
                let fm = feedback_mults
                    .get(&fragments[i].fragment_id)
                    .copied()
                    .unwrap_or(1.0);
                compute_relevance(
                    &fragments[i],
                    w_recency,
                    w_frequency,
                    w_semantic,
                    w_entropy,
                    fm,
                )
            })
            .sum();
        return SdsResult {
            selections: pinned,
            total_tokens: pinned_tokens,
            _total_value: (_total_value * 10000.0).round() / 10000.0,
            diversity_score: 1.0,
        };
    }

    // ── Best-Fit Fast Path ──────────────────────────────────────────
    // Best-fit-decreasing bin packing: when
    // ALL non-pinned fragments fit at full resolution, skip the
    // O(N×K) greedy loop. Common for small codebases or generous
    // ECDB budgets. Reduces to O(N).
    // ────────────────────────────────────────────────────────────────
    {
        let mut full_total: u32 = 0;
        let mut seen_frag = vec![false; fragments.len()];
        for c in &candidates {
            if c.resolution == Resolution::Full && !seen_frag[c.frag_idx] {
                seen_frag[c.frag_idx] = true;
                full_total += c.token_cost;
            }
        }
        if full_total <= remaining_budget && full_total > 0 {
            let mut selections = pinned.clone();
            let mut fast_tokens = pinned_tokens;
            let mut fast_value: f64 = selections
                .iter()
                .map(|&(i, _)| {
                    let fm = feedback_mults
                        .get(&fragments[i].fragment_id)
                        .copied()
                        .unwrap_or(1.0);
                    compute_relevance(
                        &fragments[i],
                        w_recency,
                        w_frequency,
                        w_semantic,
                        w_entropy,
                        fm,
                    )
                })
                .sum();
            let mut fast_hashes: Vec<u64> = pinned_hashes.clone();

            for c in &candidates {
                if c.resolution == Resolution::Full {
                    selections.push((c.frag_idx, Resolution::Full));
                    fast_tokens += c.token_cost;
                    fast_value += c.base_value;
                    fast_hashes.push(c.simhash);
                }
            }

            return SdsResult {
                selections,
                total_tokens: fast_tokens,
                _total_value: (fast_value * 10000.0).round() / 10000.0,
                diversity_score: compute_pairwise_diversity(&fast_hashes),
            };
        }
    }

    // ── Phase 3: Greedy SDS+MRK selection ──
    let mut selected: Vec<(usize, Resolution)> = pinned;
    let mut selected_hashes: Vec<u64> = pinned_hashes;
    let mut selected_frags: Vec<bool> = vec![false; fragments.len()]; // Track which fragment_idx is selected
    let mut budget_used = pinned_tokens;
    let mut _total_value: f64 = selected
        .iter()
        .map(|&(i, _)| {
            let fm = feedback_mults
                .get(&fragments[i].fragment_id)
                .copied()
                .unwrap_or(1.0);
            compute_relevance(
                &fragments[i],
                w_recency,
                w_frequency,
                w_semantic,
                w_entropy,
                fm,
            )
        })
        .sum();

    // Mark pinned as selected
    for &(idx, _) in &selected {
        selected_frags[idx] = true;
    }

    // Pre-sort candidates by base_value/cost density for faster convergence
    // (The diversity penalty will reorder, but this is a good initial ordering)
    candidates.sort_unstable_by(|a, b| {
        let da = a.base_value / a.token_cost.max(1) as f64;
        let db = b.base_value / b.token_cost.max(1) as f64;
        db.partial_cmp(&da).unwrap_or(std::cmp::Ordering::Equal)
    });

    loop {
        let budget_remaining = token_budget.saturating_sub(budget_used);
        if budget_remaining == 0 {
            break;
        }

        // Find the best candidate considering diversity
        let mut best_density = 0.0_f64;
        let mut best_idx: Option<usize> = None;

        // Precompute: for each fragment, does the full resolution fit?
        // If so, skip lower resolutions (full dominates when it fits).
        let mut full_fits: Vec<bool> = vec![false; fragments.len()];
        for cand in candidates.iter() {
            if cand.resolution == Resolution::Full
                && !selected_frags[cand.frag_idx]
                && cand.token_cost <= budget_remaining
            {
                full_fits[cand.frag_idx] = true;
            }
        }

        for (ci, cand) in candidates.iter().enumerate() {
            // Skip if this fragment already has a resolution selected
            if selected_frags[cand.frag_idx] {
                continue;
            }
            // Skip if doesn't fit
            if cand.token_cost > budget_remaining {
                continue;
            }
            // Skip lower resolutions when full fits — full dominates
            // because it carries strictly more information.
            if cand.resolution != Resolution::Full && full_fits[cand.frag_idx] {
                continue;
            }

            let div = if enable_diversity {
                diversity_factor(cand.simhash, &selected_hashes).max(diversity_floor)
            } else {
                1.0
            };

            let marginal_value = cand.base_value * div;
            let density = marginal_value / cand.token_cost.max(1) as f64;

            if density > best_density {
                best_density = density;
                best_idx = Some(ci);
            }
        }

        match best_idx {
            Some(ci) => {
                let cand = &candidates[ci];
                selected.push((cand.frag_idx, cand.resolution));
                selected_hashes.push(cand.simhash);
                selected_frags[cand.frag_idx] = true;
                budget_used += cand.token_cost;
                _total_value += cand.base_value; // Track pre-diversity value for reporting
            }
            None => break, // No more candidates fit
        }
    }

    // ── Phase 4: Compute diversity score of final selection ──
    let diversity_score = compute_pairwise_diversity(&selected_hashes);

    SdsResult {
        selections: selected,
        total_tokens: budget_used,
        _total_value: (_total_value * 10000.0).round() / 10000.0,
        diversity_score,
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::dedup::simhash;
    use crate::fragment::ContextFragment;

    fn empty_feedback() -> HashMap<String, f64> {
        HashMap::new()
    }

    fn default_factors() -> InfoFactors {
        InfoFactors::default()
    }

    const DEFAULT_DIV_FLOOR: f64 = 0.1;
    const DEFAULT_MIN_CANDIDATE_VALUE: f64 = 0.05;

    fn make_frag(id: &str, content: &str, tokens: u32, source: &str) -> ContextFragment {
        let mut f = ContextFragment::new(id.into(), content.into(), tokens, source.into());
        f.simhash = simhash(content);
        f.recency_score = 0.9;
        f.entropy_score = 0.7;
        f
    }

    #[test]
    fn test_empty_fragments() {
        let result = ios_select(
            &[],
            1000,
            0.3,
            0.25,
            0.25,
            0.2,
            &empty_feedback(),
            true,
            true,
            &default_factors(),
            DEFAULT_DIV_FLOOR,
            DEFAULT_MIN_CANDIDATE_VALUE,
        );
        assert!(result.selections.is_empty());
        assert_eq!(result.total_tokens, 0);
    }

    #[test]
    fn test_single_fragment_selected() {
        let frags = vec![make_frag("a", "def foo(): return 42", 50, "foo.py")];
        let result = ios_select(
            &frags,
            1000,
            0.3,
            0.25,
            0.25,
            0.2,
            &empty_feedback(),
            true,
            false,
            &default_factors(),
            DEFAULT_DIV_FLOOR,
            DEFAULT_MIN_CANDIDATE_VALUE,
        );
        assert_eq!(result.selections.len(), 1);
        assert_eq!(result.selections[0], (0, Resolution::Full));
    }

    #[test]
    fn test_pinned_always_included() {
        let mut frags = vec![
            make_frag("a", "pinned content", 500, "critical.py"),
            make_frag("b", "normal content", 200, "normal.py"),
        ];
        frags[0].is_pinned = true;
        frags[0].recency_score = 0.1; // Low recency shouldn't matter for pinned

        let result = ios_select(
            &frags,
            1100,
            0.3,
            0.25,
            0.25,
            0.2,
            &empty_feedback(),
            true,
            false,
            &default_factors(),
            DEFAULT_DIV_FLOOR,
            DEFAULT_MIN_CANDIDATE_VALUE,
        );
        let selected_indices: Vec<usize> = result.selections.iter().map(|s| s.0).collect();
        assert!(
            selected_indices.contains(&0),
            "Pinned fragment must be included"
        );
    }

    #[test]
    fn test_diversity_penalizes_duplicates() {
        // Three fragments: two nearly identical, one different
        let frags = vec![
            make_frag(
                "a",
                "def calculate_tax(income, rate): return income * rate",
                100,
                "tax1.py",
            ),
            make_frag(
                "b",
                "def calculate_tax(income, rate): return income * rate * 1.0",
                100,
                "tax2.py",
            ),
            make_frag(
                "c",
                "async fn connect_database(host: str, port: int): pass",
                100,
                "db.py",
            ),
        ];

        // With diversity: should prefer a + c (diverse) over a + b (redundant)
        let result_div = ios_select(
            &frags,
            200,
            0.3,
            0.25,
            0.25,
            0.2,
            &empty_feedback(),
            true,
            false,
            &default_factors(),
            DEFAULT_DIV_FLOOR,
            DEFAULT_MIN_CANDIDATE_VALUE,
        );
        let _div_indices: Vec<usize> = result_div.selections.iter().map(|s| s.0).collect();

        // Without diversity: might select a + b (both have high relevance)
        let result_no_div = ios_select(
            &frags,
            200,
            0.3,
            0.25,
            0.25,
            0.2,
            &empty_feedback(),
            false,
            false,
            &default_factors(),
            DEFAULT_DIV_FLOOR,
            DEFAULT_MIN_CANDIDATE_VALUE,
        );

        // With diversity enabled, we should have higher diversity score
        assert!(
            result_div.diversity_score >= result_no_div.diversity_score,
            "Diversity-enabled selection should have higher diversity: {} vs {}",
            result_div.diversity_score,
            result_no_div.diversity_score
        );
    }

    #[test]
    fn test_multi_resolution_fits_more() {
        // One large fragment that barely fits, and several that don't fit at full resolution
        let mut frags = vec![
            make_frag("big", "a very important function with lots of code\ndef process():\n    x = 1\n    y = 2\n    z = x + y\n    return z", 400, "big.py"),
            make_frag("med1", "def helper_one(): pass\ndef helper_two(): pass\ndef helper_three(): pass\ndef helper_four(): pass\ndef helper_five(): pass", 200, "h1.py"),
            make_frag("med2", "class Config:\n    debug = True\n    port = 8080\n    host = 'localhost'\n    timeout = 30\n    retries = 3", 200, "h2.py"),
        ];
        // Give them skeletons
        frags[1].skeleton_content = Some("def helper_one(): ...\ndef helper_two(): ...".into());
        frags[1].skeleton_token_count = Some(40);
        frags[2].skeleton_content = Some("class Config: ...".into());
        frags[2].skeleton_token_count = Some(30);

        // Budget: 500 tokens — can fit big(400) + one skeleton but not big + two full
        let result_mr = ios_select(
            &frags,
            500,
            0.3,
            0.25,
            0.25,
            0.2,
            &empty_feedback(),
            true,
            true,
            &default_factors(),
            DEFAULT_DIV_FLOOR,
            DEFAULT_MIN_CANDIDATE_VALUE,
        );
        let result_no_mr = ios_select(
            &frags,
            500,
            0.3,
            0.25,
            0.25,
            0.2,
            &empty_feedback(),
            true,
            false,
            &default_factors(),
            DEFAULT_DIV_FLOOR,
            DEFAULT_MIN_CANDIDATE_VALUE,
        );

        // Multi-resolution should cover more fragments
        let mr_frag_count = result_mr
            .selections
            .iter()
            .map(|s| s.0)
            .collect::<std::collections::HashSet<_>>()
            .len();
        let no_mr_frag_count = result_no_mr
            .selections
            .iter()
            .map(|s| s.0)
            .collect::<std::collections::HashSet<_>>()
            .len();

        assert!(
            mr_frag_count >= no_mr_frag_count,
            "Multi-resolution should cover >= fragments: {} vs {}",
            mr_frag_count,
            no_mr_frag_count
        );
    }

    #[test]
    fn test_budget_respected() {
        let frags = vec![
            make_frag("a", "content a", 300, "a.py"),
            make_frag("b", "content b", 300, "b.py"),
            make_frag("c", "content c", 300, "c.py"),
        ];

        let result = ios_select(
            &frags,
            500,
            0.3,
            0.25,
            0.25,
            0.2,
            &empty_feedback(),
            true,
            false,
            &default_factors(),
            DEFAULT_DIV_FLOOR,
            DEFAULT_MIN_CANDIDATE_VALUE,
        );
        assert!(
            result.total_tokens <= 500,
            "Budget must be respected: {} > 500",
            result.total_tokens
        );
    }

    #[test]
    fn test_feedback_multiplier_affects_selection() {
        let frags = vec![
            make_frag(
                "good",
                "useful code fragment for processing",
                200,
                "good.py",
            ),
            make_frag("bad", "unhelpful boilerplate noise padding", 200, "bad.py"),
        ];

        let mut feedback = HashMap::new();
        feedback.insert("good".to_string(), 1.8);
        feedback.insert("bad".to_string(), 0.3);

        // Budget for only one
        let result = ios_select(
            &frags,
            250,
            0.3,
            0.25,
            0.25,
            0.2,
            &feedback,
            true,
            false,
            &default_factors(),
            DEFAULT_DIV_FLOOR,
            DEFAULT_MIN_CANDIDATE_VALUE,
        );
        let selected_indices: Vec<usize> = result.selections.iter().map(|s| s.0).collect();

        assert!(
            selected_indices.contains(&0),
            "Feedback-boosted fragment should be preferred"
        );
    }

    #[test]
    fn test_reference_resolution_very_cheap() {
        let mut frags = vec![make_frag(
            "a",
            "def big_function():\n    x = 1\n    y = 2\n    z = 3\n    return x + y + z",
            500,
            "big.py",
        )];
        frags[0].skeleton_content = Some("def big_function(): ...".into());
        frags[0].skeleton_token_count = Some(50);

        // Budget so small only reference fits
        let result = ios_select(
            &frags,
            15,
            0.3,
            0.25,
            0.25,
            0.2,
            &empty_feedback(),
            true,
            true,
            &default_factors(),
            DEFAULT_DIV_FLOOR,
            DEFAULT_MIN_CANDIDATE_VALUE,
        );
        if !result.selections.is_empty() {
            assert_eq!(
                result.selections[0].1,
                Resolution::Reference,
                "With tiny budget, reference resolution should be chosen"
            );
        }
    }

    #[test]
    fn test_diversity_score_range() {
        let frags = vec![
            make_frag("a", "machine learning neural network training", 100, "a.py"),
            make_frag("b", "kubernetes docker container deployment", 100, "b.py"),
            make_frag(
                "c",
                "react component jsx virtual dom rendering",
                100,
                "c.py",
            ),
        ];

        let result = ios_select(
            &frags,
            1000,
            0.3,
            0.25,
            0.25,
            0.2,
            &empty_feedback(),
            true,
            false,
            &default_factors(),
            DEFAULT_DIV_FLOOR,
            DEFAULT_MIN_CANDIDATE_VALUE,
        );
        assert!(
            result.diversity_score >= 0.0 && result.diversity_score <= 1.0,
            "Diversity score must be in [0, 1], got {}",
            result.diversity_score
        );
    }

    #[test]
    fn test_resolution_preference_by_budget() {
        // When budget is generous, prefer full; when tight, prefer skeleton
        let mut frags = vec![
            make_frag(
                "a",
                "def foo():\n    return 1 + 2 + 3 + 4 + 5\n",
                200,
                "a.py",
            ),
            make_frag(
                "b",
                "def bar():\n    return 6 + 7 + 8 + 9 + 10\n",
                200,
                "b.py",
            ),
        ];
        frags[0].skeleton_content = Some("def foo(): ...".into());
        frags[0].skeleton_token_count = Some(30);
        frags[1].skeleton_content = Some("def bar(): ...".into());
        frags[1].skeleton_token_count = Some(30);

        // Generous budget: both full
        let result_big = ios_select(
            &frags,
            1000,
            0.3,
            0.25,
            0.25,
            0.2,
            &empty_feedback(),
            true,
            true,
            &default_factors(),
            DEFAULT_DIV_FLOOR,
            DEFAULT_MIN_CANDIDATE_VALUE,
        );
        let full_count_big = result_big
            .selections
            .iter()
            .filter(|s| s.1 == Resolution::Full)
            .count();

        // Tight budget: mix of full + skeleton/reference
        let result_tight = ios_select(
            &frags,
            250,
            0.3,
            0.25,
            0.25,
            0.2,
            &empty_feedback(),
            true,
            true,
            &default_factors(),
            DEFAULT_DIV_FLOOR,
            DEFAULT_MIN_CANDIDATE_VALUE,
        );
        let full_count_tight = result_tight
            .selections
            .iter()
            .filter(|s| s.1 == Resolution::Full)
            .count();

        assert!(
            full_count_big >= full_count_tight,
            "Generous budget should select more full-resolution fragments"
        );
    }

    #[test]
    fn test_fast_path_selects_all_when_budget_generous() {
        // 3 fragments totalling 150 tokens, budget = 500
        // Should trigger fast path: all selected at full resolution
        let frags = vec![
            make_frag("a", "def alpha(): return 1", 50, "a.py"),
            make_frag("b", "def beta(): return 2", 50, "b.py"),
            make_frag("c", "def gamma(): return 3", 50, "c.py"),
        ];

        let result = ios_select(
            &frags,
            500,
            0.3,
            0.25,
            0.25,
            0.2,
            &empty_feedback(),
            true,
            false,
            &default_factors(),
            DEFAULT_DIV_FLOOR,
            DEFAULT_MIN_CANDIDATE_VALUE,
        );

        assert_eq!(result.selections.len(), 3, "Fast path should select all 3");
        assert!(
            result.selections.iter().all(|s| s.1 == Resolution::Full),
            "Fast path should use full resolution for all"
        );
        assert_eq!(result.total_tokens, 150);
    }
}
