swarm-engine-core 0.1.6

//! LearnedProvider - 学習データへの抽象アクセス
//!
//! Swarm 全体で学習済みデータにアクセスするための抽象層。
//! Selection、Orchestrator、Worker など様々なコンポーネントから利用可能。
//!
//! # 設計思想
//!
//! - **Model**: 計算ロジックを持ち、スコアを事前計算して保持
//! - **Provider**: Model の結果を返すだけ（計算ロジックを持たない）

use std::collections::HashMap;
use std::sync::Arc;

use super::stats::LearnStats;
use super::stats_model::ScoreModel;

// ============================================================================
// LearningQuery - 学習データへのクエリ
// ============================================================================

/// 学習データへのクエリ
///
/// target を含む全てのコンテキストを統一的に扱う。
#[derive(Debug, Clone)]
pub enum LearningQuery<'a> {
    /// 遷移ボーナス (prev → action @ target)
    ///
    /// 成功エピソードで頻出する遷移は正の値、失敗エピソードで頻出は負の値。
    Transition {
        prev: &'a str,
        action: &'a str,
        target: Option<&'a str>,
    },

    /// コンテキスト条件付きボーナス
    ///
    /// prev_action との組み合わせでの成功率補正。
    Contextual {
        prev: &'a str,
        action: &'a str,
        target: Option<&'a str>,
    },

    /// N-gram ボーナス（3-gram パターンの価値）
    Ngram {
        prev_prev: &'a str,
        prev: &'a str,
        action: &'a str,
        target: Option<&'a str>,
    },

    /// Confidence スコア（コンテキスト付き学習ボーナス）
    ///
    /// prev / prev_prev がある場合、Transition + Contextual + Ngram を統合。
    /// ない場合は action 単体の平均値を返す。
    Confidence {
        action: &'a str,
        target: Option<&'a str>,
        /// 直前のアクション（任意）
        prev: Option<&'a str>,
        /// 2つ前のアクション（任意、N-gram 用）
        prev_prev: Option<&'a str>,
    },
}

impl<'a> LearningQuery<'a> {
    /// Transition クエリを作成
    pub fn transition(prev: &'a str, action: &'a str, target: Option<&'a str>) -> Self {
        Self::Transition {
            prev,
            action,
            target,
        }
    }

    /// Contextual クエリを作成
    pub fn contextual(prev: &'a str, action: &'a str, target: Option<&'a str>) -> Self {
        Self::Contextual {
            prev,
            action,
            target,
        }
    }

    /// Ngram クエリを作成
    pub fn ngram(
        prev_prev: &'a str,
        prev: &'a str,
        action: &'a str,
        target: Option<&'a str>,
    ) -> Self {
        Self::Ngram {
            prev_prev,
            prev,
            action,
            target,
        }
    }

    /// Confidence クエリを作成（コンテキストなし）
    pub fn confidence(action: &'a str, target: Option<&'a str>) -> Self {
        Self::Confidence {
            action,
            target,
            prev: None,
            prev_prev: None,
        }
    }

    /// Confidence クエリを作成（コンテキスト付き）
    ///
    /// prev / prev_prev を渡すと、Transition + Contextual + Ngram を統合計算。
    pub fn confidence_with_context(
        action: &'a str,
        target: Option<&'a str>,
        prev: Option<&'a str>,
        prev_prev: Option<&'a str>,
    ) -> Self {
        Self::Confidence {
            action,
            target,
            prev,
            prev_prev,
        }
    }
}

// ============================================================================
// LearningResult - クエリ結果
// ============================================================================

/// クエリ結果
#[derive(Debug, Clone, PartialEq, Default)]
pub enum LearningResult {
    /// スコア値
    Score(f64),
    /// データなし（事前確率を使用すべき）
    #[default]
    NotAvailable,
}

impl LearningResult {
    /// スコアを取得（NotAvailable 時はデフォルト値）
    pub fn score_or(&self, default: f64) -> f64 {
        match self {
            Self::Score(v) => *v,
            Self::NotAvailable => default,
        }
    }

    /// スコアを取得（NotAvailable 時は 0.0）
    pub fn score(&self) -> f64 {
        self.score_or(0.0)
    }

    /// データが存在するか
    pub fn is_available(&self) -> bool {
        matches!(self, Self::Score(_))
    }
}

// ============================================================================
// LearnedProvider trait
// ============================================================================

/// 学習済みデータへのアクセス Provider
///
/// Swarm の各コンポーネントから利用される統一的なインターフェース。
/// 計算ロジックは持たない。Model から結果を取得して返すだけ。
pub trait LearnedProvider: Send + Sync {
    /// クエリを実行してボーナス/スコアを取得
    fn query(&self, q: LearningQuery<'_>) -> LearningResult;

    /// 内部の LearnStats を取得（永続化用、実装がある場合のみ）
    fn stats(&self) -> Option<&LearnStats> {
        None
    }

    /// 内部の ScoreModel を取得（実装がある場合のみ）
    fn model(&self) -> Option<&ScoreModel> {
        None
    }
}

/// Provider の共有参照型
pub type SharedLearnedProvider = Arc<dyn LearnedProvider>;

// ============================================================================
// ScoreModelProvider - Model ベースの Provider
// ============================================================================

/// ScoreModel を使った Provider 実装
///
/// Model が事前計算したスコアを返すだけ。計算ロジックを持たない。
pub struct ScoreModelProvider {
    model: ScoreModel,
    stats: Option<LearnStats>,
}

impl ScoreModelProvider {
    /// Model から Provider を作成
    pub fn new(model: ScoreModel) -> Self {
        Self { model, stats: None }
    }

    /// LearnStats から Provider を作成（Model を自動構築）
    pub fn from_stats(stats: LearnStats) -> Self {
        let model = ScoreModel::from_stats(&stats);
        Self {
            model,
            stats: Some(stats),
        }
    }

    /// Model を取得
    pub fn inner(&self) -> &ScoreModel {
        &self.model
    }

    /// Model を更新
    pub fn update_model(&mut self, model: ScoreModel) {
        self.model = model;
    }
}

impl LearnedProvider for ScoreModelProvider {
    fn query(&self, q: LearningQuery<'_>) -> LearningResult {
        match q {
            LearningQuery::Transition {
                prev,
                action,
                target,
            } => match self.model.transition(prev, action, target) {
                Some(score) => LearningResult::Score(score),
                None => LearningResult::NotAvailable,
            },

            LearningQuery::Contextual {
                prev,
                action,
                target,
            } => match self.model.contextual(prev, action, target) {
                Some(score) => LearningResult::Score(score),
                None => LearningResult::NotAvailable,
            },

            LearningQuery::Ngram {
                prev_prev,
                prev,
                action,
                target,
            } => match self.model.ngram(prev_prev, prev, action, target) {
                Some(score) => LearningResult::Score(score),
                None => LearningResult::NotAvailable,
            },

            LearningQuery::Confidence {
                action,
                target,
                prev,
                prev_prev,
            } => match self.model.confidence(action, target, prev, prev_prev) {
                Some(score) => LearningResult::Score(score),
                None => LearningResult::NotAvailable,
            },
        }
    }

    fn stats(&self) -> Option<&LearnStats> {
        self.stats.as_ref()
    }

    fn model(&self) -> Option<&ScoreModel> {
        Some(&self.model)
    }
}

// ============================================================================
// NullProvider
// ============================================================================

/// ボーナスを返さない Null Provider
#[derive(Debug, Clone, Default)]
pub struct NullProvider;

impl LearnedProvider for NullProvider {
    fn query(&self, _q: LearningQuery<'_>) -> LearningResult {
        LearningResult::NotAvailable
    }
}

// ============================================================================
// ConfidenceMapProvider
// ============================================================================

/// HashMap<String, f64> ベースの静的 Provider
///
/// DependencyGraph から生成された confidence_map を使用。
#[derive(Debug, Clone, Default)]
pub struct ConfidenceMapProvider {
    confidence: HashMap<String, f64>,
}

impl ConfidenceMapProvider {
    pub fn new(confidence: HashMap<String, f64>) -> Self {
        Self { confidence }
    }

    /// confidence を取得
    pub fn get(&self, action: &str) -> Option<f64> {
        self.confidence.get(action).copied()
    }
}

impl LearnedProvider for ConfidenceMapProvider {
    fn query(&self, q: LearningQuery<'_>) -> LearningResult {
        match q {
            LearningQuery::Confidence { action, .. } => {
                match self.get(action) {
                    Some(c) => {
                        // confidence - 0.5 を返して、0.5 を中立点とする
                        LearningResult::Score(c - 0.5)
                    }
                    None => LearningResult::NotAvailable,
                }
            }

            // ConfidenceMapProvider は他のクエリには対応しない
            _ => LearningResult::NotAvailable,
        }
    }
}

// ============================================================================
// LearnStatsProvider - Stats + Model を一体管理
// ============================================================================

/// LearnStats と ScoreModel を一体管理する Provider
///
/// Stats を更新すると自動的に Model も再構築される。
/// `rebuild_model()` の呼び忘れを防止。
pub struct LearnStatsProvider {
    stats: LearnStats,
    model: ScoreModel,
}

impl LearnStatsProvider {
    pub fn new(stats: LearnStats) -> Self {
        let model = ScoreModel::from_stats(&stats);
        Self { stats, model }
    }

    /// Stats を取得（読み取り専用）
    pub fn stats(&self) -> &LearnStats {
        &self.stats
    }

    /// Model を取得
    pub fn model(&self) -> &ScoreModel {
        &self.model
    }

    /// Stats を更新し、Model を自動再構築
    ///
    /// Stats を直接変更したい場合はこのメソッドを使用。
    /// Model の再構築忘れを防止。
    pub fn update_stats<F>(&mut self, f: F)
    where
        F: FnOnce(&mut LearnStats),
    {
        f(&mut self.stats);
        self.model = ScoreModel::from_stats(&self.stats);
    }

    /// Stats を置換し、Model を自動再構築
    pub fn replace_stats(&mut self, stats: LearnStats) {
        self.stats = stats;
        self.model = ScoreModel::from_stats(&self.stats);
    }
}

impl LearnedProvider for LearnStatsProvider {
    fn query(&self, q: LearningQuery<'_>) -> LearningResult {
        match q {
            LearningQuery::Transition {
                prev,
                action,
                target,
            } => match self.model.transition(prev, action, target) {
                Some(score) => LearningResult::Score(score),
                None => LearningResult::NotAvailable,
            },

            LearningQuery::Contextual {
                prev,
                action,
                target,
            } => match self.model.contextual(prev, action, target) {
                Some(score) => LearningResult::Score(score),
                None => LearningResult::NotAvailable,
            },

            LearningQuery::Ngram {
                prev_prev,
                prev,
                action,
                target,
            } => match self.model.ngram(prev_prev, prev, action, target) {
                Some(score) => LearningResult::Score(score),
                None => LearningResult::NotAvailable,
            },

            LearningQuery::Confidence {
                action,
                target,
                prev,
                prev_prev,
            } => match self.model.confidence(action, target, prev, prev_prev) {
                Some(score) => LearningResult::Score(score),
                None => LearningResult::NotAvailable,
            },
        }
    }

    fn stats(&self) -> Option<&LearnStats> {
        Some(&self.stats)
    }

    fn model(&self) -> Option<&ScoreModel> {
        Some(&self.model)
    }
}

// ============================================================================
// Tests
// ============================================================================

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

    #[test]
    fn test_learning_result_score_or() {
        assert_eq!(LearningResult::Score(0.5).score_or(0.0), 0.5);
        assert_eq!(LearningResult::NotAvailable.score_or(0.0), 0.0);
        assert_eq!(LearningResult::NotAvailable.score_or(-1.0), -1.0);
    }

    #[test]
    fn test_learning_result_is_available() {
        assert!(LearningResult::Score(0.5).is_available());
        assert!(!LearningResult::NotAvailable.is_available());
    }

    #[test]
    fn test_null_provider() {
        let provider = NullProvider;

        assert_eq!(
            provider.query(LearningQuery::transition("A", "B", None)),
            LearningResult::NotAvailable
        );
        assert_eq!(
            provider.query(LearningQuery::contextual("A", "B", Some("svc1"))),
            LearningResult::NotAvailable
        );
        assert_eq!(
            provider.query(LearningQuery::ngram("A", "B", "C", None)),
            LearningResult::NotAvailable
        );
    }

    #[test]
    fn test_confidence_map_provider() {
        let mut map = HashMap::new();
        map.insert("grep".to_string(), 0.8);
        map.insert("restart".to_string(), 0.3);

        let provider = ConfidenceMapProvider::new(map);

        // Confidence クエリには対応
        let result = provider.query(LearningQuery::confidence("grep", None));
        let score = result.score();
        assert!((score - 0.3).abs() < 1e-10, "expected ~0.3, got {}", score); // 0.8 - 0.5

        let result = provider.query(LearningQuery::confidence("restart", None));
        let score = result.score();
        assert!(
            (score - (-0.2)).abs() < 1e-10,
            "expected ~-0.2, got {}",
            score
        ); // 0.3 - 0.5

        // 存在しないアクション
        let result = provider.query(LearningQuery::confidence("unknown", None));
        assert_eq!(result, LearningResult::NotAvailable);

        // 他のクエリには対応しない
        let result = provider.query(LearningQuery::transition("A", "B", None));
        assert_eq!(result, LearningResult::NotAvailable);
    }

    #[test]
    fn test_learning_query_constructors() {
        let q = LearningQuery::transition("A", "B", Some("svc1"));
        assert!(matches!(
            q,
            LearningQuery::Transition {
                prev: "A",
                action: "B",
                target: Some("svc1")
            }
        ));

        let q = LearningQuery::ngram("A", "B", "C", None);
        assert!(matches!(
            q,
            LearningQuery::Ngram {
                prev_prev: "A",
                prev: "B",
                action: "C",
                target: None
            }
        ));

        let q =
            LearningQuery::confidence_with_context("action", None, Some("prev"), Some("prev_prev"));
        assert!(matches!(
            q,
            LearningQuery::Confidence {
                action: "action",
                prev: Some("prev"),
                prev_prev: Some("prev_prev"),
                ..
            }
        ));
    }

    #[test]
    fn test_score_model_provider() {
        use crate::learn::stats::{ContextualActionStats, LearnStats};

        let mut stats = LearnStats::default();

        // テストデータ
        stats
            .episode_transitions
            .success_transitions
            .insert(("A".to_string(), "B".to_string()), 10);
        stats
            .episode_transitions
            .failure_transitions
            .insert(("A".to_string(), "B".to_string()), 2);
        stats.contextual_stats.insert(
            ("A".to_string(), "B".to_string()),
            ContextualActionStats {
                visits: 12,
                successes: 10,
                failures: 2,
            },
        );
        stats
            .ngram_stats
            .trigrams
            .insert(("X".to_string(), "A".to_string(), "B".to_string()), (9, 1));

        let provider = ScoreModelProvider::from_stats(stats);

        // Transition クエリ
        let result = provider.query(LearningQuery::transition("A", "B", None));
        assert!(result.is_available());

        // Contextual クエリ
        let result = provider.query(LearningQuery::contextual("A", "B", None));
        assert!(result.is_available());
        assert!(result.score() > 0.0, "成功率が高いので正のスコア");

        // Ngram クエリ
        let result = provider.query(LearningQuery::ngram("X", "A", "B", None));
        assert!(result.is_available());

        // Confidence クエリ（コンテキスト付き）
        let result = provider.query(LearningQuery::confidence_with_context(
            "B",
            None,
            Some("A"),
            Some("X"),
        ));
        assert!(result.is_available());
    }
}