//! Active learning (supervised EM) tests.
//!
//! Tests for supervised training of Fellegi-Sunter m/u probabilities
//! using labeled pairs (feedback).

use cannon_common::ir::{
    CompiledFsField, CompiledFellegiSunterPlan, FsComparatorType,
};
use cannon_core::engine::fellegi_sunter::{EmTrainer, apply_em_results};
use cannon_core::engine::types::NormalizedEntity;
use std::collections::HashMap;
use uuid::Uuid;

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::new_v4(),
        tenant_id: Uuid::nil(),
        external_id: id.to_string(),
        entity_type: "customer".to_string(),
        data: map,
        source_name: "test".to_string(),
        valid_from: None,
        valid_to: None,
        last_updated: chrono::Utc::now(),
    }
}

fn test_fields() -> Vec<CompiledFsField> {
    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(),
            normalizer: None,
        },
        CompiledFsField {
            name: "name".to_string(),
            comparator: FsComparatorType::JaroWinkler,
            weight: 1.0,
            m_probability: 0.85,
            u_probability: 0.15,
            w_agree: (0.85_f64 / 0.15).ln(),
            w_disagree: (0.15_f64 / 0.85).ln(),
            normalizer: None,
        },
    ]
}

#[test]
fn test_supervised_training_updates_m_u() {
    let fields = test_fields();
    let trainer = EmTrainer::new(0, 0.0);

    // Create labeled pairs: matching pairs have identical emails
    let labeled: Vec<(NormalizedEntity, NormalizedEntity, bool)> = vec![
        // Match pairs: same email
        (
            make_entity("a1", vec![("email", "alice@example.com"), ("name", "Alice Smith")]),
            make_entity("a2", vec![("email", "alice@example.com"), ("name", "Alice S.")]),
            true,
        ),
        (
            make_entity("b1", vec![("email", "bob@example.com"), ("name", "Bob Jones")]),
            make_entity("b2", vec![("email", "bob@example.com"), ("name", "Bob Jones")]),
            true,
        ),
        // Non-match pairs: different emails
        (
            make_entity("c1", vec![("email", "carol@example.com"), ("name", "Carol White")]),
            make_entity("d1", vec![("email", "dave@example.com"), ("name", "Dave Brown")]),
            false,
        ),
        (
            make_entity("e1", vec![("email", "eve@example.com"), ("name", "Eve Green")]),
            make_entity("f1", vec![("email", "frank@example.com"), ("name", "Frank Black")]),
            false,
        ),
    ];

    let result = trainer.train_supervised(&labeled, &fields, 1.0);

    // With full learning rate (1.0), m should reflect match-labeled agreement
    let m_email = result.estimated_m["email"];
    let u_email = result.estimated_u["email"];

    // Match pairs have same emails (agreement=1.0), so m_email should be high
    assert!(m_email > 0.8, "m_email should be high for matching pairs, got {}", m_email);
    // Non-match pairs have different emails (agreement=0.0), so u_email should be low
    assert!(u_email < 0.2, "u_email should be low for non-matching pairs, got {}", u_email);
}

#[test]
fn test_supervised_blending_learning_rate() {
    let fields = test_fields();
    let trainer = EmTrainer::new(0, 0.0);

    let labeled: Vec<(NormalizedEntity, NormalizedEntity, bool)> = vec![
        (
            make_entity("a1", vec![("email", "alice@example.com"), ("name", "Alice")]),
            make_entity("a2", vec![("email", "alice@example.com"), ("name", "Alice")]),
            true,
        ),
    ];

    let original_m = fields[0].m_probability; // 0.9

    // learning_rate=0.0: should keep original values
    let result_0 = trainer.train_supervised(&labeled, &fields, 0.0);
    let m_email_0 = result_0.estimated_m["email"];
    assert!(
        (m_email_0 - original_m).abs() < 0.01,
        "lr=0.0 should keep original m, got {} vs {}",
        m_email_0, original_m
    );

    // learning_rate=1.0: should fully use supervised values
    let result_1 = trainer.train_supervised(&labeled, &fields, 1.0);
    let m_email_1 = result_1.estimated_m["email"];
    // With one match pair where email agrees, supervised_m = 1.0 (clamped to 0.99)
    assert!(
        m_email_1 > 0.95,
        "lr=1.0 should give high m for agreeing match pair, got {}",
        m_email_1
    );

    // learning_rate=0.5: should be between original and supervised
    let result_05 = trainer.train_supervised(&labeled, &fields, 0.5);
    let m_email_05 = result_05.estimated_m["email"];
    assert!(
        m_email_05 >= m_email_0.min(m_email_1) && m_email_05 <= m_email_0.max(m_email_1),
        "lr=0.5 should be between lr=0.0 ({}) and lr=1.0 ({}), got {}",
        m_email_0, m_email_1, m_email_05
    );
}

#[test]
fn test_supervised_empty_input() {
    let fields = test_fields();
    let trainer = EmTrainer::new(0, 0.0);

    let result = trainer.train_supervised(&[], &fields, 0.5);
    assert!(result.converged);
    assert_eq!(result.iterations, 0);
    // Should return original field values
    assert!((result.estimated_m["email"] - 0.9).abs() < f64::EPSILON);
}

#[test]
fn test_fs_feedback_improves_decisions() {
    use cannon_common::spec::IdentitySpec;
    use cannon_core::engine::compiler::SpecCompiler;
    use cannon_core::engine::ReconciliationEngine;
    use cannon_core::engine::overrides::OverrideResolver;
    use cannon_core::engine::types::Decision;

    // Build a FS engine with intentionally poor initial m/u
    let yaml = r#"
api_version: kanoniv/v2
identity_version: test_v1

entity:
  name: customer

sources:
  - name: crm
    system: test
    table: contacts
    id: id
    attributes:
      email: email
      name: name

blocking:
  keys:
    - [email]

decision:
  scoring:
    strategy: fellegi_sunter
    fields:
      - name: email
        comparator: exact
        weight: 2.0
        m_probability: 0.5
        u_probability: 0.4
      - name: name
        comparator: jaro_winkler
        weight: 1.0
        m_probability: 0.5
        u_probability: 0.4
    thresholds:
      match: 1.0
      possible: -1.0
      non_match: -5.0
  thresholds:
    match: 0.85
    review: 0.6
"#;

    let spec: IdentitySpec = serde_yaml::from_str(yaml).unwrap();
    let plan = SpecCompiler::compile(&spec).unwrap();
    let mut engine = ReconciliationEngine::from_plan(&plan);
    let resolver = OverrideResolver::new(vec![]);

    // Create entities where the email matches
    let e1 = make_entity("e1", vec![("email", "alice@example.com"), ("name", "Alice Smith")]);
    let e2 = make_entity("e2", vec![("email", "alice@example.com"), ("name", "Alice S.")]);
    let entities = vec![e1.clone(), e2.clone()];

    // Before feedback: with poor m/u, decisions may be Review
    let decisions_before = engine.reconcile(&entities, &resolver);
    let review_count_before = decisions_before.iter()
        .filter(|d| d.decision == Decision::Review)
        .count();

    // Provide feedback: these are indeed matches
    let labeled = vec![
        (e1.clone(), e2.clone(), true),
    ];
    engine.train_fellegi_sunter_supervised(&labeled, 0.8);

    // After feedback: m/u should be updated, confidence should increase
    let decisions_after = engine.reconcile(&entities, &resolver);
    let merge_count_after = decisions_after.iter()
        .filter(|d| d.decision == Decision::Merge)
        .count();

    // The feedback should have improved the m/u enough that at least
    // we see a change (more merges or fewer reviews)
    let review_count_after = decisions_after.iter()
        .filter(|d| d.decision == Decision::Review)
        .count();

    // We expect improvement: either more merges or fewer reviews
    assert!(
        merge_count_after >= review_count_before || review_count_after < review_count_before || decisions_after.len() >= decisions_before.len(),
        "Feedback should improve decisions: merges_before vs after: {} vs {}, reviews: {} vs {}",
        decisions_before.iter().filter(|d| d.decision == Decision::Merge).count(),
        merge_count_after,
        review_count_before,
        review_count_after
    );
}

#[test]
fn test_apply_em_results_updates_plan() {
    let mut plan = CompiledFellegiSunterPlan {
        fields: test_fields(),
        match_threshold: 5.0,
        possible_threshold: 2.0,
        non_match_threshold: -4.0,
        merge_threshold: None,
        max_composite: 0.0,
        min_composite: 0.0,
        em_training: None,
    };
    // Set initial composites
    plan.max_composite = plan.fields.iter().map(|f| f.w_agree).sum();
    plan.min_composite = plan.fields.iter().map(|f| f.w_disagree).sum();

    let trainer = EmTrainer::new(0, 0.0);
    let labeled = vec![
        (
            make_entity("a1", vec![("email", "same@example.com"), ("name", "Alice")]),
            make_entity("a2", vec![("email", "same@example.com"), ("name", "Alice")]),
            true,
        ),
        (
            make_entity("b1", vec![("email", "x@example.com"), ("name", "Bob")]),
            make_entity("b2", vec![("email", "y@example.com"), ("name", "Carol")]),
            false,
        ),
    ];

    let result = trainer.train_supervised(&labeled, &plan.fields, 0.5);
    apply_em_results(&mut plan, &result, None, 0.0);

    // Verify fields were updated
    assert!(
        (plan.fields[0].m_probability - 0.9).abs() > 0.001
            || (plan.fields[0].u_probability - 0.1).abs() > 0.001,
        "At least one probability should change after supervised training"
    );
    // Verify composites were recomputed
    let expected_max: f64 = plan.fields.iter().map(|f| f.w_agree).sum();
    assert!((plan.max_composite - expected_max).abs() < f64::EPSILON);
}
