//! AutoTune engine: automated identity plan optimization via mutation search.
//!
//! Port of the Python SDK's `autotune.py`, `_mutations.py`, and `_metrics.py`.
//! Runs a 3-phase mutation loop over an `IdentityPlan`, using the
//! `ReconciliationEngine` to evaluate each candidate.

use crate::engine::overrides::OverrideResolver;
use crate::engine::types::{MatchDecision, NormalizedEntity};
use crate::engine::ReconciliationEngine;
use cannon_common::ir::{
    BlockingKey, ClusteringMethod, CompiledClusteringConfig, CompiledRuleType, FsComparatorType,
    IdentityPlan,
};
use rand::seq::SliceRandom;
use rand::SeedableRng;
use rand_chacha::ChaCha8Rng;
use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, HashMap, HashSet};
use tracing::{debug, info};
use uuid::Uuid;

// ---------------------------------------------------------------------------
// Configuration
// ---------------------------------------------------------------------------

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AutoTuneConfig {
    #[serde(default = "default_max_iterations")]
    pub max_iterations: usize,
    #[serde(default = "default_max_conflict")]
    pub max_conflict: f64,
    #[serde(default = "default_merge_rate_weight")]
    pub merge_rate_weight: f64,
    #[serde(default = "default_conflict_rate_weight")]
    pub conflict_rate_weight: f64,
    #[serde(default = "default_entropy_weight")]
    pub entropy_weight: f64,
    #[serde(default = "default_stability_weight")]
    pub stability_weight: f64,
    /// Optional sample size for stratified sampling. `None` means auto-calculate
    /// based on dataset size (no sampling for <= 5K records).
    #[serde(default)]
    pub sample_size: Option<usize>,
}

fn default_max_iterations() -> usize { 50 }
fn default_max_conflict() -> f64 { 0.05 }
fn default_merge_rate_weight() -> f64 { 0.3 }
fn default_conflict_rate_weight() -> f64 { 0.4 }
fn default_entropy_weight() -> f64 { 0.1 }
fn default_stability_weight() -> f64 { 0.2 }

impl Default for AutoTuneConfig {
    fn default() -> Self {
        Self {
            max_iterations: 50,
            max_conflict: 0.05,
            merge_rate_weight: 0.3,
            conflict_rate_weight: 0.4,
            entropy_weight: 0.1,
            stability_weight: 0.2,
            sample_size: None,
        }
    }
}

// ---------------------------------------------------------------------------
// Result types
// ---------------------------------------------------------------------------

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AutoTuneResult {
    pub best_plan: IdentityPlan,
    pub original_plan: IdentityPlan,
    pub metrics_before: HashMap<String, f64>,
    pub metrics_after: HashMap<String, f64>,
    pub mutations_accepted: Vec<MutationLogEntry>,
    pub mutations_total: usize,
    pub iterations_used: usize,
    pub explanation: String,
    pub total_entities: usize,
    pub sampled_entities: usize,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MutationLogEntry {
    pub mutation: Mutation,
    pub score_delta: f64,
    pub accepted: bool,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Mutation {
    pub tier: u8,
    pub category: String,
    pub description: String,
    pub field_path: String,
}

// ---------------------------------------------------------------------------
// Labeled pair for supervised tuning
// ---------------------------------------------------------------------------

#[derive(Debug, Clone)]
pub struct LabeledPair {
    pub entity_a_id: Uuid,
    pub entity_b_id: Uuid,
    pub is_match: bool,
}

// ---------------------------------------------------------------------------
// Metrics computation (port of _metrics.py)
// ---------------------------------------------------------------------------

/// Compute quality metrics for a reconciliation result.
fn compute_metrics(
    decisions: &[MatchDecision],
    entities: &[NormalizedEntity],
    identity_fields: &[String],
    baseline_clusters: Option<&[Vec<Uuid>]>,
    labels: &[LabeledPair],
) -> HashMap<String, f64> {
    let clusters = build_clusters(decisions, entities);
    let mut metrics = HashMap::new();

    // Merge rate: fraction of entities in multi-member clusters
    let multi_member_count: usize = clusters
        .iter()
        .filter(|c| c.len() > 1)
        .map(|c| c.len())
        .sum();
    let merge_rate = if entities.is_empty() {
        0.0
    } else {
        multi_member_count as f64 / entities.len() as f64
    };
    metrics.insert("merge_rate".into(), merge_rate);

    // Conflict rate
    let cr = conflict_rate(&clusters, entities, identity_fields);
    metrics.insert("conflict_rate".into(), cr);

    // Identifier entropy
    let ent = identifier_entropy(&clusters, entities, identity_fields);
    metrics.insert("entropy".into(), ent);

    // Stability (ARI vs baseline)
    if let Some(baseline) = baseline_clusters {
        let ari = adjusted_rand_index(&clusters, baseline);
        metrics.insert("stability".into(), ari);
    } else {
        metrics.insert("stability".into(), 1.0);
    }

    // F1/precision/recall from labeled pairs
    if !labels.is_empty() {
        let (f1, precision, recall) = compute_f1_from_labels(decisions, labels);
        metrics.insert("f1".into(), f1);
        metrics.insert("precision".into(), precision);
        metrics.insert("recall".into(), recall);
    }

    metrics
}

/// Build clusters from merge decisions using union-find.
fn build_clusters(decisions: &[MatchDecision], entities: &[NormalizedEntity]) -> Vec<Vec<Uuid>> {
    let mut parent: HashMap<Uuid, Uuid> = HashMap::new();
    for e in entities {
        parent.insert(e.id, e.id);
    }

    fn find(parent: &mut HashMap<Uuid, Uuid>, x: Uuid) -> Uuid {
        let mut root = x;
        while parent.get(&root).copied().unwrap_or(root) != root {
            root = parent.get(&root).copied().unwrap_or(root);
        }
        // Path compression
        let mut cur = x;
        while cur != root {
            let next = parent.get(&cur).copied().unwrap_or(cur);
            parent.insert(cur, root);
            cur = next;
        }
        root
    }

    fn union(parent: &mut HashMap<Uuid, Uuid>, a: Uuid, b: Uuid) {
        let ra = find(parent, a);
        let rb = find(parent, b);
        if ra != rb {
            parent.insert(rb, ra);
        }
    }

    for d in decisions {
        if d.decision == crate::engine::types::Decision::Merge {
            union(&mut parent, d.entity_a_id, d.entity_b_id);
        }
    }

    let mut groups: HashMap<Uuid, Vec<Uuid>> = HashMap::new();
    let ids: Vec<Uuid> = parent.keys().copied().collect();
    for id in ids {
        let root = find(&mut parent, id);
        groups.entry(root).or_default().push(id);
    }

    // Sort by external_id for determinism
    let entity_map: HashMap<Uuid, &str> = entities
        .iter()
        .map(|e| (e.id, e.external_id.as_str()))
        .collect();
    let mut clusters: Vec<Vec<Uuid>> = groups.into_values().collect();
    for cluster in &mut clusters {
        cluster.sort_by(|a, b| {
            let ea = entity_map.get(a).unwrap_or(&"");
            let eb = entity_map.get(b).unwrap_or(&"");
            ea.cmp(eb)
        });
    }
    clusters.sort_by(|a, b| {
        let ea = a.first().and_then(|id| entity_map.get(id)).unwrap_or(&"");
        let eb = b.first().and_then(|id| entity_map.get(id)).unwrap_or(&"");
        ea.cmp(eb)
    });

    clusters
}

/// Fraction of multi-member clusters with conflicting identity field values.
fn conflict_rate(
    clusters: &[Vec<Uuid>],
    entities: &[NormalizedEntity],
    identity_fields: &[String],
) -> f64 {
    let entity_data: HashMap<Uuid, &HashMap<String, String>> =
        entities.iter().map(|e| (e.id, &e.data)).collect();

    let multi: Vec<&Vec<Uuid>> = clusters.iter().filter(|c| c.len() > 1).collect();
    if multi.is_empty() {
        return 0.0;
    }

    let mut conflicting = 0usize;
    for cluster in &multi {
        let mut has_conflict = false;
        for field in identity_fields {
            let mut values: HashSet<String> = HashSet::new();
            for uid in *cluster {
                if let Some(data) = entity_data.get(uid) {
                    if let Some(val) = data.get(field) {
                        let normalized = val.trim().to_lowercase();
                        if !normalized.is_empty() {
                            values.insert(normalized);
                        }
                    }
                }
            }
            if values.len() > 1 {
                has_conflict = true;
                break;
            }
        }
        if has_conflict {
            conflicting += 1;
        }
    }

    conflicting as f64 / multi.len() as f64
}

/// Shannon entropy of identifier diversity within multi-member clusters.
fn identifier_entropy(
    clusters: &[Vec<Uuid>],
    entities: &[NormalizedEntity],
    identity_fields: &[String],
) -> f64 {
    let entity_data: HashMap<Uuid, &HashMap<String, String>> =
        entities.iter().map(|e| (e.id, &e.data)).collect();

    let multi: Vec<&Vec<Uuid>> = clusters.iter().filter(|c| c.len() > 1).collect();
    if multi.is_empty() {
        return 0.0;
    }

    let mut total_entropy = 0.0f64;
    let mut count = 0usize;

    for cluster in &multi {
        for field in identity_fields {
            let mut freq: HashMap<String, usize> = HashMap::new();
            let mut total = 0usize;
            for uid in *cluster {
                if let Some(data) = entity_data.get(uid) {
                    if let Some(val) = data.get(field) {
                        let key = val.trim().to_lowercase();
                        if !key.is_empty() {
                            *freq.entry(key).or_insert(0) += 1;
                            total += 1;
                        }
                    }
                }
            }
            if total > 0 && !freq.is_empty() {
                let mut h = 0.0f64;
                for &f in freq.values() {
                    let p = f as f64 / total as f64;
                    if p > 0.0 {
                        h -= p * p.log2();
                    }
                }
                total_entropy += h;
                count += 1;
            }
        }
    }

    if count > 0 {
        total_entropy / count as f64
    } else {
        0.0
    }
}

/// Adjusted Rand Index between two clusterings.
fn adjusted_rand_index(clusters_a: &[Vec<Uuid>], clusters_b: &[Vec<Uuid>]) -> f64 {
    let mut label_a: HashMap<Uuid, usize> = HashMap::new();
    for (i, cluster) in clusters_a.iter().enumerate() {
        for uid in cluster {
            label_a.insert(*uid, i);
        }
    }

    let mut label_b: HashMap<Uuid, usize> = HashMap::new();
    for (i, cluster) in clusters_b.iter().enumerate() {
        for uid in cluster {
            label_b.insert(*uid, i);
        }
    }

    // Only consider elements present in both
    let common: HashSet<Uuid> = label_a
        .keys()
        .filter(|k| label_b.contains_key(k))
        .copied()
        .collect();

    if common.len() < 2 {
        return if common.len() <= 1 { 1.0 } else { 0.0 };
    }

    // Build contingency table
    let mut contingency: HashMap<(usize, usize), usize> = HashMap::new();
    let mut row_sums: HashMap<usize, usize> = HashMap::new();
    let mut col_sums: HashMap<usize, usize> = HashMap::new();

    for uid in &common {
        let a = label_a[uid];
        let b = label_b[uid];
        *contingency.entry((a, b)).or_insert(0) += 1;
        *row_sums.entry(a).or_insert(0) += 1;
        *col_sums.entry(b).or_insert(0) += 1;
    }

    let n = common.len();
    let comb2 = |x: usize| -> usize { x * x.saturating_sub(1) / 2 };

    let sum_comb_nij: usize = contingency.values().map(|&v| comb2(v)).sum();
    let sum_comb_a: usize = row_sums.values().map(|&v| comb2(v)).sum();
    let sum_comb_b: usize = col_sums.values().map(|&v| comb2(v)).sum();

    let comb_n = comb2(n);
    if comb_n == 0 {
        return 1.0;
    }

    let expected = (sum_comb_a as f64 * sum_comb_b as f64) / comb_n as f64;
    let max_index = (sum_comb_a as f64 + sum_comb_b as f64) / 2.0;

    if (max_index - expected).abs() < f64::EPSILON {
        return 1.0;
    }

    (sum_comb_nij as f64 - expected) / (max_index - expected)
}

/// Compute F1, precision, recall from labeled pairs.
fn compute_f1_from_labels(
    decisions: &[MatchDecision],
    labels: &[LabeledPair],
) -> (f64, f64, f64) {
    // Build set of merged pairs
    let mut merged_pairs: HashSet<(Uuid, Uuid)> = HashSet::new();
    for d in decisions {
        if d.decision == crate::engine::types::Decision::Merge {
            let pair = if d.entity_a_id < d.entity_b_id {
                (d.entity_a_id, d.entity_b_id)
            } else {
                (d.entity_b_id, d.entity_a_id)
            };
            merged_pairs.insert(pair);
        }
    }

    let mut tp = 0usize;
    let mut fp = 0usize;
    let mut fn_ = 0usize;

    for label in labels {
        let pair = if label.entity_a_id < label.entity_b_id {
            (label.entity_a_id, label.entity_b_id)
        } else {
            (label.entity_b_id, label.entity_a_id)
        };
        let predicted = merged_pairs.contains(&pair);

        match (label.is_match, predicted) {
            (true, true) => tp += 1,
            (false, true) => fp += 1,
            (true, false) => fn_ += 1,
            (false, false) => {}
        }
    }

    let precision = if tp + fp > 0 {
        tp as f64 / (tp + fp) as f64
    } else {
        0.0
    };
    let recall = if tp + fn_ > 0 {
        tp as f64 / (tp + fn_) as f64
    } else {
        0.0
    };
    let f1 = if precision + recall > 0.0 {
        2.0 * precision * recall / (precision + recall)
    } else {
        0.0
    };

    (f1, precision, recall)
}

/// Infer identity fields from a plan's rules and FS fields.
fn infer_identity_fields(plan: &IdentityPlan) -> Vec<String> {
    let mut fields = Vec::new();
    let mut seen: HashSet<String> = HashSet::new();

    // From match graph rules
    for rule in &plan.match_graph.rules {
        let field = match &rule.rule_type {
            CompiledRuleType::Exact { field, .. } => Some(field.clone()),
            CompiledRuleType::Similarity { field, .. } => Some(field.clone()),
            CompiledRuleType::Range { field, .. } => Some(field.clone()),
            _ => None,
        };
        if let Some(f) = field {
            if !seen.contains(&f) {
                // Check if it's identity-oriented
                let normalizer = match &rule.rule_type {
                    CompiledRuleType::Exact { normalizer, .. } => normalizer.clone(),
                    CompiledRuleType::Similarity { normalizer, .. } => normalizer.clone(),
                    _ => None,
                };
                let is_exact = matches!(rule.rule_type, CompiledRuleType::Exact { .. });
                let is_identity_normalizer = normalizer
                    .as_deref()
                    .is_some_and(|n| {
                        matches!(n, "email" | "phone" | "name" | "nickname")
                    });
                let is_identity_name = {
                    let lower = f.to_lowercase();
                    lower.contains("email")
                        || lower.contains("phone")
                        || lower.contains("ssn")
                        || lower.contains("national_id")
                };

                if is_exact || is_identity_normalizer || is_identity_name {
                    fields.push(f.clone());
                    seen.insert(f);
                }
            }
        }
    }

    // From FS fields
    if let Some(ref fs) = plan.decision.fellegi_sunter {
        for ff in &fs.fields {
            if !seen.contains(&ff.name) {
                fields.push(ff.name.clone());
                seen.insert(ff.name.clone());
            }
        }
    }

    // Fallback
    if fields.is_empty() {
        for fallback in &["email", "phone"] {
            fields.push(fallback.to_string());
        }
    }

    fields
}

// ---------------------------------------------------------------------------
// Stratified sampling
// ---------------------------------------------------------------------------

/// Determine the automatic sample size based on dataset size.
///
/// Tiers:
/// - <= 5K records: no sampling (use all)
/// - 5K-50K: sample 5K
/// - 50K-500K: sample 10K
/// - 500K+: sample 20K
pub fn auto_sample_size(total: usize) -> usize {
    if total <= 5_000 {
        total
    } else if total <= 50_000 {
        5_000
    } else if total <= 500_000 {
        10_000
    } else {
        20_000
    }
}

/// Sample entities using stratified sampling by source, preserving source
/// distribution with a minimum of 100 records per source guaranteed.
///
/// Uses a seeded ChaCha8 RNG for determinism (same data = same sample).
/// Result is sorted by `external_id` for deterministic downstream behavior.
pub fn sample_stratified(
    entities: &[NormalizedEntity],
    sample_size: Option<usize>,
    seed: u64,
) -> Vec<NormalizedEntity> {
    let total = entities.len();
    let target = sample_size
        .unwrap_or_else(|| auto_sample_size(total))
        .min(total);

    if target >= total {
        return entities.to_vec();
    }

    // Group by source_name using BTreeMap for deterministic iteration
    let mut by_source: BTreeMap<&str, Vec<usize>> = BTreeMap::new();
    for (i, e) in entities.iter().enumerate() {
        by_source.entry(&e.source_name).or_default().push(i);
    }

    let num_sources = by_source.len();
    let min_per_source: usize = 100;

    let mut rng = ChaCha8Rng::seed_from_u64(seed);
    let mut selected_indices: Vec<usize> = Vec::with_capacity(target);

    // Phase 1: guarantee minimum per source
    let mut remaining_budget = target;
    let mut source_counts: BTreeMap<&str, usize> = BTreeMap::new();

    for (source, indices) in &mut by_source {
        let guarantee = min_per_source.min(indices.len()).min(remaining_budget);
        indices.shuffle(&mut rng);
        selected_indices.extend_from_slice(&indices[..guarantee]);
        source_counts.insert(source, guarantee);
        remaining_budget = remaining_budget.saturating_sub(guarantee);
    }

    // Phase 2: distribute remaining budget proportionally
    if remaining_budget > 0 {
        let total_remaining_available: usize = by_source
            .iter()
            .map(|(src, indices)| indices.len().saturating_sub(*source_counts.get(src).unwrap_or(&0)))
            .sum();

        if total_remaining_available > 0 {
            let mut extra_allocated = 0usize;
            let sources: Vec<&str> = by_source.keys().copied().collect();

            for (i, source) in sources.iter().enumerate() {
                let indices = &by_source[source];
                let already = *source_counts.get(source).unwrap_or(&0);
                let available = indices.len().saturating_sub(already);

                if available == 0 {
                    continue;
                }

                // Last source gets whatever is left to avoid rounding gaps
                let extra = if i == num_sources - 1 {
                    remaining_budget.saturating_sub(extra_allocated)
                } else {
                    let proportion = available as f64 / total_remaining_available as f64;
                    ((remaining_budget as f64 * proportion).round() as usize).min(available)
                };

                let extra = extra.min(available);
                selected_indices.extend_from_slice(&indices[already..already + extra]);
                extra_allocated += extra;
            }
        }
    }

    // Deduplicate (shouldn't happen but defensive)
    selected_indices.sort_unstable();
    selected_indices.dedup();

    // Build result sorted by external_id for determinism
    let mut result: Vec<NormalizedEntity> = selected_indices
        .into_iter()
        .map(|i| entities[i].clone())
        .collect();
    result.sort_by(|a, b| a.external_id.cmp(&b.external_id));
    result
}

// ---------------------------------------------------------------------------
// Scoring (port of compute_score from _metrics.py)
// ---------------------------------------------------------------------------

/// Compute composite score with z-score normalization.
fn compute_score(
    metrics: &HashMap<String, f64>,
    weights: &HashMap<String, f64>,
    all_candidates: &[HashMap<String, f64>],
) -> f64 {
    // Metrics where lower is better
    let negate: HashSet<&str> = ["conflict_rate", "entropy"].iter().copied().collect();

    if all_candidates.is_empty() {
        let mut score = 0.0;
        for (key, w) in weights {
            let val = metrics.get(key).copied().unwrap_or(0.0);
            if negate.contains(key.as_str()) {
                score += w * (-val);
            } else {
                score += w * val;
            }
        }
        return score;
    }

    // Compute mean/std for each metric across all candidates
    let mut stats: HashMap<String, (f64, f64)> = HashMap::new();
    for key in weights.keys() {
        let values: Vec<f64> = all_candidates
            .iter()
            .map(|c| c.get(key).copied().unwrap_or(0.0))
            .collect();
        let n = values.len() as f64;
        let mean = values.iter().sum::<f64>() / n;
        let variance = values.iter().map(|v| (v - mean).powi(2)).sum::<f64>() / n;
        let std = if variance > 0.0 { variance.sqrt() } else { 1.0 };
        stats.insert(key.clone(), (mean, std));
    }

    let mut score = 0.0;
    for (key, w) in weights {
        let val = metrics.get(key).copied().unwrap_or(0.0);
        let (mean, std) = stats.get(key).copied().unwrap_or((0.0, 1.0));
        let mut z = (val - mean) / std;
        if negate.contains(key.as_str()) {
            z = -z;
        }
        score += w * z;
    }

    score
}

// ---------------------------------------------------------------------------
// Mutation generators (port of _mutations.py)
// ---------------------------------------------------------------------------

/// Generate candidate mutations for a plan at a given tier.
fn generate_candidates(plan: &IdentityPlan, tier: u8) -> Vec<(Mutation, IdentityPlan)> {
    let mut candidates = Vec::new();
    match tier {
        1 => {
            candidates.extend(tier1_add_normalizers(plan));
            candidates.extend(tier1_case_insensitive(plan));
            candidates.extend(tier1_add_blocking_keys(plan));
            candidates.extend(tier1_enable_graph_clustering(plan));
        }
        2 => {
            candidates.extend(tier2_adjust_fs_probabilities(plan));
            candidates.extend(tier2_adjust_match_threshold(plan));
            candidates.extend(tier2_adjust_clustering_params(plan));
        }
        3 => {
            candidates.extend(tier3_remove_blocking_keys(plan));
            candidates.extend(tier3_large_threshold_swings(plan));
            candidates.extend(tier3_switch_comparators(plan));
        }
        _ => {}
    }
    candidates
}

// -- Tier 1: Safe --

fn infer_normalizer(field: &str) -> Option<&'static str> {
    let lower = field.to_lowercase();
    if lower.contains("email") {
        Some("email")
    } else if lower.contains("phone") || lower.contains("tel") || lower.contains("mobile") {
        Some("phone")
    } else if lower.contains("first_name") || lower.contains("fname") || lower.contains("given") {
        Some("nickname")
    } else if lower.contains("last_name") || lower.contains("surname") || lower.contains("lname") {
        Some("name")
    } else if lower.contains("company") || lower.contains("org") {
        Some("generic")
    } else if lower.contains("name") {
        Some("name")
    } else if lower.contains("domain") || lower.contains("website") {
        Some("domain")
    } else {
        None
    }
}

fn tier1_add_normalizers(plan: &IdentityPlan) -> Vec<(Mutation, IdentityPlan)> {
    let mut results = Vec::new();
    for (i, rule) in plan.match_graph.rules.iter().enumerate() {
        let (field, has_normalizer) = match &rule.rule_type {
            CompiledRuleType::Exact {
                field, normalizer, ..
            } => (field.clone(), normalizer.is_some()),
            CompiledRuleType::Similarity {
                field, normalizer, ..
            } => (field.clone(), normalizer.is_some()),
            _ => continue,
        };

        if has_normalizer {
            continue;
        }

        if let Some(norm) = infer_normalizer(&field) {
            let mut new_plan = plan.clone();
            let new_rule = &mut new_plan.match_graph.rules[i];
            match &mut new_rule.rule_type {
                CompiledRuleType::Exact { normalizer, .. } => {
                    *normalizer = Some(norm.to_string());
                }
                CompiledRuleType::Similarity { normalizer, .. } => {
                    *normalizer = Some(norm.to_string());
                }
                _ => continue,
            }

            // Also add normalizer to FS field if present
            if let Some(ref mut fs) = new_plan.decision.fellegi_sunter {
                for ff in &mut fs.fields {
                    if ff.name == field && ff.normalizer.is_none() {
                        ff.normalizer = Some(norm.to_string());
                    }
                }
            }

            results.push((
                Mutation {
                    tier: 1,
                    category: "normalizer".into(),
                    description: format!("Add {} normalizer to rule '{}'", norm, rule.name),
                    field_path: format!("rules.{}.normalizer", i),
                },
                new_plan,
            ));
        }
    }
    results
}

fn tier1_case_insensitive(plan: &IdentityPlan) -> Vec<(Mutation, IdentityPlan)> {
    let mut results = Vec::new();
    for (i, rule) in plan.match_graph.rules.iter().enumerate() {
        if let CompiledRuleType::Exact {
            case_insensitive, ..
        } = &rule.rule_type
        {
            if *case_insensitive {
                continue;
            }
            let mut new_plan = plan.clone();
            if let CompiledRuleType::Exact {
                case_insensitive, ..
            } = &mut new_plan.match_graph.rules[i].rule_type
            {
                *case_insensitive = true;
            }
            results.push((
                Mutation {
                    tier: 1,
                    category: "case_insensitive".into(),
                    description: format!("Enable case-insensitive on rule '{}'", rule.name),
                    field_path: format!("rules.{}.case_insensitive", i),
                },
                new_plan,
            ));
        }
    }
    results
}

fn tier1_add_blocking_keys(plan: &IdentityPlan) -> Vec<(Mutation, IdentityPlan)> {
    // Find fields used in rules but not currently blocked
    let existing_blocked: HashSet<String> = plan
        .blocking
        .keys
        .iter()
        .flat_map(|k| k.fields.iter().cloned())
        .collect();

    let mut results = Vec::new();
    for rule in &plan.match_graph.rules {
        let field = match &rule.rule_type {
            CompiledRuleType::Exact { field, .. } => field.clone(),
            CompiledRuleType::Similarity { field, .. } => field.clone(),
            _ => continue,
        };

        if existing_blocked.contains(&field) {
            continue;
        }

        // Only add single-field blocks for high-discriminator fields
        let lower = field.to_lowercase();
        if lower.contains("email") || lower.contains("phone") || lower.contains("ssn") {
            let mut new_plan = plan.clone();
            let transformation = if lower.contains("email") {
                cannon_common::ir::BlockingTransformation::NormalizeEmail
            } else if lower.contains("phone") {
                cannon_common::ir::BlockingTransformation::NormalizePhone
            } else {
                cannon_common::ir::BlockingTransformation::Lowercase
            };
            new_plan.blocking.keys.push(BlockingKey {
                fields: vec![field.clone()],
                transformations: vec![transformation],
            });
            results.push((
                Mutation {
                    tier: 1,
                    category: "blocking".into(),
                    description: format!("Add blocking key [{}]", field),
                    field_path: format!(
                        "blocking.keys.{}",
                        new_plan.blocking.keys.len() - 1
                    ),
                },
                new_plan,
            ));
        }
    }
    results
}

fn tier1_enable_graph_clustering(plan: &IdentityPlan) -> Vec<(Mutation, IdentityPlan)> {
    // Only if clustering is not already graph
    if let Some(ref cc) = plan.decision.clustering {
        if cc.method == ClusteringMethod::Graph {
            return Vec::new();
        }
    }

    let mut new_plan = plan.clone();
    new_plan.decision.clustering = Some(CompiledClusteringConfig {
        method: ClusteringMethod::Graph,
        min_edge_weight: 0.6,
        min_cluster_coherence: 0.6,
        max_cluster_size: 100,
    });

    vec![(
        Mutation {
            tier: 1,
            category: "clustering".into(),
            description: "Enable graph clustering with coherence checks".into(),
            field_path: "decision.clustering".into(),
        },
        new_plan,
    )]
}

// -- Tier 2: Moderate --

fn tier2_adjust_fs_probabilities(plan: &IdentityPlan) -> Vec<(Mutation, IdentityPlan)> {
    let mut results = Vec::new();
    let fs = match &plan.decision.fellegi_sunter {
        Some(fs) => fs,
        None => return results,
    };

    for (i, field) in fs.fields.iter().enumerate() {
        // m_probability adjustments
        for delta in [0.05, -0.05] {
            let new_m = field.m_probability + delta;
            if !(0.5..=0.999).contains(&new_m) {
                continue;
            }
            let mut new_plan = plan.clone();
            let new_fs = new_plan.decision.fellegi_sunter.as_mut().unwrap();
            let ff = &mut new_fs.fields[i];
            ff.m_probability = new_m;
            // Clamp before ln() to prevent ln(0) if range guards are ever relaxed
            let safe_m = ff.m_probability.clamp(1e-4, 1.0 - 1e-4);
            let safe_u = ff.u_probability.clamp(1e-4, 1.0 - 1e-4);
            ff.w_agree = (safe_m / safe_u).ln() * ff.weight;
            ff.w_disagree = ((1.0 - safe_m) / (1.0 - safe_u)).ln() * ff.weight;
            recompute_fs_bounds(new_fs);

            results.push((
                Mutation {
                    tier: 2,
                    category: "fs_probability".into(),
                    description: format!(
                        "Adjust {}.m_probability by {:+.02} to {:.3}",
                        field.name, delta, new_m
                    ),
                    field_path: format!("decision.fellegi_sunter.fields.{}.m_probability", i),
                },
                new_plan,
            ));
        }

        // u_probability adjustments
        for delta in [0.005, -0.005] {
            let new_u = field.u_probability + delta;
            if !(0.0001..=0.5).contains(&new_u) {
                continue;
            }
            let mut new_plan = plan.clone();
            let new_fs = new_plan.decision.fellegi_sunter.as_mut().unwrap();
            let ff = &mut new_fs.fields[i];
            ff.u_probability = new_u;
            // Clamp before ln() to prevent ln(0) if range guards are ever relaxed
            let safe_m = ff.m_probability.clamp(1e-4, 1.0 - 1e-4);
            let safe_u = ff.u_probability.clamp(1e-4, 1.0 - 1e-4);
            ff.w_agree = (safe_m / safe_u).ln() * ff.weight;
            ff.w_disagree = ((1.0 - safe_m) / (1.0 - safe_u)).ln() * ff.weight;
            recompute_fs_bounds(new_fs);

            results.push((
                Mutation {
                    tier: 2,
                    category: "fs_probability".into(),
                    description: format!(
                        "Adjust {}.u_probability by {:+.4} to {:.4}",
                        field.name, delta, new_u
                    ),
                    field_path: format!("decision.fellegi_sunter.fields.{}.u_probability", i),
                },
                new_plan,
            ));
        }
    }

    results
}

fn tier2_adjust_match_threshold(plan: &IdentityPlan) -> Vec<(Mutation, IdentityPlan)> {
    let mut results = Vec::new();

    // Adjust decision match threshold
    for delta in [0.05, -0.05] {
        let new_t = plan.decision.match_threshold + delta;
        if !(0.1..=1.0).contains(&new_t) {
            continue;
        }
        let mut new_plan = plan.clone();
        new_plan.decision.match_threshold = new_t;
        results.push((
            Mutation {
                tier: 2,
                category: "threshold".into(),
                description: format!(
                    "Adjust decision match_threshold by {:+.2} to {:.2}",
                    delta, new_t
                ),
                field_path: "decision.match_threshold".into(),
            },
            new_plan,
        ));
    }

    // Adjust FS match threshold
    if let Some(ref fs) = plan.decision.fellegi_sunter {
        for delta in [1.0, -1.0, 2.0, -2.0] {
            let new_t = fs.match_threshold + delta;
            if new_t < 0.0 {
                continue;
            }
            let mut new_plan = plan.clone();
            let new_fs = new_plan.decision.fellegi_sunter.as_mut().unwrap();
            new_fs.match_threshold = new_t;
            results.push((
                Mutation {
                    tier: 2,
                    category: "threshold".into(),
                    description: format!(
                        "Adjust FS match_threshold by {:+.1} to {:.1}",
                        delta, new_t
                    ),
                    field_path: "decision.fellegi_sunter.match_threshold".into(),
                },
                new_plan,
            ));
        }
    }

    results
}

fn tier2_adjust_clustering_params(plan: &IdentityPlan) -> Vec<(Mutation, IdentityPlan)> {
    let mut results = Vec::new();
    let cc = match &plan.decision.clustering {
        Some(cc) => cc,
        None => return results,
    };

    for delta in [0.05, -0.05] {
        // min_edge_weight
        let new_w = cc.min_edge_weight + delta;
        if (0.1..=0.95).contains(&new_w) {
            let mut new_plan = plan.clone();
            new_plan
                .decision
                .clustering
                .as_mut()
                .unwrap()
                .min_edge_weight = new_w;
            results.push((
                Mutation {
                    tier: 2,
                    category: "clustering".into(),
                    description: format!(
                        "Adjust min_edge_weight by {:+.2} to {:.2}",
                        delta, new_w
                    ),
                    field_path: "decision.clustering.min_edge_weight".into(),
                },
                new_plan,
            ));
        }

        // min_cluster_coherence
        let new_c = cc.min_cluster_coherence + delta;
        if (0.1..=0.95).contains(&new_c) {
            let mut new_plan = plan.clone();
            new_plan
                .decision
                .clustering
                .as_mut()
                .unwrap()
                .min_cluster_coherence = new_c;
            results.push((
                Mutation {
                    tier: 2,
                    category: "clustering".into(),
                    description: format!(
                        "Adjust min_cluster_coherence by {:+.2} to {:.2}",
                        delta, new_c
                    ),
                    field_path: "decision.clustering.min_cluster_coherence".into(),
                },
                new_plan,
            ));
        }
    }

    results
}

// -- Tier 3: Risky --

fn tier3_remove_blocking_keys(plan: &IdentityPlan) -> Vec<(Mutation, IdentityPlan)> {
    let mut results = Vec::new();
    if plan.blocking.keys.len() <= 1 {
        return results;
    }

    for i in 0..plan.blocking.keys.len() {
        let mut new_plan = plan.clone();
        let removed = new_plan.blocking.keys.remove(i);
        results.push((
            Mutation {
                tier: 3,
                category: "blocking".into(),
                description: format!(
                    "Remove blocking key [{}]",
                    removed.fields.join(", ")
                ),
                field_path: format!("blocking.keys.{}", i),
            },
            new_plan,
        ));
    }

    results
}

fn tier3_large_threshold_swings(plan: &IdentityPlan) -> Vec<(Mutation, IdentityPlan)> {
    let mut results = Vec::new();

    if let Some(ref fs) = plan.decision.fellegi_sunter {
        for delta in [3.0, -3.0, 5.0, -5.0] {
            let new_t = fs.match_threshold + delta;
            if new_t < 0.0 {
                continue;
            }
            let mut new_plan = plan.clone();
            new_plan
                .decision
                .fellegi_sunter
                .as_mut()
                .unwrap()
                .match_threshold = new_t;
            results.push((
                Mutation {
                    tier: 3,
                    category: "threshold".into(),
                    description: format!(
                        "Large FS threshold swing by {:+.1} to {:.1}",
                        delta, new_t
                    ),
                    field_path: "decision.fellegi_sunter.match_threshold".into(),
                },
                new_plan,
            ));
        }
    }

    results
}

fn tier3_switch_comparators(plan: &IdentityPlan) -> Vec<(Mutation, IdentityPlan)> {
    let mut results = Vec::new();

    if let Some(ref fs) = plan.decision.fellegi_sunter {
        let alternatives = [
            FsComparatorType::JaroWinkler,
            FsComparatorType::Levenshtein,
            FsComparatorType::Soundex,
        ];

        for (i, field) in fs.fields.iter().enumerate() {
            for alt in &alternatives {
                if &field.comparator == alt {
                    continue;
                }
                // Only switch similarity-type comparators (not exact on email/phone)
                if field.comparator == FsComparatorType::Exact {
                    continue;
                }
                let mut new_plan = plan.clone();
                let new_fs = new_plan.decision.fellegi_sunter.as_mut().unwrap();
                new_fs.fields[i].comparator = alt.clone();

                results.push((
                    Mutation {
                        tier: 3,
                        category: "comparator".into(),
                        description: format!(
                            "Switch {} comparator to {:?}",
                            field.name, alt
                        ),
                        field_path: format!("decision.fellegi_sunter.fields.{}.comparator", i),
                    },
                    new_plan,
                ));
            }
        }
    }

    results
}

// -- Helpers --

fn recompute_fs_bounds(fs: &mut cannon_common::ir::CompiledFellegiSunterPlan) {
    fs.max_composite = fs.fields.iter().map(|f| f.w_agree).sum();
    fs.min_composite = fs.fields.iter().map(|f| f.w_disagree).sum();
}

// ---------------------------------------------------------------------------
// Main autotune algorithm (port of autotune.py, 3-phase)
// ---------------------------------------------------------------------------

/// Run the 3-phase autotune optimization on a set of entities.
///
/// Applies stratified sampling before the mutation loop when the dataset
/// exceeds the auto-calculated (or user-specified) sample size. The winning
/// spec is returned for the caller to apply to the full dataset.
pub fn run(
    entities: &[NormalizedEntity],
    plan: &IdentityPlan,
    config: &AutoTuneConfig,
    labels: &[LabeledPair],
) -> AutoTuneResult {
    let total_entities = entities.len();

    // Stratified sampling: use a deterministic seed based on entity count
    // so the same dataset always produces the same sample.
    let seed = total_entities as u64;
    let sampled = sample_stratified(entities, config.sample_size, seed);
    let sampled_entities = sampled.len();

    if sampled_entities < total_entities {
        info!(
            total = total_entities,
            sampled = sampled_entities,
            "Sampled {} of {} entities for autotune",
            sampled_entities,
            total_entities
        );
    }

    // Run the mutation loop on the (possibly sampled) entities
    let entities = &sampled;

    let identity_fields = infer_identity_fields(plan);
    let resolver = OverrideResolver::new(Vec::new());

    // Build weight map
    let weights: HashMap<String, f64> = if !labels.is_empty() {
        [
            ("f1".into(), 0.7),
            ("conflict_rate".into(), 0.2),
            ("stability".into(), 0.1),
        ]
        .into()
    } else {
        [
            ("merge_rate".into(), config.merge_rate_weight),
            ("conflict_rate".into(), config.conflict_rate_weight),
            ("entropy".into(), config.entropy_weight),
            ("stability".into(), config.stability_weight),
        ]
        .into()
    };

    // Run baseline
    let engine = ReconciliationEngine::from_plan(plan);
    let (baseline_decisions, _telemetry) = engine.reconcile_with_telemetry(entities, &resolver);
    let baseline_clusters = build_clusters(&baseline_decisions, entities);
    let baseline_metrics = compute_metrics(
        &baseline_decisions,
        entities,
        &identity_fields,
        None,
        labels,
    );
    let baseline_score = compute_score(&baseline_metrics, &weights, &[]);

    info!(
        merge_rate = baseline_metrics.get("merge_rate").copied().unwrap_or(0.0),
        conflict_rate = baseline_metrics.get("conflict_rate").copied().unwrap_or(0.0),
        "AutoTune baseline computed"
    );

    let mut current_plan = plan.clone();
    let mut current_metrics = baseline_metrics.clone();
    let mut current_score = baseline_score;
    let mut accepted_mutations: Vec<MutationLogEntry> = Vec::new();
    let mut iterations_used = 0usize;
    let mut mutations_total = 0usize;

    // Phase 1: Tier 1 sweep - cumulative acceptance
    let mut tier1_accepted_this_round = true;
    while tier1_accepted_this_round && iterations_used < config.max_iterations {
        tier1_accepted_this_round = false;
        let candidates = generate_candidates(&current_plan, 1);
        mutations_total += candidates.len();

        for (mutation, candidate_plan) in candidates {
            iterations_used += 1;
            if iterations_used > config.max_iterations {
                break;
            }

            let cand_engine = ReconciliationEngine::from_plan(&candidate_plan);
            let (cand_decisions, _) = cand_engine.reconcile_with_telemetry(entities, &resolver);
            let cand_metrics = compute_metrics(
                &cand_decisions,
                entities,
                &identity_fields,
                Some(&baseline_clusters),
                labels,
            );
            let cand_score = compute_score(&cand_metrics, &weights, &[]);
            let score_delta = cand_score - current_score;
            let passes_conflict = cand_metrics
                .get("conflict_rate")
                .copied()
                .unwrap_or(1.0)
                <= config.max_conflict;

            let accepted = score_delta > 0.0 && passes_conflict;

            debug!(
                mutation = mutation.description,
                score_delta,
                accepted,
                "Tier 1 mutation evaluated"
            );

            accepted_mutations.push(MutationLogEntry {
                mutation: mutation.clone(),
                score_delta,
                accepted,
            });

            if accepted {
                current_plan = candidate_plan;
                current_metrics = cand_metrics;
                current_score = cand_score;
                tier1_accepted_this_round = true;
                break; // Re-generate candidates from new plan
            }
        }
    }

    let tier1_count = accepted_mutations.iter().filter(|m| m.accepted && m.mutation.tier == 1).count();
    info!(tier1_accepted = tier1_count, "Phase 1 complete");

    // Phase 2: Tier 2 greedy - keep single best per pass
    let mut tier2_pass_improved = true;
    let mut tier2_ever_improved = false;

    while tier2_pass_improved && iterations_used < config.max_iterations {
        tier2_pass_improved = false;
        let candidates = generate_candidates(&current_plan, 2);
        mutations_total += candidates.len();

        let mut best_delta = 0.0f64;
        let mut best_entry: Option<(IdentityPlan, HashMap<String, f64>, MutationLogEntry)> = None;

        for (mutation, candidate_plan) in candidates {
            iterations_used += 1;
            if iterations_used > config.max_iterations {
                break;
            }

            let cand_engine = ReconciliationEngine::from_plan(&candidate_plan);
            let (cand_decisions, _) = cand_engine.reconcile_with_telemetry(entities, &resolver);
            let cand_metrics = compute_metrics(
                &cand_decisions,
                entities,
                &identity_fields,
                Some(&baseline_clusters),
                labels,
            );
            let cand_score = compute_score(&cand_metrics, &weights, &[]);
            let score_delta = cand_score - current_score;
            let passes_conflict = cand_metrics
                .get("conflict_rate")
                .copied()
                .unwrap_or(1.0)
                <= config.max_conflict;

            let log_entry = MutationLogEntry {
                mutation: mutation.clone(),
                score_delta,
                accepted: false, // Updated below if best
            };

            if score_delta > best_delta && passes_conflict {
                best_delta = score_delta;
                best_entry = Some((candidate_plan, cand_metrics, log_entry));
            } else {
                accepted_mutations.push(log_entry);
            }
        }

        if let Some((best_plan, best_metrics, mut entry)) = best_entry {
            entry.accepted = true;
            debug!(
                mutation = entry.mutation.description,
                score_delta = best_delta,
                "Tier 2 best mutation accepted"
            );
            accepted_mutations.push(entry);
            current_plan = best_plan;
            current_metrics = best_metrics;
            current_score += best_delta;
            tier2_pass_improved = true;
            tier2_ever_improved = true;
        }
    }

    let tier2_count = accepted_mutations.iter().filter(|m| m.accepted && m.mutation.tier == 2).count();
    info!(tier2_accepted = tier2_count, "Phase 2 complete");

    // Phase 3: Tier 3 opportunistic - only if Tier 2 found nothing
    if !tier2_ever_improved && iterations_used < config.max_iterations {
        let candidates = generate_candidates(&current_plan, 3);
        mutations_total += candidates.len();

        for (mutation, candidate_plan) in candidates {
            iterations_used += 1;
            if iterations_used > config.max_iterations {
                break;
            }

            let cand_engine = ReconciliationEngine::from_plan(&candidate_plan);
            let (cand_decisions, _) = cand_engine.reconcile_with_telemetry(entities, &resolver);
            let cand_metrics = compute_metrics(
                &cand_decisions,
                entities,
                &identity_fields,
                Some(&baseline_clusters),
                labels,
            );
            let cand_score = compute_score(&cand_metrics, &weights, &[]);
            let score_delta = cand_score - current_score;
            let passes_conflict = cand_metrics
                .get("conflict_rate")
                .copied()
                .unwrap_or(1.0)
                <= config.max_conflict;

            let accepted = score_delta > 0.0 && passes_conflict;

            accepted_mutations.push(MutationLogEntry {
                mutation: mutation.clone(),
                score_delta,
                accepted,
            });

            if accepted {
                current_plan = candidate_plan;
                current_metrics = cand_metrics;
                current_score = cand_score;
            }
        }
    }

    let total_accepted = accepted_mutations.iter().filter(|m| m.accepted).count();
    let explanation = if total_accepted == 0 {
        "AutoTune found no improvements. The current plan is already well-calibrated.".into()
    } else {
        let descs: Vec<&str> = accepted_mutations
            .iter()
            .filter(|m| m.accepted)
            .map(|m| m.mutation.description.as_str())
            .collect();
        format!(
            "AutoTune accepted {} mutations: {}",
            total_accepted,
            descs.join("; ")
        )
    };

    info!(
        accepted = total_accepted,
        total = mutations_total,
        iterations = iterations_used,
        "AutoTune complete"
    );

    AutoTuneResult {
        best_plan: current_plan,
        original_plan: plan.clone(),
        metrics_before: baseline_metrics,
        metrics_after: current_metrics,
        mutations_accepted: accepted_mutations
            .into_iter()
            .filter(|m| m.accepted)
            .collect(),
        mutations_total,
        iterations_used,
        explanation,
        total_entities,
        sampled_entities,
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    fn make_test_entity(
        id: Uuid,
        external_id: &str,
        data: HashMap<String, String>,
    ) -> NormalizedEntity {
        NormalizedEntity {
            id,
            tenant_id: Uuid::nil(),
            external_id: external_id.to_string(),
            entity_type: "person".to_string(),
            data,
            source_name: "test".to_string(),
            valid_from: None,
            valid_to: None,
            last_updated: chrono::Utc::now(),
        }
    }

    fn make_test_plan() -> IdentityPlan {
        let fs_plan = CompiledFellegiSunterPlan {
            fields: vec![
                CompiledFsField {
                    name: "email".into(),
                    comparator: FsComparatorType::Exact,
                    weight: 2.0,
                    m_probability: 0.95,
                    u_probability: 0.001,
                    w_agree: (0.95_f64 / 0.001).ln() * 2.0,
                    w_disagree: ((1.0 - 0.95_f64) / (1.0 - 0.001)).ln() * 2.0,
                    normalizer: Some("email".into()),
                },
                CompiledFsField {
                    name: "last_name".into(),
                    comparator: FsComparatorType::JaroWinkler,
                    weight: 1.0,
                    m_probability: 0.88,
                    u_probability: 0.02,
                    w_agree: (0.88_f64 / 0.02).ln() * 1.0,
                    w_disagree: ((1.0 - 0.88_f64) / (1.0 - 0.02)).ln() * 1.0,
                    normalizer: Some("name".into()),
                },
            ],
            match_threshold: 10.0,
            possible_threshold: 4.0,
            non_match_threshold: -5.0,
            merge_threshold: None,
            max_composite: 0.0, // Will be recomputed
            min_composite: 0.0,
            em_training: None,
        };

        let mut plan = IdentityPlan {
            entity: "person".into(),
            plan_hash: "test".into(),
            identity_version: "v1".into(),
            sources: vec![ResolvedSource {
                name: "test".into(),
                adapter: "database".into(),
                location: "external_entities".into(),
                id_field: "external_id".into(),
                field_mappings: HashMap::new(),
                temporal: None,
                freshness: None,
                schema: None,
                sampling: None,
                tags: Vec::new(),
                pii_fields: Vec::new(),
                reliability: 1.0,
            }],
            blocking: CompiledBlocking {
                strategy: BlockingStrategy::Composite,
                keys: vec![BlockingKey {
                    fields: vec!["email".into()],
                    transformations: vec![BlockingTransformation::NormalizeEmail],
                }],
                fallback_sample_rate: None,
                lsh_bands: None,
                lsh_rows: None,
                neighborhood_window: None,
                sort_fields: Vec::new(),
            },
            match_graph: MatchGraph {
                rules: vec![
                    CompiledRule {
                        name: "email_match".into(),
                        rule_type: CompiledRuleType::Exact {
                            field: "email".into(),
                            weight: 2.0,
                            case_insensitive: true,
                            normalize: true,
                            normalizer: Some("email".into()),
                        },
                    },
                    CompiledRule {
                        name: "last_name_match".into(),
                        rule_type: CompiledRuleType::Similarity {
                            field: "last_name".into(),
                            algorithm: SimilarityAlgorithm::JaroWinkler,
                            threshold: 0.85,
                            weight: 1.0,
                            normalizer: Some("name".into()),
                        },
                    },
                ],
                leaf_count: 2,
            },
            survivorship: SurvivorshipPolicy {
                default_strategy: SurvivorshipStrategy::MostComplete,
                field_policies: HashMap::new(),
            },
            decision: DecisionConfig {
                match_threshold: 0.9,
                review_threshold: Some(0.7),
                reject_threshold: None,
                conflict_strategy: ConflictStrategy::PreferHighConfidence,
                review_webhook: None,
                scoring_method: ScoringMethod::FellegiSunter,
                ml_ensemble_config: None,
                custom_scoring_config: None,
                fellegi_sunter: Some(fs_plan),
                tie_breaking: Vec::new(),
                clustering: None,
                min_total_weight: 0.0,
                allow_single_field: vec!["email_match".into()],
            },
            reference_identifiers: Vec::new(),
            relations: Vec::new(),
            exclusions: Vec::new(),
            compliance: None,
            hierarchy: None,
            audit: None,
            metadata: None,
            governance: None,
        };

        // Recompute bounds
        if let Some(ref mut fs) = plan.decision.fellegi_sunter {
            recompute_fs_bounds(fs);
        }

        plan
    }

    #[test]
    fn test_conflict_rate_no_conflicts() {
        let id1 = Uuid::new_v4();
        let id2 = Uuid::new_v4();
        let entities = vec![
            make_test_entity(id1, "1", [("email".into(), "a@b.com".into())].into()),
            make_test_entity(id2, "2", [("email".into(), "a@b.com".into())].into()),
        ];
        let clusters = vec![vec![id1, id2]];
        let cr = conflict_rate(&clusters, &entities, &["email".into()]);
        assert!((cr - 0.0).abs() < 0.001);
    }

    #[test]
    fn test_conflict_rate_with_conflict() {
        let id1 = Uuid::new_v4();
        let id2 = Uuid::new_v4();
        let entities = vec![
            make_test_entity(id1, "1", [("email".into(), "a@b.com".into())].into()),
            make_test_entity(id2, "2", [("email".into(), "c@d.com".into())].into()),
        ];
        let clusters = vec![vec![id1, id2]];
        let cr = conflict_rate(&clusters, &entities, &["email".into()]);
        assert!((cr - 1.0).abs() < 0.001);
    }

    #[test]
    fn test_entropy_uniform() {
        let id1 = Uuid::new_v4();
        let id2 = Uuid::new_v4();
        let entities = vec![
            make_test_entity(id1, "1", [("email".into(), "a@b.com".into())].into()),
            make_test_entity(id2, "2", [("email".into(), "c@d.com".into())].into()),
        ];
        let clusters = vec![vec![id1, id2]];
        let e = identifier_entropy(&clusters, &entities, &["email".into()]);
        assert!(e > 0.0); // Two different values -> positive entropy
        assert!((e - 1.0).abs() < 0.001); // log2(2) = 1.0
    }

    #[test]
    fn test_entropy_single_value() {
        let id1 = Uuid::new_v4();
        let id2 = Uuid::new_v4();
        let entities = vec![
            make_test_entity(id1, "1", [("email".into(), "a@b.com".into())].into()),
            make_test_entity(id2, "2", [("email".into(), "a@b.com".into())].into()),
        ];
        let clusters = vec![vec![id1, id2]];
        let e = identifier_entropy(&clusters, &entities, &["email".into()]);
        assert!((e - 0.0).abs() < 0.001); // Same value -> 0 entropy
    }

    #[test]
    fn test_ari_identical() {
        let ids: Vec<Uuid> = (0..4).map(|_| Uuid::new_v4()).collect();
        let clusters = vec![vec![ids[0], ids[1]], vec![ids[2], ids[3]]];
        let ari = adjusted_rand_index(&clusters, &clusters);
        assert!((ari - 1.0).abs() < 0.001);
    }

    #[test]
    fn test_ari_different() {
        let ids: Vec<Uuid> = (0..4).map(|_| Uuid::new_v4()).collect();
        let clusters_a = vec![vec![ids[0], ids[1]], vec![ids[2], ids[3]]];
        let clusters_b = vec![vec![ids[0], ids[2]], vec![ids[1], ids[3]]];
        let ari = adjusted_rand_index(&clusters_a, &clusters_b);
        assert!(ari < 0.5); // Very different clusterings
    }

    #[test]
    fn test_compute_score_defaults() {
        let metrics: HashMap<String, f64> = [
            ("merge_rate".into(), 0.5),
            ("conflict_rate".into(), 0.02),
            ("entropy".into(), 0.3),
            ("stability".into(), 0.9),
        ]
        .into();

        let weights: HashMap<String, f64> = [
            ("merge_rate".into(), 0.3),
            ("conflict_rate".into(), 0.4),
            ("entropy".into(), 0.1),
            ("stability".into(), 0.2),
        ]
        .into();

        let score = compute_score(&metrics, &weights, &[]);
        // 0.3*0.5 + 0.4*(-0.02) + 0.1*(-0.3) + 0.2*0.9
        // = 0.15 - 0.008 - 0.03 + 0.18 = 0.292
        assert!((score - 0.292).abs() < 0.001);
    }

    #[test]
    fn test_f1_computation() {
        let id1 = Uuid::new_v4();
        let id2 = Uuid::new_v4();
        let id3 = Uuid::new_v4();

        let decisions = vec![MatchDecision {
            entity_a_id: id1,
            entity_b_id: id2,
            decision: crate::engine::types::Decision::Merge,
            confidence: 0.95,
            posterior_probability: Some(0.95),
            rule_results: Vec::new(),
            matched_on: Vec::new(),
            blocking_key: None,
            temporal_adjustment: None,
        }];

        let labels = vec![
            LabeledPair {
                entity_a_id: id1,
                entity_b_id: id2,
                is_match: true,
            },
            LabeledPair {
                entity_a_id: id1,
                entity_b_id: id3,
                is_match: false,
            },
        ];

        let (f1, precision, recall) = compute_f1_from_labels(&decisions, &labels);
        assert!((precision - 1.0).abs() < 0.001); // 1 TP, 0 FP
        assert!((recall - 1.0).abs() < 0.001); // 1 TP, 0 FN
        assert!((f1 - 1.0).abs() < 0.001);
    }

    #[test]
    fn test_generate_tier1_candidates() {
        let plan = make_test_plan();
        let candidates = generate_candidates(&plan, 1);
        // Should at least have graph clustering mutation (since clustering is None)
        assert!(!candidates.is_empty());
    }

    #[test]
    fn test_generate_tier2_candidates() {
        let plan = make_test_plan();
        let candidates = generate_candidates(&plan, 2);
        // Should have FS probability adjustments + threshold adjustments
        assert!(!candidates.is_empty());
    }

    #[test]
    fn test_generate_tier3_candidates() {
        let plan = make_test_plan();
        let candidates = generate_candidates(&plan, 3);
        // Should have comparator switches for last_name (jaro_winkler -> levenshtein, soundex)
        assert!(!candidates.is_empty());
    }

    #[test]
    fn test_infer_identity_fields() {
        let plan = make_test_plan();
        let fields = infer_identity_fields(&plan);
        assert!(fields.contains(&"email".to_string()));
        assert!(fields.contains(&"last_name".to_string()));
    }

    #[test]
    fn test_infer_normalizer() {
        assert_eq!(infer_normalizer("email"), Some("email"));
        assert_eq!(infer_normalizer("phone_number"), Some("phone"));
        assert_eq!(infer_normalizer("first_name"), Some("nickname"));
        assert_eq!(infer_normalizer("last_name"), Some("name"));
        assert_eq!(infer_normalizer("company_name"), Some("generic"));
        assert_eq!(infer_normalizer("random_field"), None);
    }

    #[test]
    fn test_recompute_fs_bounds() {
        let mut fs = CompiledFellegiSunterPlan {
            fields: vec![
                CompiledFsField {
                    name: "a".into(),
                    comparator: FsComparatorType::Exact,
                    weight: 1.0,
                    m_probability: 0.9,
                    u_probability: 0.1,
                    w_agree: 2.0,
                    w_disagree: -1.0,
                    normalizer: None,
                },
                CompiledFsField {
                    name: "b".into(),
                    comparator: FsComparatorType::Exact,
                    weight: 1.0,
                    m_probability: 0.8,
                    u_probability: 0.2,
                    w_agree: 1.5,
                    w_disagree: -0.5,
                    normalizer: None,
                },
            ],
            match_threshold: 5.0,
            possible_threshold: 2.0,
            non_match_threshold: -3.0,
            merge_threshold: None,
            max_composite: 0.0,
            min_composite: 0.0,
            em_training: None,
        };

        recompute_fs_bounds(&mut fs);
        assert!((fs.max_composite - 3.5).abs() < 0.001);
        assert!((fs.min_composite - (-1.5)).abs() < 0.001);
    }

    #[test]
    fn test_autotune_run_no_entities() {
        let plan = make_test_plan();
        let config = AutoTuneConfig::default();
        let result = run(&[], &plan, &config, &[]);
        // With zero entities, mutations are evaluated but none improve
        assert!(result.mutations_accepted.is_empty());
        assert_eq!(result.total_entities, 0);
        assert_eq!(result.sampled_entities, 0);
    }

    // -- Stratified sampling tests --

    fn make_source_entities(source: &str, count: usize, offset: usize) -> Vec<NormalizedEntity> {
        (0..count)
            .map(|i| NormalizedEntity {
                id: Uuid::new_v4(),
                tenant_id: Uuid::nil(),
                external_id: format!("{}-{}", source, offset + i),
                entity_type: "person".to_string(),
                data: [("email".into(), format!("{}@test.com", offset + i))].into(),
                source_name: source.to_string(),
                valid_from: None,
                valid_to: None,
                last_updated: chrono::Utc::now(),
            })
            .collect()
    }

    #[test]
    fn test_sampling_no_sample_for_small_datasets() {
        let entities = make_source_entities("crm", 100, 0);
        let result = sample_stratified(&entities, None, 42);
        assert_eq!(result.len(), 100);
    }

    #[test]
    fn test_sampling_preserves_source_distribution() {
        // 6000 entities: 4000 from source_a, 2000 from source_b
        let mut entities = make_source_entities("source_a", 4000, 0);
        entities.extend(make_source_entities("source_b", 2000, 4000));

        let result = sample_stratified(&entities, None, 42);
        // auto_sample_size(6000) = 5000
        assert_eq!(result.len(), 5000);

        let count_a = result.iter().filter(|e| e.source_name == "source_a").count();
        let count_b = result.iter().filter(|e| e.source_name == "source_b").count();

        // Both sources should be represented
        assert!(count_a > 0);
        assert!(count_b > 0);

        // Proportions should be roughly preserved (a ~67%, b ~33%)
        // After min 100 guarantee, remaining 4800 distributed proportionally
        let ratio_a = count_a as f64 / result.len() as f64;
        assert!(ratio_a > 0.5 && ratio_a < 0.85, "source_a ratio {:.2} out of range", ratio_a);
    }

    #[test]
    fn test_sampling_min_per_source_guaranteed() {
        // 3 sources: one large (5500), two small (150, 150)
        let mut entities = make_source_entities("big", 5500, 0);
        entities.extend(make_source_entities("small_a", 150, 5500));
        entities.extend(make_source_entities("small_b", 150, 5650));
        // total = 5800, auto_sample_size = 5000

        let result = sample_stratified(&entities, None, 42);
        assert_eq!(result.len(), 5000);

        let count_small_a = result.iter().filter(|e| e.source_name == "small_a").count();
        let count_small_b = result.iter().filter(|e| e.source_name == "small_b").count();

        // Min 100 per source guaranteed
        assert!(count_small_a >= 100, "small_a only got {}", count_small_a);
        assert!(count_small_b >= 100, "small_b only got {}", count_small_b);
    }

    #[test]
    fn test_sampling_deterministic() {
        let mut entities = make_source_entities("crm", 3000, 0);
        entities.extend(make_source_entities("erp", 3000, 3000));

        let result1 = sample_stratified(&entities, None, 42);
        let result2 = sample_stratified(&entities, None, 42);

        assert_eq!(result1.len(), result2.len());
        for (a, b) in result1.iter().zip(result2.iter()) {
            assert_eq!(a.external_id, b.external_id);
        }
    }

    #[test]
    fn test_sampling_explicit_size() {
        let entities = make_source_entities("crm", 10000, 0);
        let result = sample_stratified(&entities, Some(500), 42);
        assert_eq!(result.len(), 500);
    }

    #[test]
    fn test_auto_sample_size_tiers() {
        assert_eq!(auto_sample_size(100), 100);
        assert_eq!(auto_sample_size(5000), 5000);
        assert_eq!(auto_sample_size(5001), 5000);
        assert_eq!(auto_sample_size(50000), 5000);
        assert_eq!(auto_sample_size(50001), 10000);
        assert_eq!(auto_sample_size(500000), 10000);
        assert_eq!(auto_sample_size(500001), 20000);
        assert_eq!(auto_sample_size(1000000), 20000);
    }
}
