swarm-engine-core 0.1.6

//! SwarmStats - オンライン統計
//!
//! Swarm 内でリアルタイムに更新される基本統計。
//! Selection（UCB1/Thompson）が参照する visits/success_rate を提供。
//!
//! # 責務
//!
//! - アクション別統計（visits, successes, failures）
//! - アクション×ターゲット別統計
//! - グローバル統計
//!
//! # 分離された責務
//!
//! - 学習統計（EpisodeTransitions, NgramStats 等）→ `learn::LearnStats`

use std::collections::HashMap;
use std::time::Duration;

use serde::{Deserialize, Serialize};

use crate::events::ActionEvent;

/// Swarm オンライン統計
///
/// ActionEvent を受け取り、Selection が参照する基本統計を提供。
/// UCB1/Thompson が必要とする visits/success_rate をリアルタイムに更新。
#[derive(Debug, Clone, Default)]
pub struct SwarmStats {
    /// アクション別統計
    action_stats: HashMap<String, ActionStats>,
    /// アクション × ターゲット別統計
    action_target_stats: HashMap<(String, String), ActionStats>,
    /// グローバル統計（Worker アクションのみ）
    global: GlobalStats,
}

/// アクション単位の統計
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct ActionStats {
    /// 訪問回数
    pub visits: u32,
    /// 成功回数
    pub successes: u32,
    /// 失敗回数
    pub failures: u32,
    /// 発見数の合計
    pub discoveries: u32,
    /// 総実行時間
    #[serde(
        serialize_with = "serialize_duration",
        deserialize_with = "deserialize_duration"
    )]
    pub total_duration: Duration,
    /// KPI貢献度の累計
    #[serde(default)]
    pub kpi_total: f64,
    /// KPI貢献度が設定されたイベント数
    #[serde(default)]
    pub kpi_count: u32,
}

fn serialize_duration<S>(duration: &Duration, serializer: S) -> Result<S::Ok, S::Error>
where
    S: serde::Serializer,
{
    serializer.serialize_u64(duration.as_millis() as u64)
}

fn deserialize_duration<'de, D>(deserializer: D) -> Result<Duration, D::Error>
where
    D: serde::Deserializer<'de>,
{
    let millis = u64::deserialize(deserializer)?;
    Ok(Duration::from_millis(millis))
}

impl ActionStats {
    /// 成功率（訪問なしの場合は 0.5）
    pub fn success_rate(&self) -> f64 {
        if self.visits == 0 {
            0.5
        } else {
            self.successes as f64 / self.visits as f64
        }
    }

    /// 平均実行時間
    pub fn avg_duration(&self) -> Duration {
        if self.visits == 0 {
            Duration::ZERO
        } else {
            self.total_duration / self.visits
        }
    }

    /// 平均発見数
    pub fn avg_discoveries(&self) -> f64 {
        if self.visits == 0 {
            0.0
        } else {
            self.discoveries as f64 / self.visits as f64
        }
    }

    /// 平均KPI貢献度（KPIが設定されたイベントのみ）
    pub fn avg_kpi(&self) -> f64 {
        if self.kpi_count == 0 {
            0.0
        } else {
            self.kpi_total / self.kpi_count as f64
        }
    }
}

/// グローバル統計
#[derive(Debug, Clone, Default)]
pub struct GlobalStats {
    /// 総訪問回数
    pub total_visits: u32,
    /// 総成功回数
    pub total_successes: u32,
    /// 総失敗回数
    pub total_failures: u32,
    /// 総発見数
    pub total_discoveries: u32,
    /// 総実行時間
    pub total_duration: Duration,
    /// KPI貢献度の累計
    pub total_kpi: f64,
    /// KPI貢献度が設定されたイベント数
    pub kpi_count: u32,
}

impl GlobalStats {
    /// 成功率（訪問なしの場合は 1.0）
    pub fn success_rate(&self) -> f64 {
        if self.total_visits == 0 {
            1.0
        } else {
            self.total_successes as f64 / self.total_visits as f64
        }
    }

    /// 失敗率（訪問なしの場合は 0.0）
    pub fn failure_rate(&self) -> f64 {
        if self.total_visits == 0 {
            0.0
        } else {
            self.total_failures as f64 / self.total_visits as f64
        }
    }

    /// 平均KPI貢献度（KPIが設定されたイベントのみ）
    pub fn avg_kpi(&self) -> f64 {
        if self.kpi_count == 0 {
            0.0
        } else {
            self.total_kpi / self.kpi_count as f64
        }
    }
}

impl SwarmStats {
    pub fn new() -> Self {
        Self::default()
    }

    /// イベントを記録
    ///
    /// Selection 用の統計のみ更新。
    /// - LLM イベント、Manager イベントはスキップ
    /// - Worker の通常アクションのみ統計に反映
    pub fn record(&mut self, event: &ActionEvent) {
        let action = &event.action;
        let success = event.result.success;
        let duration = event.duration;

        // LLM 呼び出しイベントはスキップ
        if action == "llm_invoke" {
            return;
        }

        // Manager イベントはスキップ
        if event.worker_id.is_manager() {
            return;
        }

        // tick_start/tick_end はスキップ
        if action == "tick_start" || action == "tick_end" {
            return;
        }

        let target = event.target.as_deref();
        let discoveries = event.result.discoveries;
        let kpi = event.result.kpi_contribution;

        // グローバル統計を更新
        self.global.total_visits += 1;
        self.global.total_duration += duration;
        if success {
            self.global.total_successes += 1;
            self.global.total_discoveries += discoveries;
        } else {
            self.global.total_failures += 1;
        }
        if let Some(k) = kpi {
            self.global.total_kpi += k;
            self.global.kpi_count += 1;
        }

        // アクション統計を更新
        let action_stat = self.action_stats.entry(action.clone()).or_default();
        action_stat.visits += 1;
        action_stat.total_duration += duration;
        if success {
            action_stat.successes += 1;
            action_stat.discoveries += discoveries;
        } else {
            action_stat.failures += 1;
        }
        if let Some(k) = kpi {
            action_stat.kpi_total += k;
            action_stat.kpi_count += 1;
        }

        // アクション × ターゲット統計を更新
        if let Some(t) = target {
            let at_stat = self
                .action_target_stats
                .entry((action.clone(), t.to_string()))
                .or_default();
            at_stat.visits += 1;
            at_stat.total_duration += duration;
            if success {
                at_stat.successes += 1;
                at_stat.discoveries += discoveries;
            } else {
                at_stat.failures += 1;
            }
            if let Some(k) = kpi {
                at_stat.kpi_total += k;
                at_stat.kpi_count += 1;
            }
        }
    }

    /// アクションの統計を取得
    pub fn get_action_stats(&self, action: &str) -> ActionStats {
        self.action_stats.get(action).cloned().unwrap_or_default()
    }

    /// アクション × ターゲットの統計を取得
    pub fn get_action_target_stats(&self, action: &str, target: &str) -> ActionStats {
        self.action_target_stats
            .get(&(action.to_string(), target.to_string()))
            .cloned()
            .unwrap_or_default()
    }

    /// グローバル統計を取得
    pub fn global(&self) -> &GlobalStats {
        &self.global
    }

    /// 総訪問回数
    pub fn total_visits(&self) -> u32 {
        self.global.total_visits
    }

    /// 総成功回数
    pub fn total_successes(&self) -> u32 {
        self.global.total_successes
    }

    /// 総失敗回数
    pub fn total_failures(&self) -> u32 {
        self.global.total_failures
    }

    /// 全体の成功率
    pub fn success_rate(&self) -> f64 {
        self.global.success_rate()
    }

    /// 全体の失敗率
    pub fn failure_rate(&self) -> f64 {
        self.global.failure_rate()
    }

    /// 全アクションの統計一覧
    pub fn all_action_stats(&self) -> impl Iterator<Item = (&String, &ActionStats)> {
        self.action_stats.iter()
    }

    /// 統計をリセット
    pub fn reset(&mut self) {
        self.action_stats.clear();
        self.action_target_stats.clear();
        self.global = GlobalStats::default();
    }
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use super::*;
    use crate::events::{ActionEventBuilder, ActionEventResult};
    use crate::types::WorkerId;

    fn make_event(action: &str, target: Option<&str>, success: bool) -> ActionEvent {
        let mut builder =
            ActionEventBuilder::new(1, WorkerId(0), action).duration(Duration::from_millis(100));

        if let Some(t) = target {
            builder = builder.target(t);
        }

        let result = if success {
            ActionEventResult::success()
        } else {
            ActionEventResult::failure("error")
        };

        builder.result(result).build()
    }

    #[test]
    fn test_swarm_stats_basic() {
        let mut stats = SwarmStats::new();

        stats.record(&make_event("CheckStatus", Some("svc1"), true));
        stats.record(&make_event("CheckStatus", Some("svc1"), true));
        stats.record(&make_event("CheckStatus", Some("svc2"), false));

        let action_stats = stats.get_action_stats("CheckStatus");
        assert_eq!(action_stats.visits, 3);
        assert_eq!(action_stats.successes, 2);
        assert_eq!(action_stats.failures, 1);
        assert!((action_stats.success_rate() - 0.666).abs() < 0.01);

        let at_stats = stats.get_action_target_stats("CheckStatus", "svc1");
        assert_eq!(at_stats.visits, 2);
        assert_eq!(at_stats.successes, 2);
    }

    #[test]
    fn test_swarm_stats_global() {
        let mut stats = SwarmStats::new();

        stats.record(&make_event("A", None, true));
        stats.record(&make_event("B", None, true));
        stats.record(&make_event("C", None, false));

        assert_eq!(stats.total_visits(), 3);
        assert_eq!(stats.total_successes(), 2);
        assert_eq!(stats.total_failures(), 1);
        assert!((stats.success_rate() - 0.666).abs() < 0.01);
    }

    #[test]
    fn test_llm_invoke_skipped() {
        let mut stats = SwarmStats::new();

        let e = ActionEventBuilder::new(1, WorkerId(0), "llm_invoke")
            .result(ActionEventResult::success())
            .build();
        stats.record(&e);

        assert_eq!(stats.total_visits(), 0);
    }

    #[test]
    fn test_manager_skipped() {
        let mut stats = SwarmStats::new();

        let e = ActionEventBuilder::new(1, WorkerId::MANAGER, "decide")
            .result(ActionEventResult::success())
            .build();
        stats.record(&e);

        assert_eq!(stats.total_visits(), 0);
    }

    #[test]
    fn test_kpi_contribution() {
        let mut stats = SwarmStats::new();

        // KPIなしのイベント
        stats.record(&make_event("A", None, true));

        // KPIありのイベント
        let e1 = ActionEventBuilder::new(1, WorkerId(0), "B")
            .result(ActionEventResult::success().with_kpi(0.5))
            .build();
        stats.record(&e1);

        let e2 = ActionEventBuilder::new(2, WorkerId(0), "B")
            .result(ActionEventResult::success().with_kpi(0.8))
            .build();
        stats.record(&e2);

        // グローバル統計
        assert_eq!(stats.global().kpi_count, 2);
        assert!((stats.global().total_kpi - 1.3).abs() < 0.01);
        assert!((stats.global().avg_kpi() - 0.65).abs() < 0.01);

        // アクション別統計
        let a_stats = stats.get_action_stats("A");
        assert_eq!(a_stats.kpi_count, 0);
        assert_eq!(a_stats.avg_kpi(), 0.0);

        let b_stats = stats.get_action_stats("B");
        assert_eq!(b_stats.kpi_count, 2);
        assert!((b_stats.kpi_total - 1.3).abs() < 0.01);
        assert!((b_stats.avg_kpi() - 0.65).abs() < 0.01);
    }
}