//! Fellegi-Sunter probabilistic matching: EM training and sensitivity analysis.
//!
//! This module provides unsupervised estimation of m/u probabilities via the
//! Expectation-Maximization algorithm, plus sensitivity analysis for understanding
//! how parameter changes affect match decisions.

use cannon_common::ir::{CompiledFsField, CompiledFellegiSunterPlan};
use super::types::NormalizedEntity;
use std::collections::HashMap;

// ============================================================================
// EM Training
// ============================================================================

/// Configuration for the Expectation-Maximization trainer.
pub struct EmTrainer {
    pub max_iterations: usize,
    pub convergence_threshold: f64,
}

/// Result of EM training with estimated probabilities.
#[derive(Debug, Clone)]
pub struct EmResult {
    pub estimated_m: HashMap<String, f64>,
    pub estimated_u: HashMap<String, f64>,
    pub iterations: usize,
    pub converged: bool,
    pub log_likelihood_history: Vec<f64>,
}

impl EmTrainer {
    pub fn new(max_iterations: usize, convergence_threshold: f64) -> Self {
        Self { max_iterations, convergence_threshold }
    }

    /// Run unsupervised EM to estimate m/u probabilities from candidate pairs.
    ///
    /// Algorithm:
    /// 1. Initialize m/u from spec values (or defaults m=0.9, u=0.1)
    /// 2. E-step: compute P(match|agreement) for each pair using current m/u
    /// 3. M-step: re-estimate m/u from weighted agreement counts
    /// 4. Repeat until convergence or max iterations
    pub fn train(
        &self,
        pairs: &[(NormalizedEntity, NormalizedEntity)],
        fields: &[CompiledFsField],
    ) -> EmResult {
        if pairs.is_empty() || fields.is_empty() {
            return EmResult {
                estimated_m: HashMap::new(),
                estimated_u: HashMap::new(),
                iterations: 0,
                converged: true,
                log_likelihood_history: vec![],
            };
        }

        let n_fields = fields.len();
        let n_pairs = pairs.len();

        // Initialize m and u from the compiled field values (spec priors)
        let mut m: Vec<f64> = fields.iter().map(|f| f.m_probability).collect();
        let mut u: Vec<f64> = fields.iter().map(|f| f.u_probability).collect();

        // Prior probability of a random pair being a true match
        let mut p_match = 0.1_f64; // conservative initial estimate

        // Precompute agreement scores for each pair and field
        let agreement_matrix: Vec<Vec<f64>> = pairs.iter()
            .map(|(a, b)| {
                fields.iter().map(|field| {
                    Self::compute_agreement(a, b, &field.name)
                }).collect()
            })
            .collect();

        let mut log_likelihood_history = Vec::new();

        let mut iteration = 0;
        let mut converged = false;

        while iteration < self.max_iterations {
            // E-step: compute posterior P(match|agreement pattern) for each pair
            let mut weights: Vec<f64> = Vec::with_capacity(n_pairs);
            let mut total_log_likelihood = 0.0_f64;

            for pair_agreements in &agreement_matrix {
                // P(agreement | match) = product of m_j^a_j * (1-m_j)^(1-a_j)
                let mut log_p_agree_match = 0.0_f64;
                let mut log_p_agree_nonmatch = 0.0_f64;

                for j in 0..n_fields {
                    let a = pair_agreements[j]; // agreement score [0,1]
                    let mj = m[j].clamp(1e-10, 1.0 - 1e-10);
                    let uj = u[j].clamp(1e-10, 1.0 - 1e-10);

                    // For continuous agreement scores, interpolate
                    log_p_agree_match += a * mj.ln() + (1.0 - a) * (1.0 - mj).ln();
                    log_p_agree_nonmatch += a * uj.ln() + (1.0 - a) * (1.0 - uj).ln();
                }

                let p_m = p_match.clamp(1e-10, 1.0 - 1e-10);
                let log_numerator = log_p_agree_match + p_m.ln();
                let log_denominator_match = log_p_agree_match + p_m.ln();
                let log_denominator_nonmatch = log_p_agree_nonmatch + (1.0 - p_m).ln();

                // Log-sum-exp for numerical stability
                let max_log = log_denominator_match.max(log_denominator_nonmatch);
                let log_total = max_log + (
                    (log_denominator_match - max_log).exp() +
                    (log_denominator_nonmatch - max_log).exp()
                ).ln();

                let w = (log_numerator - log_total).exp().clamp(0.0, 1.0);
                weights.push(w);
                total_log_likelihood += log_total;
            }

            log_likelihood_history.push(total_log_likelihood);

            // Check convergence
            if log_likelihood_history.len() >= 2 {
                let prev = log_likelihood_history[log_likelihood_history.len() - 2];
                let delta = (total_log_likelihood - prev).abs();
                if delta < self.convergence_threshold {
                    converged = true;
                    iteration += 1;
                    break;
                }
            }

            // M-step: re-estimate m, u, and p_match
            let sum_w: f64 = weights.iter().sum();
            let sum_1_w: f64 = weights.iter().map(|w| 1.0 - w).sum();

            // Update p_match
            p_match = (sum_w / n_pairs as f64).clamp(0.01, 0.99);

            // Update m and u for each field
            for j in 0..n_fields {
                let mut numerator_m = 0.0_f64;
                let mut numerator_u = 0.0_f64;

                for i in 0..n_pairs {
                    let a = agreement_matrix[i][j];
                    numerator_m += weights[i] * a;
                    numerator_u += (1.0 - weights[i]) * a;
                }

                m[j] = if sum_w > 1e-10 {
                    (numerator_m / sum_w).clamp(0.01, 0.99)
                } else {
                    m[j]
                };
                u[j] = if sum_1_w > 1e-10 {
                    (numerator_u / sum_1_w).clamp(0.001, 0.99)
                } else {
                    u[j]
                };
            }

            iteration += 1;
        }

        let estimated_m: HashMap<String, f64> = fields.iter().enumerate()
            .map(|(j, f)| (f.name.clone(), m[j]))
            .collect();
        let estimated_u: HashMap<String, f64> = fields.iter().enumerate()
            .map(|(j, f)| (f.name.clone(), u[j]))
            .collect();

        EmResult {
            estimated_m,
            estimated_u,
            iterations: iteration,
            converged,
            log_likelihood_history,
        }
    }

    /// Compute agreement score between two entities on a given field.
    /// Returns 1.0 for exact match, 0.0 for missing/no match, or similarity score.
    fn compute_agreement(a: &NormalizedEntity, b: &NormalizedEntity, field_name: &str) -> f64 {
        match (a.data.get(field_name), b.data.get(field_name)) {
            (Some(va), Some(vb)) if !va.is_empty() && !vb.is_empty() => {
                if va == vb {
                    1.0
                } else {
                    // Use Jaro-Winkler as a reasonable default similarity for EM
                    strsim::jaro_winkler(va, vb)
                }
            }
            _ => 0.0,
        }
    }
}

/// Update a compiled FS plan's field probabilities and weights from EM results.
pub fn apply_em_results(plan: &mut CompiledFellegiSunterPlan, result: &EmResult) {
    for field in &mut plan.fields {
        if let Some(&new_m) = result.estimated_m.get(&field.name) {
            if let Some(&new_u) = result.estimated_u.get(&field.name) {
                field.m_probability = new_m;
                field.u_probability = new_u;
                field.w_agree = if new_u > 0.0 { (new_m / new_u).ln() * field.weight } else { f64::MAX };
                field.w_disagree = if (1.0 - new_u) > 0.0 {
                    ((1.0 - new_m) / (1.0 - new_u)).ln() * field.weight
                } else {
                    f64::MIN
                };
            }
        }
    }
    // Recompute composite bounds
    plan.max_composite = plan.fields.iter().map(|f| f.w_agree).sum();
    plan.min_composite = plan.fields.iter().map(|f| f.w_disagree).sum();
}

// ============================================================================
// Sensitivity Analysis (A12)
// ============================================================================

/// Per-field sensitivity metrics.
#[derive(Debug, Clone)]
pub struct FieldSensitivity {
    pub name: String,
    /// How match rate changes per unit change in m probability
    pub m_elasticity: f64,
    /// How match rate changes per unit change in u probability
    pub u_elasticity: f64,
    /// Match rate change if this field is removed entirely
    pub removal_impact: f64,
}

/// Threshold sensitivity metrics.
#[derive(Debug, Clone)]
pub struct ThresholdSensitivity {
    /// Number of entities within ±0.5 of match threshold
    pub entities_at_boundary: usize,
    /// Rate of match rate change per unit threshold change
    pub match_rate_per_unit_threshold: f64,
}

/// Complete sensitivity analysis report.
#[derive(Debug, Clone)]
pub struct SensitivityReport {
    pub fields: Vec<FieldSensitivity>,
    pub threshold_sensitivity: ThresholdSensitivity,
}

impl SensitivityReport {
    /// Run sensitivity analysis by scoring pairs and perturbing parameters.
    pub fn analyze(
        plan: &CompiledFellegiSunterPlan,
        pair_scores: &[f64],
    ) -> Self {
        let n = pair_scores.len() as f64;
        if n == 0.0 {
            return Self {
                fields: vec![],
                threshold_sensitivity: ThresholdSensitivity {
                    entities_at_boundary: 0,
                    match_rate_per_unit_threshold: 0.0,
                },
            };
        }

        let base_match_rate = pair_scores.iter()
            .filter(|&&s| s >= plan.match_threshold)
            .count() as f64 / n;

        // Per-field sensitivity: perturb m/u by ±10% and measure match rate change
        let perturbation = 0.10;
        let mut field_sensitivities = Vec::new();

        for field in &plan.fields {
            // m elasticity: increase m by 10%
            let m_up = (field.m_probability * (1.0 + perturbation)).min(0.999);
            let m_delta = m_up - field.m_probability;
            let w_agree_up = if field.u_probability > 0.0 {
                (m_up / field.u_probability).ln() * field.weight
            } else {
                field.w_agree
            };
            let composite_shift_m = w_agree_up - field.w_agree;
            let match_rate_m = pair_scores.iter()
                .filter(|&&s| s + composite_shift_m >= plan.match_threshold)
                .count() as f64 / n;
            let m_elasticity = if m_delta.abs() > 1e-10 {
                (match_rate_m - base_match_rate) / m_delta
            } else {
                0.0
            };

            // u elasticity: increase u by 10%
            let u_up = (field.u_probability * (1.0 + perturbation)).min(0.999);
            let u_delta = u_up - field.u_probability;
            let w_agree_u = if u_up > 0.0 {
                (field.m_probability / u_up).ln() * field.weight
            } else {
                field.w_agree
            };
            let composite_shift_u = w_agree_u - field.w_agree;
            let match_rate_u = pair_scores.iter()
                .filter(|&&s| s + composite_shift_u >= plan.match_threshold)
                .count() as f64 / n;
            let u_elasticity = if u_delta.abs() > 1e-10 {
                (match_rate_u - base_match_rate) / u_delta
            } else {
                0.0
            };

            // Removal impact: remove this field's contribution entirely
            let match_rate_removed = pair_scores.iter()
                .filter(|&&s| {
                    let adjusted = s - field.w_agree; // remove best-case contribution
                    adjusted >= plan.match_threshold
                })
                .count() as f64 / n;
            let removal_impact = match_rate_removed - base_match_rate;

            field_sensitivities.push(FieldSensitivity {
                name: field.name.clone(),
                m_elasticity,
                u_elasticity,
                removal_impact,
            });
        }

        // Threshold sensitivity
        let boundary_range = 0.5;
        let entities_at_boundary = pair_scores.iter()
            .filter(|&&s| (s - plan.match_threshold).abs() <= boundary_range)
            .count();

        // Match rate change per unit threshold shift
        let threshold_up = plan.match_threshold + 1.0;
        let match_rate_up = pair_scores.iter()
            .filter(|&&s| s >= threshold_up)
            .count() as f64 / n;
        let match_rate_per_unit_threshold = match_rate_up - base_match_rate;

        Self {
            fields: field_sensitivities,
            threshold_sensitivity: ThresholdSensitivity {
                entities_at_boundary,
                match_rate_per_unit_threshold,
            },
        }
    }
}

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

    fn make_entity(id: &str, data: Vec<(&str, &str)>) -> NormalizedEntity {
        let mut map = HashMap::new();
        for (k, v) in data {
            map.insert(k.to_string(), v.to_string());
        }
        NormalizedEntity {
            id: uuid::Uuid::new_v4(),
            tenant_id: uuid::Uuid::nil(),
            external_id: id.to_string(),
            entity_type: "test".to_string(),
            data: map,
            source_name: "test".to_string(),
            valid_from: None,
            valid_to: None,
            last_updated: chrono::Utc::now(),
        }
    }

    #[test]
    fn em_converges_with_identical_pairs() {
        let fields = vec![
            CompiledFsField {
                name: "email".to_string(),
                comparator: FsComparatorType::Exact,
                weight: 1.0,
                m_probability: 0.9,
                u_probability: 0.1,
                w_agree: (0.9_f64 / 0.1).ln(),
                w_disagree: (0.1_f64 / 0.9).ln(),
            },
        ];

        let pairs: Vec<(NormalizedEntity, NormalizedEntity)> = (0..50).map(|i| {
            let a = make_entity(&format!("a{}", i), vec![("email", "test@example.com")]);
            let b = make_entity(&format!("b{}", i), vec![("email", "test@example.com")]);
            (a, b)
        }).collect();

        let trainer = EmTrainer::new(25, 0.001);
        let result = trainer.train(&pairs, &fields);

        assert!(result.converged || result.iterations <= 25);
        // With all-matching pairs, m should be high
        let m_email = result.estimated_m["email"];
        assert!(m_email > 0.5, "m_email should be high for all-matching pairs, got {}", m_email);
    }

    #[test]
    fn em_handles_empty_input() {
        let fields = vec![];
        let pairs: Vec<(NormalizedEntity, NormalizedEntity)> = vec![];
        let trainer = EmTrainer::new(10, 0.001);
        let result = trainer.train(&pairs, &fields);
        assert!(result.converged);
        assert_eq!(result.iterations, 0);
    }

    #[test]
    fn sensitivity_report_detects_boundary_entities() {
        let plan = CompiledFellegiSunterPlan {
            fields: vec![
                CompiledFsField {
                    name: "name".to_string(),
                    comparator: FsComparatorType::JaroWinkler,
                    weight: 1.0,
                    m_probability: 0.9,
                    u_probability: 0.1,
                    w_agree: (0.9_f64 / 0.1).ln(),
                    w_disagree: (0.1_f64 / 0.9).ln(),
                },
            ],
            match_threshold: 5.0,
            possible_threshold: 2.0,
            non_match_threshold: -4.0,
            max_composite: (0.9_f64 / 0.1).ln(),
            min_composite: (0.1_f64 / 0.9).ln(),
            em_training: None,
        };

        // Scores: some near the threshold, some far
        let scores = vec![5.1, 4.8, 4.5, 3.0, 1.0, -2.0, 7.0, 5.5];
        let report = SensitivityReport::analyze(&plan, &scores);

        // 5.1, 4.8, 4.5, 5.5 are within ±0.5 of threshold 5.0
        assert!(report.threshold_sensitivity.entities_at_boundary > 0);
        assert_eq!(report.fields.len(), 1);
    }
}
