swarm-engine-core 0.1.6

//! Exploration Operator - Mutation と Selection のペア
//!
//! # 設計思想
//!
//! Mutation（Map展開ロジック）と Selection（ノード選択ロジック）は
//! アルゴリズム的にセットである。
//!
//! - Discover フェーズで発見した結果を Selection の統計に反映
//! - その統計が次の Selection に影響
//!
//! Operator はこの連携を担保する。
//!
//! # 責務
//!
//! | コンポーネント | 責務 |
//! |---------------|------|
//! | MutationLogic | 「結果をどう Map に反映するか」 |
//! | SelectionLogic | 「次にどの Node を選ぶか」（selection モジュール） |
//! | Operator | M と Sel の連携、共有状態(stats)の管理 |
//! | Space | Map + Operator を保持、API 提供 |
//!
//! # 使用例
//!
//! ```ignore
//! use swarm_engine_core::exploration::{
//!     Operator, RulesBasedMutation,
//!     selection::{Ucb1, Fifo},
//! };
//!
//! // 静的型付け Operator
//! let operator = Operator::new(
//!     RulesBasedMutation::new(),
//!     Ucb1::new(1.41),
//!     rules,
//! );
//!
//! // 動的 Selection の Operator
//! let operator = Operator::new(
//!     RulesBasedMutation::new(),
//!     AnySelection::from_kind(SelectionKind::Ucb1, 1.41),
//!     rules,
//! );
//! ```

use std::fmt::Debug;
use std::marker::PhantomData;
use std::sync::Arc;

use super::map::{ExplorationMap, GraphMap, MapNodeId, MapState};
use super::mutation::{ActionNodeData, ExplorationResult, MapUpdate, MutationInput};
use super::node_rules::Rules;
use super::selection::{AnySelection, Fifo, SelectionLogic, Ucb1};
use crate::actions::ActionsConfig;
use crate::learn::{LearnedProvider, NullProvider, SharedLearnedProvider};
use crate::online_stats::SwarmStats;

// ============================================================================
// MutationLogic trait
// ============================================================================

/// Map 展開ロジック
///
/// 入力を MapUpdate に変換する責務のみを持つ。
/// Selection とは分離されている。
pub trait MutationLogic<N, E, S, R>: Send + Sync
where
    N: Debug + Clone,
    E: Debug + Clone,
    S: MapState,
    R: Rules,
{
    /// 入力を MapUpdate に変換
    ///
    /// stats を参照して展開ロジックを調整できる（Optional）。
    fn interpret(
        &self,
        input: &dyn MutationInput,
        map: &GraphMap<N, E, S>,
        actions: &ActionsConfig,
        rules: &R,
        stats: &SwarmStats,
    ) -> Vec<MapUpdate<N, E, S>>;

    /// 初期ノードを展開
    fn initialize(
        &self,
        root_id: MapNodeId,
        initial_contexts: &[&str],
        rules: &R,
    ) -> Vec<MapUpdate<N, E, S>>;

    /// ノードデータを生成
    fn create_node_data(&self, input: &dyn MutationInput) -> N;

    /// エッジデータを生成
    fn create_edge_data(&self, input: &dyn MutationInput) -> E;

    /// 初期状態を生成
    fn initial_state(&self) -> S;

    /// 名前
    fn name(&self) -> &str;
}

// ============================================================================
// Operator - Mutation と Selection のペア
// ============================================================================

/// Mutation と Selection を組み合わせた Operator
///
/// SwarmStats を参照して Selection を行う。
/// 統計は ActionEventPublisher 経由で記録されるため、Operator 自身は統計を管理しない。
pub struct Operator<M, Sel, N, E, S, R>
where
    N: Debug + Clone,
    E: Debug + Clone,
    S: MapState,
    R: Rules,
    M: MutationLogic<N, E, S, R>,
    Sel: SelectionLogic<N, E, S>,
{
    mutation: M,
    pub selection: Sel,
    rules: R,
    /// 学習済みデータ Provider
    provider: SharedLearnedProvider,
    _phantom: PhantomData<(N, E, S)>,
}

impl<M, Sel, N, E, S, R> Operator<M, Sel, N, E, S, R>
where
    N: Debug + Clone,
    E: Debug + Clone,
    S: MapState,
    R: Rules,
    M: MutationLogic<N, E, S, R>,
    Sel: SelectionLogic<N, E, S>,
{
    /// 新しい Operator を作成
    pub fn new(mutation: M, selection: Sel, rules: R) -> Self {
        Self {
            mutation,
            selection,
            rules,
            provider: Arc::new(NullProvider),
            _phantom: PhantomData,
        }
    }

    /// Provider を設定
    pub fn with_provider(mut self, provider: SharedLearnedProvider) -> Self {
        self.provider = provider;
        self
    }

    /// Provider への参照
    pub fn provider(&self) -> &dyn LearnedProvider {
        self.provider.as_ref()
    }

    /// Provider を設定
    pub fn set_provider(&mut self, provider: SharedLearnedProvider) {
        self.provider = provider;
    }

    /// Rules への参照
    pub fn rules(&self) -> &R {
        &self.rules
    }

    /// Selection への参照
    pub fn selection(&self) -> &Sel {
        &self.selection
    }

    /// Selection への可変参照
    pub fn selection_mut(&mut self) -> &mut Sel {
        &mut self.selection
    }

    /// Selection を置き換え
    pub fn set_selection(&mut self, selection: Sel) {
        self.selection = selection;
    }

    /// 入力を MapUpdate に変換
    ///
    /// Note: 統計は ActionEventPublisher 経由で別途記録されるため、
    /// この関数では統計を更新しない。
    pub fn interpret(
        &self,
        input: &dyn MutationInput,
        map: &GraphMap<N, E, S>,
        actions: &ActionsConfig,
        stats: &SwarmStats,
    ) -> Vec<MapUpdate<N, E, S>> {
        self.mutation
            .interpret(input, map, actions, &self.rules, stats)
    }

    /// 初期ノードを展開
    pub fn initialize(
        &self,
        root_id: MapNodeId,
        initial_contexts: &[&str],
    ) -> Vec<MapUpdate<N, E, S>> {
        self.mutation
            .initialize(root_id, initial_contexts, &self.rules)
    }

    /// 次のノードを1つ選択
    pub fn next(&self, map: &GraphMap<N, E, S>, stats: &SwarmStats) -> Option<MapNodeId> {
        self.selection.next(map, stats, self.provider.as_ref())
    }

    /// 次のノードを複数選択
    pub fn select(
        &self,
        map: &GraphMap<N, E, S>,
        count: usize,
        stats: &SwarmStats,
    ) -> Vec<MapNodeId> {
        self.selection
            .select(map, count, stats, self.provider.as_ref())
    }

    /// ノードのスコアを計算
    pub fn score(&self, action: &str, target: Option<&str>, stats: &SwarmStats) -> f64 {
        self.selection
            .score(action, target, stats, self.provider.as_ref())
    }

    /// 完了判定
    ///
    /// デフォルト: フロンティア枯渇で完了
    pub fn is_complete(&self, map: &GraphMap<N, E, S>) -> bool {
        map.frontiers().is_empty()
    }

    /// ノードデータを生成
    pub fn create_node_data(&self, input: &dyn MutationInput) -> N {
        self.mutation.create_node_data(input)
    }

    /// エッジデータを生成
    pub fn create_edge_data(&self, input: &dyn MutationInput) -> E {
        self.mutation.create_edge_data(input)
    }

    /// 初期状態を生成
    pub fn initial_state(&self) -> S {
        self.mutation.initial_state()
    }

    /// Operator 名
    pub fn name(&self) -> String {
        format!("{}+{}", self.mutation.name(), self.selection.name())
    }
}

// ============================================================================
// RulesBasedMutation - Rules ベースの基本 Mutation
// ============================================================================

/// Rules ベースの基本 Mutation
///
/// NodeRules を使用してノード遷移を制御する標準的な MutationLogic 実装。
#[derive(Debug, Clone, Default)]
pub struct RulesBasedMutation;

impl RulesBasedMutation {
    pub fn new() -> Self {
        Self
    }
}

impl<E, S, R> MutationLogic<ActionNodeData, E, S, R> for RulesBasedMutation
where
    E: Debug + Clone + Default,
    S: MapState + Default,
    R: Rules,
{
    fn interpret(
        &self,
        input: &dyn MutationInput,
        _map: &GraphMap<ActionNodeData, E, S>,
        _actions: &ActionsConfig,
        rules: &R,
        _stats: &SwarmStats,
    ) -> Vec<MapUpdate<ActionNodeData, E, S>> {
        let node_id = input.node_id();
        let action_name = input.action_name();

        match input.result() {
            ExplorationResult::Discover(children) => {
                let successors = rules.successors(action_name);
                tracing::debug!(
                    node_id = node_id.0,
                    action = %action_name,
                    target = ?input.target(),
                    children_count = children.len(),
                    successors = ?successors,
                    "ExpMap: Discover result"
                );
                if successors.is_empty() || children.is_empty() {
                    tracing::debug!(
                        node_id = node_id.0,
                        "ExpMap: Close (no successors or children)"
                    );
                    return vec![MapUpdate::Close(node_id)];
                }

                let mut updates: Vec<MapUpdate<ActionNodeData, E, S>> = Vec::new();

                for child in children {
                    for next_action in &successors {
                        let dedup_key = format!("{}:{}", next_action, child);
                        let node_data = ActionNodeData::new(*next_action).with_target(child);
                        tracing::debug!(
                            parent = node_id.0,
                            next_action = %next_action,
                            child = %child,
                            dedup_key = %dedup_key,
                            "ExpMap: AddChild (from Discover)"
                        );

                        updates.push(MapUpdate::AddChild {
                            parent: node_id,
                            edge_data: E::default(),
                            node_data,
                            node_state: S::default(),
                            dedup_key,
                        });
                    }
                }

                updates.push(MapUpdate::Close(node_id));
                updates
            }

            ExplorationResult::Success => {
                tracing::debug!(
                    node_id = node_id.0,
                    action = %action_name,
                    target = ?input.target(),
                    is_terminal = rules.is_terminal(action_name),
                    "ExpMap: Success result"
                );
                if rules.is_terminal(action_name) {
                    tracing::debug!(node_id = node_id.0, "ExpMap: Close (terminal action)");
                    return vec![MapUpdate::Close(node_id)];
                }

                let successors = rules.successors(action_name);
                if successors.is_empty() {
                    tracing::debug!(node_id = node_id.0, "ExpMap: Close (no successors)");
                    return vec![MapUpdate::Close(node_id)];
                }

                let target = input.target();
                let mut updates: Vec<MapUpdate<ActionNodeData, E, S>> = Vec::new();

                tracing::debug!(
                    node_id = node_id.0,
                    successors = ?successors,
                    target = ?target,
                    "ExpMap: expanding successors"
                );
                for next_action in successors {
                    if let Some((param_key, param_values)) = rules.param_variants(next_action) {
                        for param_value in param_values {
                            let dedup_key = match target {
                                Some(t) => {
                                    format!("{}:{}:{}:{}", next_action, t, param_key, param_value)
                                }
                                None => format!("{}:_:{}:{}", next_action, param_key, param_value),
                            };
                            tracing::debug!(
                                parent = node_id.0,
                                next_action = %next_action,
                                param = %format!("{}={}", param_key, param_value),
                                dedup_key = %dedup_key,
                                "ExpMap: AddChild (with param variant)"
                            );

                            let mut node_data = ActionNodeData::new(next_action);
                            if let Some(t) = target {
                                node_data = node_data.with_target(t);
                            }
                            node_data = node_data.with_arg(param_key, param_value);

                            updates.push(MapUpdate::AddChild {
                                parent: node_id,
                                edge_data: E::default(),
                                node_data,
                                node_state: S::default(),
                                dedup_key,
                            });
                        }
                    } else {
                        let dedup_key = match target {
                            Some(t) => format!("{}:{}", next_action, t),
                            None => format!("{}:_", next_action),
                        };
                        tracing::debug!(
                            parent = node_id.0,
                            next_action = %next_action,
                            dedup_key = %dedup_key,
                            "ExpMap: AddChild (from Success)"
                        );
                        let mut node_data = ActionNodeData::new(next_action);
                        if let Some(t) = target {
                            node_data = node_data.with_target(t);
                        }
                        updates.push(MapUpdate::AddChild {
                            parent: node_id,
                            edge_data: E::default(),
                            node_data,
                            node_state: S::default(),
                            dedup_key,
                        });
                    }
                }

                updates.push(MapUpdate::Close(node_id));
                updates
            }

            ExplorationResult::Fail(ref err) => {
                tracing::debug!(
                    node_id = node_id.0,
                    action = %action_name,
                    target = ?input.target(),
                    error = ?err,
                    "ExpMap: Fail result, closing node"
                );
                vec![MapUpdate::Close(node_id)]
            }
        }
    }

    fn initialize(
        &self,
        root_id: MapNodeId,
        initial_contexts: &[&str],
        rules: &R,
    ) -> Vec<MapUpdate<ActionNodeData, E, S>> {
        let root_actions = rules.roots();
        tracing::debug!(
            root_id = root_id.0,
            root_actions = ?root_actions,
            initial_contexts = ?initial_contexts,
            "ExpMap: initialize"
        );
        let mut updates = Vec::new();

        for action in root_actions {
            for ctx in initial_contexts {
                let dedup_key = format!("{}:{}", action, ctx);
                let node_data = ActionNodeData::new(action).with_target(*ctx);
                tracing::debug!(
                    parent = root_id.0,
                    action = %action,
                    context = %ctx,
                    dedup_key = %dedup_key,
                    "ExpMap: AddChild (initial)"
                );
                updates.push(MapUpdate::AddChild {
                    parent: root_id,
                    edge_data: E::default(),
                    node_data,
                    node_state: S::default(),
                    dedup_key,
                });
            }
        }

        updates
    }

    fn create_node_data(&self, input: &dyn MutationInput) -> ActionNodeData {
        ActionNodeData::from_input(input)
    }

    fn create_edge_data(&self, _input: &dyn MutationInput) -> E {
        E::default()
    }

    fn initial_state(&self) -> S {
        S::default()
    }

    fn name(&self) -> &str {
        "RulesBased"
    }
}

// ============================================================================
// Type Aliases for Common Operators
// ============================================================================

/// FIFO Operator（FIFO選択 + RulesBasedMutation）
pub type FifoOperator<R> =
    Operator<RulesBasedMutation, Fifo, ActionNodeData, String, super::map::MapNodeState, R>;

/// UCB1 Operator
pub type Ucb1Operator<R> =
    Operator<RulesBasedMutation, Ucb1, ActionNodeData, String, super::map::MapNodeState, R>;

/// 動的に Selection を切り替え可能な Operator
///
/// Config や LLM から Selection を決定できる。
pub type ConfigurableOperator<R> =
    Operator<RulesBasedMutation, AnySelection, ActionNodeData, String, super::map::MapNodeState, R>;

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

#[cfg(test)]
mod tests {
    use super::super::map::MapNodeState;
    use super::super::selection::FifoSelection;
    use super::super::NodeRules;
    use super::*;

    struct TestInput {
        node_id: MapNodeId,
        action_name: String,
        target: Option<String>,
        result: ExplorationResult,
    }

    impl MutationInput for TestInput {
        fn node_id(&self) -> MapNodeId {
            self.node_id
        }
        fn action_name(&self) -> &str {
            &self.action_name
        }
        fn target(&self) -> Option<&str> {
            self.target.as_deref()
        }
        fn result(&self) -> &ExplorationResult {
            &self.result
        }
    }

    #[test]
    fn test_operator_basic() {
        let rules = NodeRules::for_testing();
        let operator: FifoOperator<NodeRules> =
            Operator::new(RulesBasedMutation::new(), FifoSelection::new(), rules);
        let stats = SwarmStats::new();

        let mut map: GraphMap<ActionNodeData, String, MapNodeState> = GraphMap::new();
        let root = map.create_root(ActionNodeData::new("root"), MapNodeState::Open);

        // Initialize
        let updates = operator.initialize(root, &["auth", "login"]);
        assert_eq!(updates.len(), 4); // 2 roots × 2 contexts

        // Apply updates
        for update in updates {
            map.apply_update(update, |k| k.to_string());
        }
        assert_eq!(map.node_count(), 5); // root + 4

        // Select
        let selected = operator.select(&map, 2, &stats);
        assert_eq!(selected.len(), 2);
    }

    #[test]
    fn test_operator_interpret() {
        let rules = NodeRules::for_testing();
        let operator: FifoOperator<NodeRules> =
            Operator::new(RulesBasedMutation::new(), FifoSelection::new(), rules);
        let stats = SwarmStats::new();

        let map: GraphMap<ActionNodeData, String, MapNodeState> = GraphMap::new();
        let actions = ActionsConfig::new();

        let input = TestInput {
            node_id: MapNodeId::new(0),
            action_name: "grep".to_string(),
            target: Some("auth.rs".to_string()),
            result: ExplorationResult::Success,
        };

        // interpret 実行（統計は ActionEventPublisher 経由で記録される）
        let updates = operator.interpret(&input, &map, &actions, &stats);
        // Success なので Close が返る
        assert!(!updates.is_empty());
    }

    #[test]
    fn test_operator_name() {
        let rules = NodeRules::for_testing();
        let operator: FifoOperator<NodeRules> =
            Operator::new(RulesBasedMutation::new(), FifoSelection::new(), rules);

        assert_eq!(operator.name(), "RulesBased+FIFO");
    }

    #[test]
    fn test_configurable_operator_works_like_fifo() {
        use super::super::selection::SelectionKind;

        let rules = NodeRules::for_testing();
        let operator: ConfigurableOperator<NodeRules> = Operator::new(
            RulesBasedMutation::new(),
            AnySelection::from_kind(SelectionKind::Fifo, 1.0),
            rules,
        );
        let stats = SwarmStats::new();

        let mut map: GraphMap<ActionNodeData, String, MapNodeState> = GraphMap::new();
        let root = map.create_root(ActionNodeData::new("root"), MapNodeState::Open);

        let updates = operator.initialize(root, &["auth"]);
        assert_eq!(updates.len(), 2); // grep:auth, glob:auth

        for update in updates {
            map.apply_update(update, |k| k.to_string());
        }

        let selected = operator.select(&map, 1, &stats);
        assert_eq!(selected.len(), 1);
    }
}